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

(do-template [<name>]
  [(exception: #export (<name> {message Text})
     message)]

  [Cannot-Infer]
  [Cannot-Infer-Argument]
  [Smaller-Variant-Than-Expected]
  [Invalid-Type-Application]
  [Not-A-Record-Type]
  [Not-A-Variant-Type]
  )

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

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

(def: new-named-type
  (Meta Type)
  (do macro.Monad<Meta>
    [[_module _line _column] macro.cursor
     [ex-id exT] (&.with-type-env tc.existential)]
    (wrap (#.Primitive (format "{New Type @ " (%t _module)
                               "," (%n _line)
                               "," (%n _column)
                               "} " (%n ex-id))
                       (list)))))

## 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>
        [[var-id varT] (&.with-type-env tc.var)
         output (general analyse
                         (maybe.assume (type.apply (list varT) inferT))
                         args)
         bound? (&.with-type-env
                  (tc.bound? var-id))
         _ (if bound?
             (wrap [])
             (do @
               [newT new-named-type]
               (&.with-type-env
                 (tc.check varT newT))))]
        (wrap output))

      (#.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 (_ _)
                  (ex.construct 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)))))