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

(do-template [<name> <type> <tag>]
  [(def: #export (<name> cursor value)
     (-> Cursor <type> (Lux Analysis))
     (do Monad<Lux>
       [expected macro;expected-type
        _ (&;within-type-env
           (TC;check expected <type>))]
       (wrap [[expected cursor]
              (<tag> value)])))]

  [analyse-bool Bool #la;Bool]
  [analyse-nat  Nat  #la;Nat]
  [analyse-int  Int  #la;Int]
  [analyse-deg  Deg  #la;Deg]
  [analyse-real Real #la;Real]
  [analyse-char Char #la;Char]
  [analyse-text Text #la;Text]
  )

(def: #export (analyse-unit cursor)
  (-> Cursor (Lux Analysis))
  (do Monad<Lux>
    [expected macro;expected-type
     _ (&;within-type-env
        (TC;check expected Unit))]
    (wrap [[expected cursor]
           #la;Unit])))

(def: (analyse-definition cursor def-name)
  (-> Cursor Ident (Lux Analysis))
  (do Monad<Lux>
    [actual (macro;find-def-type def-name)
     expected macro;expected-type
     _ (&;within-type-env
        (TC;check expected actual))]
    (wrap [[expected cursor]
           (#la;Absolute def-name)])))

(def: (analyse-variable cursor var-name)
  (-> Cursor Text (Lux (Maybe Analysis)))
  (do Monad<Lux>
    [?var (&env;find var-name)]
    (case ?var
      (#;Some [actual ref])
      (do @
        [expected macro;expected-type
         _ (&;within-type-env
            (TC;check expected actual))
         #let [analysis [[expected cursor]
                         (#la;Relative ref)]]]
        (wrap (#;Some analysis)))

      #;None
      (wrap #;None))))

(def: #export (analyse-reference cursor reference)
  (-> Cursor Ident (Lux Analysis))
  (case reference
    ["" simple-name]
    (do Monad<Lux>
      [?var (analyse-variable cursor simple-name)]
      (case ?var
        (#;Some analysis)
        (wrap analysis)

        #;None
        (do @
          [this-module macro;current-module-name]
          (analyse-definition cursor [this-module simple-name]))))

    _
    (analyse-definition cursor reference)))

(def: #export (analyse-check analyse eval cursor type value)
  (-> &;Analyser &;Eval Cursor Code Code (Lux Analysis))
  (do Monad<Lux>
    [actual (eval Type type)
     #let [actual (:! Type actual)]
     expected macro;expected-type
     _ (&;within-type-env
        (TC;check expected actual))]
    (&;with-expected-type actual
      (analyse eval value))))

(def: #export (analyse-coerce analyse eval cursor type value)
  (-> &;Analyser &;Eval Cursor Code Code (Lux Analysis))
  (do Monad<Lux>
    [actual (eval Type type)
     #let [actual (:! Type actual)]
     expected macro;expected-type
     _ (&;within-type-env
        (TC;check expected actual))
     =value (&;with-expected-type Top
              (analyse eval value))]
    (wrap (&common;replace-type expected =value))))

(def: (analyse-typed-tuple analyse cursor members)
  (-> (-> Code (Lux Analysis)) Cursor
      (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 [[expected cursor]
                     (#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 [[expected cursor]
                     (#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 cursor tail-xs))]
              (wrap [[expected cursor]
                     (#la;Tuple (L/append =prevs (list =last)))]))
            ))))

(def: (tuple cursor members)
  (-> Cursor (List Analysis) Analysis)
  (let [tuple-type (type;tuple (L/map la;get-type members))]
    [[tuple-type cursor]
     (#la;Tuple members)]))

(def: #export (analyse-tuple analyse cursor members)
  (-> (-> Code (Lux Analysis)) Cursor
      (List Code)
      (Lux Analysis))
  (do Monad<Lux>
    [expected macro;expected-type]
    (case expected
      (#;Product _)
      (analyse-typed-tuple analyse cursor members)

      (#;Var id)
      (do @
        [bound? (&;within-type-env
                 (TC;bound? id))]
        (if bound?
          (do @
            [expected' (&;within-type-env
                        (TC;read-var id))
             =tuple (&;with-expected-type expected'
                      (analyse-tuple analyse cursor members))]
            (wrap (&common;replace-type expected =tuple)))
          (do @
            [=members (mapM @ (<|. &common;with-unknown-type
                                   analyse)
                            members)
             #let [=tuple (tuple cursor =members)]
             _ (&;within-type-env
                (TC;check expected (la;get-type =tuple)))]
            (wrap (&common;replace-type expected =tuple)))))

      _
      (if (type;quantified? expected)
        (do @
          [[bindings expected'] (&;within-type-env
                                 (&common;realize expected))
           =tuple (&;with-expected-type expected'
                    (analyse-tuple analyse cursor members))
           =tuple (foldM @ &common;clean =tuple bindings)
           _ (&;within-type-env
              (TC;check expected (la;get-type =tuple)))]
          (wrap (&common;replace-type expected =tuple)))
        (&;fail (format "Invalid type for tuple: " (%type expected))))
      )))

(def: #export (analyse-variant analyse cursor tag value)
  (-> (-> Code (Lux Analysis)) Cursor
      Nat Code
      (Lux Analysis))
  (do Monad<Lux>
    [expected macro;expected-type]
    (case expected
      (#;Sum _)
      (let [flat (type;flatten-variant expected)
            type-size (list;size flat)]
        (if (n.< type-size tag)
          (do @
            [#let [last? (n.= tag (n.dec type-size))
                   variant-type (default (undefined)
                                  (list;nth tag flat))]
             =value (&;with-expected-type variant-type
                      (analyse value))]
            (wrap [[expected cursor]
                   (#la;Variant tag last? =value)]))
          (&;fail (format "Trying to create variant with tag beyond type's limitations." "\n"
                          "      Tag: " (%n tag) "\n"
                          "Type size: " (%n type-size) "\n"
                          "     Type: " (%type expected) "\n"))))

      _
      (if (type;quantified? expected)
        (do @
          [[bindings expected'] (&;within-type-env
                                 (&common;realize expected))
           =variant (&;with-expected-type expected'
                      (analyse-variant analyse cursor tag value))
           =variant (foldM @ &common;clean =variant bindings)
           _ (&;within-type-env
              (TC;check expected (la;get-type =variant)))]
          (wrap (&common;replace-type expected =variant)))
        (&;fail (format "Invalid type for variant: " (%type expected)))))))

## Functions
(def: (maybe-to-lux input)
  (All [a] (-> (Maybe a) (Lux a)))
  (case input
    #;None
    (&;fail "")

    (#;Some value)
    (:: Monad<Lux> wrap value)))

(def: (with-var body)
  (All [a] (-> (-> [Nat Type] (Lux a)) (Lux a)))
  (do Monad<Lux>
    [[id var] (&;within-type-env TC;create-var)
     output (body [id var])
     _ (&;within-type-env (TC;delete-var id))]
    (wrap output)))

(def: (bind-var var-id bound-idx type)
  (-> Nat Nat Type Type)
  (case type
    (#;Host name params)
    (#;Host name (L/map (bind-var var-id bound-idx) params))

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

    (^template [<tag>]
      (<tag> env quantified)
      (<tag> (L/map (bind-var var-id bound-idx) env)
             (bind-var var-id (n.+ +2 bound-idx) quantified)))
    ([#;UnivQ]
     [#;ExQ])
    
    (#;Named name unnamed)
    (#;Named name
             (bind-var var-id bound-idx unnamed))

    _
    type))

(def: #export (analyse-function analyse cursor func-name arg-name body)
  (-> (-> Code (Lux Analysis)) Cursor
      Text Text Code
      (Lux Analysis))
  (do Monad<Lux>
    [expected macro;expected-type]
    (&;with-try
      (function [error]
        (let [raw (format "Functions require function types: " (type;to-text expected))]
          (&;fail (if (T/= "" error)
                    raw
                    (format error "\n" raw)))))
      (case expected
        (#;Named name unnamed)
        (do @
          [=function (&;with-expected-type unnamed
                       (analyse-function analyse cursor func-name arg-name body))]
          (wrap (&common;replace-type expected =function)))

        (#;App funT argT)
        (do @
          [fully-applied (maybe-to-lux (type;apply-type funT argT))
           =function (&;with-expected-type fully-applied
                       (analyse-function analyse cursor func-name arg-name body))]
          (wrap (&common;replace-type expected =function)))
        
        (#;UnivQ _)
        (do @
          [[var-id var] (&;within-type-env
                         TC;existential)
           expected' (maybe-to-lux (type;apply-type expected var))
           =function (&;with-expected-type expected'
                       (analyse-function analyse cursor func-name arg-name body))]
          (wrap (&common;replace-type expected =function)))

        (#;ExQ _)
        (with-var
          (function [[var-id var]]
            (do @
              [expected' (maybe-to-lux (type;apply-type expected var))
               =function (&;with-expected-type expected'
                           (analyse-function analyse cursor func-name arg-name body))]
              (&common;clean var =function))))
        
        (#;Var id)
        (do @
          [? (&;within-type-env
              (TC;bound? id))]
          (if ?
            (do @
              [expected' (&;within-type-env
                          (TC;read-var id))]
              (&;with-expected-type expected'
                (analyse-function analyse cursor func-name arg-name body)))
            ## Inference
            (with-var
              (function [[input-id inputT]]
                (with-var
                  (function [[output-id outputT]]
                    (do @
                      [#let [funT (#;Function inputT outputT)]
                       =function (&;with-expected-type funT
                                   (analyse-function analyse cursor func-name arg-name body))
                       funT' (&;within-type-env
                              (TC;clean output-id funT))
                       concrete-input? (&;within-type-env
                                        (TC;bound? input-id))
                       funT'' (if concrete-input?
                                (&;within-type-env
                                 (TC;clean input-id funT'))
                                (wrap (#;UnivQ (list) (bind-var input-id +1 funT'))))
                       _ (&;within-type-env
                          (TC;check expected funT''))]
                      (wrap (&common;replace-type expected =function)))
                    ))))))

        (#;Function inputT outputT)
        (do @
          [[=scope =body] (&;with-scope
                            (&env;with-local [func-name expected]
                              (&env;with-local [arg-name inputT]
                                (&;with-expected-type outputT
                                  (analyse body)))))]
          (wrap [[expected cursor]
                 (#la;Function =scope =body)]))
        
        _
        (&;fail "")
        ))))