(;module:
  lux
  (lux (control monad
                pipe)
       [io #- run]
       (concurrency ["A" atom])
       (data [text "T/" Eq<Text>]
             text/format
             [ident]
             (coll [list "L/" Fold<List> Monoid<List> Monad<List>]
                   ["D" dict]
                   ["S" set])
             [number]
             [product])
       [macro #+ Monad<Lux>]
       (macro [code])
       [type]
       (type ["TC" check]))
  (luxc ["&" base]
        (lang ["la" analysis #+ Analysis])
        ["&;" module]
        ["&;" env]
        (analyser ["&;" common]
                  ["&;" inference])))

## [Analysers]
(def: (analyse-typed-tuple analyse members)
  (-> &;Analyser (List Code) (Lux Analysis))
  (do Monad<Lux>
    [expected macro;expected-type]
    (let [member-types (type;flatten-tuple expected)
          num-types (list;size member-types)
          num-members (list;size members)]
      (cond (n.= num-types num-members)
            (do @
              [=tuple (: (Lux (List Analysis))
                         (mapM @
                               (function [[expected member]]
                                 (&;with-expected-type expected
                                   (analyse member)))
                               (list;zip2 member-types members)))]
              (wrap (#la;Tuple =tuple)))

            (n.< num-types num-members)
            (do @
              [#let [[head-ts tail-ts] (list;split (n.- +2 num-members)
                                                   member-types)]
               =prevs (mapM @
                            (function [[expected member]]
                              (&;with-expected-type expected
                                (analyse member)))
                            (list;zip2 head-ts members))
               =last (&;with-expected-type (type;tuple tail-ts)
                       (analyse (default (undefined)
                                  (list;last members))))]
              (wrap (#la;Tuple (L/append =prevs (list =last)))))

            ## (n.> num-types num-members)
            (do @
              [#let [[head-xs tail-xs] (list;split (n.- +2 num-types)
                                                   members)]
               =prevs (mapM @
                            (function [[expected member]]
                              (&;with-expected-type expected
                                (analyse member)))
                            (list;zip2 member-types head-xs))
               =last (&;with-expected-type (default (undefined)
                                             (list;last member-types))
                       (analyse-typed-tuple analyse tail-xs))]
              (wrap (#la;Tuple (L/append =prevs (list =last)))))
            ))))

(def: #export (normalize-record pairs)
  (-> (List [Code Code]) (Lux (List [Ident Code])))
  (mapM Monad<Lux>
        (function [[key val]]
          (case key
            [_ (#;Tag key)]
            (do Monad<Lux>
              [key (macro;normalize key)]
              (wrap [key val]))

            _
            (&;fail (format "Cannot use non-tag tokens in key positions in records: " (%code key)))))
        pairs))

(def: #export (order-record pairs)
  (-> (List [Ident Code]) (Lux [(List Code) Type]))
  (case pairs
    (#;Cons [head-k head-v] _)
    (do Monad<Lux>
      [head-k (macro;normalize head-k)
       [_ tag-set recordT] (macro;resolve-tag head-k)
       #let [size-record (list;size pairs)
             size-ts (list;size tag-set)]
       _ (if (n.= size-ts size-record)
           (wrap [])
           (&;fail (format "Record size does not match tag-set size." "\n"
                           "Expected: " (|> size-ts nat-to-int %i) "\n"
                           "  Actual: " (|> size-record nat-to-int %i) "\n"
                           "For type: " (%type recordT))))
       #let [tuple-range (list;n.range +0 size-ts)
             tag->idx (D;from-list ident;Hash<Ident> (list;zip2 tag-set tuple-range))]
       idx->val (foldM @
                       (function [[key val] idx->val]
                         (do @
                           [key (macro;normalize key)]
                           (case (D;get key tag->idx)
                             #;None
                             (&;fail (format "Tag " (%code (code;tag key))
                                             " does not belong to tag-set for type " (%type recordT)))

                             (#;Some idx)
                             (if (D;contains? idx idx->val)
                               (&;fail (format "Cannot repeat tag inside record: " (%code (code;tag key))))
                               (wrap (D;put idx val idx->val))))))
                       (: (D;Dict Nat Code)
                          (D;new number;Hash<Nat>))
                       pairs)
       #let [ordered-tuple (L/map (function [idx]
                                    (assume (D;get idx idx->val)))
                                  tuple-range)]]
      (wrap [ordered-tuple recordT]))

    _
    (:: Monad<Lux> wrap [(list) Unit])))

(def: #export (analyse-tuple analyse members)
  (-> &;Analyser (List Code) (Lux Analysis))
  (do Monad<Lux>
    [expected macro;expected-type]
    (&;with-stacked-errors
      (function [_] (format "Invalid type for tuple: " (%type expected)))
      (case expected
        (#;Product _)
        (analyse-typed-tuple analyse members)

        (#;Named name unnamedT)
        (&;with-expected-type unnamedT
          (analyse-tuple analyse members))

        (#;Var id)
        (do @
          [bound? (&;within-type-env
                   (TC;bound? id))]
          (if bound?
            (do @
              [expected' (&;within-type-env
                          (TC;read-var id))]
              (&;with-expected-type expected'
                (analyse-tuple analyse members)))
            (do @
              [=members (mapM @ (|>. analyse &common;with-unknown-type)
                              members)
               #let [tuple-type (type;tuple (L/map product;left =members))]
               _ (&;within-type-env
                  (TC;check expected tuple-type))]
              (wrap (#la;Tuple (L/map product;right =members))))))

        (#;UnivQ _)
        (do @
          [[var-id var] (&;within-type-env
                         TC;existential)]
          (&;with-expected-type (assume (type;apply-type expected var))
            (analyse-tuple analyse members)))

        (#;ExQ _)
        (&common;with-var
          (function [[var-id var]]
            (&;with-expected-type (assume (type;apply-type expected var))
              (analyse-tuple analyse members))))
        
        _
        (&;fail "")
        ))))

(def: (record-function-type type)
  (-> Type (Lux Type))
  (case type
    (#;Named name unnamedT)
    (do Monad<Lux>
      [unnamedT+ (record-function-type unnamedT)]
      (wrap (#;Named name unnamedT+)))

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

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

    _
    (&;fail (format "Not a record type: " (%type type)))))

(def: (out-of-bounds-error type size tag)
  (All [a] (-> Type Nat Nat (Lux a)))
  (&;fail (format "Trying to create variant with tag beyond type's limitations." "\n"
                  " Tag: " (%i (nat-to-int tag)) "\n"
                  "Size: " (%i (nat-to-int size)) "\n"
                  "Type: " (%type type))))

(def: (variant-function-type tag expected-size type)
  (-> Nat Nat Type (Lux Type))
  (case type
    (#;Named name unnamedT)
    (do Monad<Lux>
      [unnamedT+ (record-function-type unnamedT)]
      (wrap (#;Named name unnamedT+)))

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

    (#;Sum _)
    (let [cases (type;flatten-variant type)
          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<Lux> wrap (type;function (list caseT) type))

              #;None
              (out-of-bounds-error type expected-size tag))
            
            (n.< expected-size actual-size)
            (&;fail (format "Variant type is smaller than expected." "\n"
                            "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<Lux> wrap (type;function (list caseT) type)))
            
            ## else
            (out-of-bounds-error type expected-size tag)))

    _
    (&;fail (format "Not a variant type: " (%type type)))))

(def: #export (analyse-record analyse members)
  (-> &;Analyser (List [Code Code]) (Lux Analysis))
  (do Monad<Lux>
    [members (normalize-record members)
     [members recordT] (order-record members)
     expectedT macro;expected-type
     functionT (record-function-type recordT)
     [inferredT membersA] (&inference;apply-function analyse functionT members)
     _ (&;within-type-env
        (TC;check expectedT inferredT))]
    (wrap (#la;Tuple membersA))))

(def: #export (analyse-tagged-variant analyse tag value)
  (-> &;Analyser Ident Code (Lux Analysis))
  (do Monad<Lux>
    [tag (macro;normalize tag)
     [idx group variantT] (macro;resolve-tag tag)
     #let [case-size (list;size group)]
     functionT (variant-function-type idx case-size variantT)
     [inferredT valueA+] (&inference;apply-function analyse functionT (list value))
     expectedT macro;expected-type
     _ (&;within-type-env
        (TC;check expectedT inferredT))]
    (wrap (#la;Variant idx case-size (|> valueA+ list;head assume)))))

(def: #export (analyse-variant analyse tag value)
  (-> &;Analyser Nat Code (Lux Analysis))
  (do Monad<Lux>
    [expected macro;expected-type]
    (&;with-stacked-errors
      (function [_] (format "Invalid type for variant: " (%type expected)))
      (case expected
        (#;Sum _)
        (let [flat (type;flatten-variant expected)
              type-size (list;size flat)]
          (case (list;nth tag flat)
            (#;Some variant-type)
            (do @
              [=value (&;with-expected-type variant-type
                        (analyse value))]
              (wrap (#la;Variant tag type-size =value)))

            #;None
            (out-of-bounds-error expected type-size tag)))

        (#;Named name unnamedT)
        (&;with-expected-type unnamedT
          (analyse-variant analyse tag value))

        (#;Var id)
        (do @
          [bound? (&;within-type-env
                   (TC;bound? id))]
          (if bound?
            (do @
              [expected' (&;within-type-env
                          (TC;read-var id))]
              (&;with-expected-type expected'
                (analyse-variant analyse tag value)))
            (&;fail (format "Invalid type for variant: " (%type expected)))))

        (#;UnivQ _)
        (do @
          [[var-id var] (&;within-type-env
                         TC;existential)]
          (&;with-expected-type (assume (type;apply-type expected var))
            (analyse-variant analyse tag value)))

        (#;ExQ _)
        (&common;with-var
          (function [[var-id var]]
            (&;with-expected-type (assume (type;apply-type expected var))
              (analyse-variant analyse tag value))))
        
        _
        (&;fail "")))))