(;module:
  lux
  (lux (control [monad #+ do]
                ["ex" exception #+ exception:])
       (data [maybe]
             [text]
             text/format
             (coll [list "list/" Functor<List>]))
       [meta #+ Monad<Meta>]
       (meta [type]
             (type ["tc" check])))
  (luxc ["&" base]
        (lang ["la" analysis #+ Analysis]
              (analysis ["&;" common]))))

(exception: #export Cannot-Infer)
(exception: #export Cannot-Infer-Argument)
(exception: #export Smaller-Variant-Than-Expected)
(exception: #export Invalid-Type-Application)
(exception: #export Not-A-Record-Type)
(exception: #export Not-A-Variant-Type)

## When doing inference, type-variables often need to be created in
## order to figure out which types are present in the expression being
## inferred.
## If a type-variable never gets bound/resolved to a type, then that
## means the expression can be generalized through universal
## quantification.
## When that happens, the type-variable must be replaced by an
## argument to the universally-quantified type.
(def: #export (replace-var var-id bound-idx type)
  (-> Nat Nat Type Type)
  (case type
    (#;Primitive name params)
    (#;Primitive name (list/map (replace-var var-id bound-idx) params))

    (^template [<tag>]
      (<tag> left right)
      (<tag> (replace-var var-id bound-idx left)
             (replace-var var-id bound-idx right)))
    ([#;Sum]
     [#;Product]
     [#;Function]
     [#;Apply])
    
    (#;Var id)
    (if (n.= var-id id)
      (#;Bound bound-idx)
      type)

    (^template [<tag>]
      (<tag> env quantified)
      (<tag> (list/map (replace-var var-id bound-idx) env)
             (replace-var var-id (n.+ +2 bound-idx) quantified)))
    ([#;UnivQ]
     [#;ExQ])
    
    _
    type))

(def: (replace-bound bound-idx replacementT type)
  (-> Nat Type Type Type)
  (case type
    (#;Primitive name params)
    (#;Primitive name (list/map (replace-bound bound-idx replacementT) params))

    (^template [<tag>]
      (<tag> left right)
      (<tag> (replace-bound bound-idx replacementT left)
             (replace-bound bound-idx replacementT right)))
    ([#;Sum]
     [#;Product]
     [#;Function]
     [#;Apply])
    
    (#;Bound idx)
    (if (n.= bound-idx idx)
      replacementT
      type)

    (^template [<tag>]
      (<tag> env quantified)
      (<tag> (list/map (replace-bound bound-idx replacementT) env)
             (replace-bound (n.+ +2 bound-idx) replacementT quantified)))
    ([#;UnivQ]
     [#;ExQ])
    
    _
    type))

## Type-inference works by applying some (potentially quantified) type
## to a sequence of values.
## Function types are used for this, although inference is not always
## done for function application (alternative uses may be records and
## tagged variants).
## But, so long as the type being used for the inference can be treated
## as a function type, this method of inference should work.
(def: #export (apply-function analyse inferT args)
  (-> &;Analyser Type (List Code) (Meta [Type (List Analysis)]))
  (case args
    #;Nil
    (do Monad<Meta>
      [_ (&;infer inferT)]
      (wrap [inferT (list)]))
    
    (#;Cons argC args')
    (case inferT
      (#;Named name unnamedT)
      (apply-function analyse unnamedT args)

      (#;UnivQ _)
      (&common;with-var
        (function [[var-id varT]]
          (do Monad<Meta>
            [[outputT argsA] (apply-function analyse (maybe;assume (type;apply (list varT) inferT)) args)]
            (do @
              [? (&;with-type-env
                   (tc;concrete? var-id))
               ## Quantify over the type if genericity/parametricity
               ## is discovered.
               outputT' (if ?
                          (&;with-type-env
                            (tc;clean var-id outputT))
                          (wrap (type;univ-q +1 (replace-var var-id +1 outputT))))]
              (wrap [outputT' argsA])))))

      (#;ExQ _)
      (do Monad<Meta>
        [[ex-id exT] (&;with-type-env
                       tc;existential)]
        (apply-function analyse (maybe;assume (type;apply (list exT) inferT)) args))

      (#;Apply inputT transT)
      (case (type;apply (list inputT) transT)
        (#;Some outputT)
        (apply-function analyse outputT args)

        #;None
        (&;throw Invalid-Type-Application (%type inferT)))

      ## Arguments are inferred back-to-front because, by convention,
      ## Lux functions take the most important arguments *last*, which
      ## means that the most information for doing proper inference is
      ## located in the last arguments to a function call.
      ## By inferring back-to-front, a lot of type-annotations can be
      ## avoided in Lux code, since the inference algorithm can piece
      ## things together more easily.
      (#;Function inputT outputT)
      (do Monad<Meta>
        [[outputT' args'A] (apply-function analyse outputT args')
         argA (&;with-stacked-errors
                (function [_] (Cannot-Infer-Argument
                               (format "Inferred Type: " (%type inputT) "\n"
                                       "     Argument: " (%code argC))))
                (&;with-expected-type inputT
                  (analyse argC)))]
        (wrap [outputT' (list& argA args'A)]))

      _
      (&;throw Cannot-Infer (format "     Type: " (%type inferT) "\n"
                                    "Arguments:"
                                    (|> args
                                        list;enumerate
                                        (list/map (function [[idx argC]]
                                                    (format "\n  " (%n idx) " " (%code argC))))
                                        (text;join-with "")))))
    ))

## Turns a record type into the kind of function type suitable for inference.
(def: #export (record type)
  (-> Type (Meta Type))
  (case type
    (#;Named name unnamedT)
    (record unnamedT)

    (^template [<tag>]
      (<tag> env bodyT)
      (do Monad<Meta>
        [bodyT+ (record bodyT)]
        (wrap (<tag> env bodyT+))))
    ([#;UnivQ]
     [#;ExQ])

    (#;Product _)
    (:: Monad<Meta> wrap (type;function (type;flatten-tuple type) type))

    _
    (&;throw Not-A-Record-Type (%type type))))

## Turns a variant type into the kind of function type suitable for inference.
(def: #export (variant tag expected-size type)
  (-> Nat Nat Type (Meta Type))
  (loop [depth +0
         currentT type]
    (case currentT
      (#;Named name unnamedT)
      (do Monad<Meta>
        [unnamedT+ (recur depth unnamedT)]
        (wrap unnamedT+))

      (^template [<tag>]
        (<tag> env bodyT)
        (do Monad<Meta>
          [bodyT+ (recur (n.inc depth) bodyT)]
          (wrap (<tag> env bodyT+))))
      ([#;UnivQ]
       [#;ExQ])

      (#;Sum _)
      (let [cases (type;flatten-variant currentT)
            actual-size (list;size cases)
            boundary (n.dec expected-size)]
        (cond (or (n.= expected-size actual-size)
                  (and (n.> expected-size actual-size)
                       (n.< boundary tag)))
              (case (list;nth tag cases)
                (#;Some caseT)
                (:: Monad<Meta> wrap (if (n.= +0 depth)
                                       (type;function (list caseT) currentT)
                                       (let [replace! (replace-bound (|> depth n.dec (n.* +2)) type)]
                                         (type;function (list (replace! caseT))
                                           (replace! currentT)))))

                #;None
                (&common;variant-out-of-bounds-error type expected-size tag))
              
              (n.< expected-size actual-size)
              (&;throw Smaller-Variant-Than-Expected
                       (format "Expected: " (%i (nat-to-int expected-size)) "\n"
                               "  Actual: " (%i (nat-to-int actual-size))))

              (n.= boundary tag)
              (let [caseT (type;variant (list;drop boundary cases))]
                (:: Monad<Meta> wrap (if (n.= +0 depth)
                                       (type;function (list caseT) currentT)
                                       (let [replace! (replace-bound (|> depth n.dec (n.* +2)) type)]
                                         (type;function (list (replace! caseT))
                                           (replace! currentT))))))
              
              ## else
              (&common;variant-out-of-bounds-error type expected-size tag)))

      _
      (&;throw Not-A-Variant-Type (%type type)))))