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

(exception: #export Cannot-Infer)
(def: (cannot-infer type args)
  (-> Type (List Code) Text)
  (format "     Type: " (%type type) "\n"
          "Arguments:"
          (|> args
              list;enumerate
              (list/map (function [[idx argC]]
                          (format "\n  " (%n idx) " " (%code argC))))
              (text;join-with ""))))

(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)

(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 (general analyse inferT args)
  (-> &;Analyser Type (List Code) (Meta [Type (List Analysis)]))
  (case args
    #;Nil
    (do macro;Monad<Meta>
      [_ (&;infer inferT)]
      (wrap [inferT (list)]))
    
    (#;Cons argC args')
    (case inferT
      (#;Named name unnamedT)
      (general analyse unnamedT args)

      (#;UnivQ _)
      (do macro;Monad<Meta>
        [[var-id varT] (&;with-type-env tc;var)]
        (general analyse (maybe;assume (type;apply (list varT) inferT)) args))

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

      (#;Apply inputT transT)
      (case (type;apply (list inputT) transT)
        (#;Some outputT)
        (general 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 macro;Monad<Meta>
        [[outputT' args'A] (general analyse outputT args')
         argA (&;with-stacked-errors
                (function [_] (Cannot-Infer-Argument
                               (format "Inferred Type: " (%type inputT) "\n"
                                       "     Argument: " (%code argC))))
                (&;with-type inputT
                  (analyse argC)))]
        (wrap [outputT' (list& argA args'A)]))

      (#;Var infer-id)
      (do macro;Monad<Meta>
        [?inferT' (&;with-type-env (tc;read infer-id))]
        (case ?inferT'
          (#;Some inferT')
          (general analyse inferT' args)

          _
          (&;throw Cannot-Infer (cannot-infer inferT args))))

      _
      (&;throw Cannot-Infer (cannot-infer inferT args)))
    ))

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

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

    (#;Apply inputT funcT)
    (case (type;apply (list inputT) funcT)
      (#;Some outputT)
      (record outputT)

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

    (#;Product _)
    (macro/wrap (type;function (type;flatten-tuple inferT) inferT))

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

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

      (^template [<tag>]
        (<tag> env bodyT)
        (do macro;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)
                (macro/wrap (if (n.= +0 depth)
                              (type;function (list caseT) currentT)
                              (let [replace! (replace-bound (|> depth n.dec (n.* +2)) inferT)]
                                (type;function (list (replace! caseT))
                                  (replace! currentT)))))

                #;None
                (&common;variant-out-of-bounds-error inferT 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))]
                (macro/wrap (if (n.= +0 depth)
                              (type;function (list caseT) currentT)
                              (let [replace! (replace-bound (|> depth n.dec (n.* +2)) inferT)]
                                (type;function (list (replace! caseT))
                                  (replace! currentT))))))
              
              ## else
              (&common;variant-out-of-bounds-error inferT expected-size tag)))

      (#;Apply inputT funcT)
      (case (type;apply (list inputT) funcT)
        (#;Some outputT)
        (variant tag expected-size outputT)

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

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