(.module:
  [lux #*
   [abstract
    ["." monad (#+ do)]]
   [control
    ["ex" exception (#+ exception:)]
    ["." state]]
   [data
    ["." name]
    ["." product]
    ["." maybe]
    [number
     ["n" nat]]
    [text
     ["%" format (#+ format)]]
    [collection
     ["." list ("#@." functor)]
     ["." dictionary (#+ Dictionary)]]]
   ["." type
    ["." check]]
   ["." macro
    ["." code]]]
  ["." // #_
   ["#." type]
   ["#." primitive]
   ["#." inference]
   ["#/" //
    ["#." extension]
    [//
     ["/" analysis (#+ Tag Analysis Operation Phase)]]]])

(exception: #export (invalid-variant-type {type Type} {tag Tag} {code Code})
  (ex.report ["Type" (%.type type)]
             ["Tag" (%.nat tag)]
             ["Expression" (%.code code)]))

(template [<name>]
  [(exception: #export (<name> {type Type} {members (List Code)})
     (ex.report ["Type" (%.type type)]
                ["Expression" (%.code (` [(~+ members)]))]))]

  [invalid-tuple-type]
  [cannot-analyse-tuple]
  )

(exception: #export (not-a-quantified-type {type Type})
  (%.type type))

(template [<name>]
  [(exception: #export (<name> {type Type} {tag Tag} {code Code})
     (ex.report ["Type" (%.type type)]
                ["Tag" (%.nat tag)]
                ["Expression" (%.code code)]))]

  [cannot-analyse-variant]
  [cannot-infer-numeric-tag]
  )

(exception: #export (record-keys-must-be-tags {key Code} {record (List [Code Code])})
  (ex.report ["Key" (%.code key)]
             ["Record" (%.code (code.record record))]))

(template [<name>]
  [(exception: #export (<name> {key Name} {record (List [Name Code])})
     (ex.report ["Tag" (%.code (code.tag key))]
                ["Record" (%.code (code.record (list@map (function (_ [keyI valC])
                                                           [(code.tag keyI) valC])
                                                         record)))]))]

  [cannot-repeat-tag]
  )

(exception: #export (tag-does-not-belong-to-record {key Name} {type Type})
  (ex.report ["Tag" (%.code (code.tag key))]
             ["Type" (%.type type)]))

(exception: #export (record-size-mismatch {expected Nat} {actual Nat} {type Type} {record (List [Name Code])})
  (ex.report ["Expected" (%.nat expected)]
             ["Actual" (%.nat actual)]
             ["Type" (%.type type)]
             ["Expression" (%.code (|> record
                                       (list@map (function (_ [keyI valueC])
                                                   [(code.tag keyI) valueC]))
                                       code.record))]))

(def: #export (sum analyse tag valueC)
  (-> Phase Nat Code (Operation Analysis))
  (do ///.monad
    [expectedT (///extension.lift macro.expected-type)
     expectedT' (//type.with-env
                  (check.clean expectedT))]
    (/.with-stack cannot-analyse-variant [expectedT' tag valueC]
      (case expectedT
        (#.Sum _)
        (let [flat (type.flatten-variant expectedT)
              type-size (list.size flat)
              right? (n.= (dec type-size)
                          tag)
              lefts (if right?
                      (dec tag)
                      tag)]
          (case (list.nth tag flat)
            (#.Some variant-type)
            (do @
              [valueA (//type.with-type variant-type
                        (analyse valueC))]
              (wrap (/.variant [lefts right? valueA])))

            #.None
            (/.throw //inference.variant-tag-out-of-bounds [type-size tag expectedT])))

        (#.Named name unnamedT)
        (//type.with-type unnamedT
          (sum analyse tag valueC))

        (#.Var id)
        (do @
          [?expectedT' (//type.with-env
                         (check.read id))]
          (case ?expectedT'
            (#.Some expectedT')
            (//type.with-type expectedT'
              (sum analyse tag valueC))

            _
            ## Cannot do inference when the tag is numeric.
            ## This is because there is no way of knowing how many
            ## cases the inferred sum type would have.
            (/.throw cannot-infer-numeric-tag [expectedT tag valueC])
            ))

        (^template [<tag> <instancer>]
          (<tag> _)
          (do @
            [[instance-id instanceT] (//type.with-env <instancer>)]
            (//type.with-type (maybe.assume (type.apply (list instanceT) expectedT))
              (sum analyse tag valueC))))
        ([#.UnivQ check.existential]
         [#.ExQ check.var])

        (#.Apply inputT funT)
        (case funT
          (#.Var funT-id)
          (do @
            [?funT' (//type.with-env (check.read funT-id))]
            (case ?funT'
              (#.Some funT')
              (//type.with-type (#.Apply inputT funT')
                (sum analyse tag valueC))

              _
              (/.throw invalid-variant-type [expectedT tag valueC])))

          _
          (case (type.apply (list inputT) funT)
            (#.Some outputT)
            (//type.with-type outputT
              (sum analyse tag valueC))

            #.None
            (/.throw not-a-quantified-type funT)))
        
        _
        (/.throw invalid-variant-type [expectedT tag valueC])))))

(def: (typed-product analyse members)
  (-> Phase (List Code) (Operation Analysis))
  (do ///.monad
    [expectedT (///extension.lift macro.expected-type)
     membersA+ (: (Operation (List Analysis))
                  (loop [membersT+ (type.flatten-tuple expectedT)
                         membersC+ members]
                    (case [membersT+ membersC+]
                      [(#.Cons memberT #.Nil) _]
                      (//type.with-type memberT
                        (:: @ map (|>> list) (analyse (code.tuple membersC+))))
                      
                      [_ (#.Cons memberC #.Nil)]
                      (//type.with-type (type.tuple membersT+)
                        (:: @ map (|>> list) (analyse memberC)))
                      
                      [(#.Cons memberT membersT+') (#.Cons memberC membersC+')]
                      (do @
                        [memberA (//type.with-type memberT
                                   (analyse memberC))
                         memberA+ (recur membersT+' membersC+')]
                        (wrap (#.Cons memberA memberA+)))

                      _
                      (/.throw cannot-analyse-tuple [expectedT members]))))]
    (wrap (/.tuple membersA+))))

(def: #export (product analyse membersC)
  (-> Phase (List Code) (Operation Analysis))
  (do ///.monad
    [expectedT (///extension.lift macro.expected-type)]
    (/.with-stack cannot-analyse-tuple [expectedT membersC]
      (case expectedT
        (#.Product _)
        (..typed-product analyse membersC)

        (#.Named name unnamedT)
        (//type.with-type unnamedT
          (product analyse membersC))

        (#.Var id)
        (do @
          [?expectedT' (//type.with-env
                         (check.read id))]
          (case ?expectedT'
            (#.Some expectedT')
            (//type.with-type expectedT'
              (product analyse membersC))

            _
            ## Must do inference...
            (do @
              [membersTA (monad.map @ (|>> analyse //type.with-inference)
                                    membersC)
               _ (//type.with-env
                   (check.check expectedT
                                (type.tuple (list@map product.left membersTA))))]
              (wrap (/.tuple (list@map product.right membersTA))))))

        (^template [<tag> <instancer>]
          (<tag> _)
          (do @
            [[instance-id instanceT] (//type.with-env <instancer>)]
            (//type.with-type (maybe.assume (type.apply (list instanceT) expectedT))
              (product analyse membersC))))
        ([#.UnivQ check.existential]
         [#.ExQ check.var])

        (#.Apply inputT funT)
        (case funT
          (#.Var funT-id)
          (do @
            [?funT' (//type.with-env (check.read funT-id))]
            (case ?funT'
              (#.Some funT')
              (//type.with-type (#.Apply inputT funT')
                (product analyse membersC))

              _
              (/.throw invalid-tuple-type [expectedT membersC])))

          _
          (case (type.apply (list inputT) funT)
            (#.Some outputT)
            (//type.with-type outputT
              (product analyse membersC))

            #.None
            (/.throw not-a-quantified-type funT)))
        
        _
        (/.throw invalid-tuple-type [expectedT membersC])
        ))))

(def: #export (tagged-sum analyse tag valueC)
  (-> Phase Name Code (Operation Analysis))
  (do ///.monad
    [tag (///extension.lift (macro.normalize tag))
     [idx group variantT] (///extension.lift (macro.resolve-tag tag))
     expectedT (///extension.lift macro.expected-type)]
    (case expectedT
      (#.Var _)
      (do @
        [#let [case-size (list.size group)]
         inferenceT (//inference.variant idx case-size variantT)
         [inferredT valueA+] (//inference.general analyse inferenceT (list valueC))
         #let [right? (n.= (dec case-size) idx)
               lefts (if right?
                       (dec idx)
                       idx)]]
        (wrap (/.variant [lefts right? (|> valueA+ list.head maybe.assume)])))

      _
      (..sum analyse idx valueC))))

## There cannot be any ambiguity or improper syntax when analysing
## records, so they must be normalized for further analysis.
## Normalization just means that all the tags get resolved to their
## canonical form (with their corresponding module identified).
(def: #export (normalize record)
  (-> (List [Code Code]) (Operation (List [Name Code])))
  (monad.map ///.monad
             (function (_ [key val])
               (case key
                 [_ (#.Tag key)]
                 (do ///.monad
                   [key (///extension.lift (macro.normalize key))]
                   (wrap [key val]))

                 _
                 (/.throw record-keys-must-be-tags [key record])))
             record))

## Lux already possesses the means to analyse tuples, so
## re-implementing the same functionality for records makes no sense.
## Records, thus, get transformed into tuples by ordering the elements.
(def: #export (order record)
  (-> (List [Name Code]) (Operation [(List Code) Type]))
  (case record
    ## empty-record = empty-tuple = unit = []
    #.Nil
    (:: ///.monad wrap [(list) Any])

    (#.Cons [head-k head-v] _)
    (do ///.monad
      [head-k (///extension.lift (macro.normalize head-k))
       [_ tag-set recordT] (///extension.lift (macro.resolve-tag head-k))
       #let [size-record (list.size record)
             size-ts (list.size tag-set)]
       _ (if (n.= size-ts size-record)
           (wrap [])
           (/.throw record-size-mismatch [size-ts size-record recordT record]))
       #let [tuple-range (list.indices size-ts)
             tag->idx (dictionary.from-list name.hash (list.zip2 tag-set tuple-range))]
       idx->val (monad.fold @
                            (function (_ [key val] idx->val)
                              (do @
                                [key (///extension.lift (macro.normalize key))]
                                (case (dictionary.get key tag->idx)
                                  (#.Some idx)
                                  (if (dictionary.contains? idx idx->val)
                                    (/.throw cannot-repeat-tag [key record])
                                    (wrap (dictionary.put idx val idx->val)))

                                  #.None
                                  (/.throw tag-does-not-belong-to-record [key recordT]))))
                            (: (Dictionary Nat Code)
                               (dictionary.new n.hash))
                            record)
       #let [ordered-tuple (list@map (function (_ idx) (maybe.assume (dictionary.get idx idx->val)))
                                     tuple-range)]]
      (wrap [ordered-tuple recordT]))
    ))

(def: #export (record analyse members)
  (-> Phase (List [Code Code]) (Operation Analysis))
  (case members
    (^ (list))
    //primitive.unit

    (^ (list [_ singletonC]))
    (analyse singletonC)

    _
    (do ///.monad
      [members (normalize members)
       [membersC recordT] (order members)
       expectedT (///extension.lift macro.expected-type)]
      (case expectedT
        (#.Var _)
        (do @
          [inferenceT (//inference.record recordT)
           [inferredT membersA] (//inference.general analyse inferenceT membersC)]
          (wrap (/.tuple membersA)))

        _
        (..product analyse membersC)))))