(.module:
  [lux (#- function)
   [control
    ["." monad (#+ do Monad)]
    [equivalence]
    ["p" parser]
    ["ex" exception (#+ exception:)]]
   ["." function]
   [data
    ["." product]
    ["." bit]
    ["." maybe]
    [name ("name/." Codec<Text,Name>)]
    ["e" error]
    ["." number ("nat/." Codec<Text,Nat>)]
    ["." text ("text/." Monoid<Text>)
     format]
    [collection
     ["." list ("list/." Fold<List> Monad<List> Monoid<List>)]
     ["dict" dictionary (#+ Dictionary)]]]
   ["." macro (#+ with-gensyms)
    ["." code]
    ["s" syntax (#+ syntax: Syntax)]
    [syntax
     ["cs" common]
     [common
      ["csr" reader]
      ["csw" writer]]]]
   ["." type ("type/." Equivalence<Type>)
    ["." check]]])

(do-template [<name>]
  [(exception: #export (<name> {type Type})
     (%type type))]

  [not-existential]
  [not-recursive]
  [not-named]
  [not-parameter]
  [unknown-parameter]
  [not-function]
  [not-application]
  [not-polymorphic]
  [not-variant]
  [not-tuple]
  )

(do-template [<name>]
  [(exception: #export (<name> {expected Type} {actual Type})
     (ex.report ["Expected" (%type expected)]
                ["Actual" (%type actual)]))]

  [types-do-not-match]
  [wrong-parameter]
  )

(exception: #export (unconsumed {remaining (List Type)})
  (ex.report ["Types" (|> remaining
                          (list/map (|>> %type (format "\n* ")))
                          (text.join-with ""))]))

(type: #export Env (Dictionary Nat [Type Code]))

(type: #export (Poly a)
  (p.Parser [Env (List Type)] a))

(def: #export fresh Env (dict.new number.Hash<Nat>))

(def: (run' env types poly)
  (All [a] (-> Env (List Type) (Poly a) (e.Error a)))
  (case (p.run [env types] poly)
    (#e.Error error)
    (#e.Error error)

    (#e.Success [[env' remaining] output])
    (case remaining
      #.Nil
      (#e.Success output)

      _
      (ex.throw unconsumed remaining))))

(def: #export (run type poly)
  (All [a] (-> Type (Poly a) (e.Error a)))
  (run' fresh (list type) poly))

(def: #export env
  (Poly Env)
  (.function (_ [env inputs])
    (#e.Success [[env inputs] env])))

(def: (with-env temp poly)
  (All [a] (-> Env (Poly a) (Poly a)))
  (.function (_ [env inputs])
    (case (p.run [temp inputs] poly)
      (#e.Error error)
      (#e.Error error)

      (#e.Success [[_ remaining] output])
      (#e.Success [[env remaining] output]))))

(def: #export peek
  (Poly Type)
  (.function (_ [env inputs])
    (case inputs
      #.Nil
      (#e.Error "Empty stream of types.")

      (#.Cons headT tail)
      (#e.Success [[env inputs] headT]))))

(def: #export any
  (Poly Type)
  (.function (_ [env inputs])
    (case inputs
      #.Nil
      (#e.Error "Empty stream of types.")

      (#.Cons headT tail)
      (#e.Success [[env tail] headT]))))

(def: #export (local types poly)
  (All [a] (-> (List Type) (Poly a) (Poly a)))
  (.function (_ [env pass-through])
    (case (run' env types poly)
      (#e.Error error)
      (#e.Error error)

      (#e.Success output)
      (#e.Success [[env pass-through] output]))))

(def: (label idx)
  (-> Nat Code)
  (code.local-identifier (text/compose "label\u0000" (nat/encode idx))))

(def: #export (with-extension type poly)
  (All [a] (-> Type (Poly a) (Poly [Code a])))
  (.function (_ [env inputs])
    (let [current-id (dict.size env)
          g!var (label current-id)]
      (case (p.run [(dict.put current-id [type g!var] env)
                    inputs]
                   poly)
        (#e.Error error)
        (#e.Error error)

        (#e.Success [[_ inputs'] output])
        (#e.Success [[env inputs'] [g!var output]])))))

(do-template [<name> <flattener> <tag> <exception>]
  [(def: #export (<name> poly)
     (All [a] (-> (Poly a) (Poly a)))
     (do p.Monad<Parser>
       [headT any]
       (let [members (<flattener> (type.un-name headT))]
         (if (n/> +1 (list.size members))
           (local members poly)
           (p.fail (ex.construct <exception> headT))))))]

  [variant type.flatten-variant #.Sum     not-variant]
  [tuple   type.flatten-tuple   #.Product not-tuple]
  )

(def: polymorphic'
  (Poly [Nat Type])
  (do p.Monad<Parser>
    [headT any
     #let [[num-arg bodyT] (type.flatten-univ-q (type.un-name headT))]]
    (if (n/= +0 num-arg)
      (p.fail (ex.construct not-polymorphic headT))
      (wrap [num-arg bodyT]))))

(def: #export (polymorphic poly)
  (All [a] (-> (Poly a) (Poly [Code (List Code) a])))
  (do p.Monad<Parser>
    [headT any
     funcI (:: @ map dict.size ..env)
     [num-args non-poly] (local (list headT) polymorphic')
     env ..env
     #let [funcL (label funcI)
           [all-varsL env'] (loop [current-arg +0
                                   env' env
                                   all-varsL (: (List Code) (list))]
                              (if (n/< num-args current-arg)
                                (if (n/= +0 current-arg)
                                  (let [varL (label (inc funcI))]
                                    (recur (inc current-arg)
                                           (|> env'
                                               (dict.put funcI [headT funcL])
                                               (dict.put (inc funcI) [(#.Parameter (inc funcI)) varL]))
                                           (#.Cons varL all-varsL)))
                                  (let [partialI (|> current-arg (n/* +2) (n/+ funcI))
                                        partial-varI (inc partialI)
                                        partial-varL (label partial-varI)
                                        partialC (` ((~ funcL) (~+ (|> (list.n/range +0 (dec num-args))
                                                                       (list/map (|>> (n/* +2) inc (n/+ funcI) label))
                                                                       list.reverse))))]
                                    (recur (inc current-arg)
                                           (|> env'
                                               (dict.put partialI [.Nothing partialC])
                                               (dict.put partial-varI [(#.Parameter partial-varI) partial-varL]))
                                           (#.Cons partial-varL all-varsL))))
                                [all-varsL env']))]]
    (|> (do @
          [output poly]
          (wrap [funcL all-varsL output]))
        (local (list non-poly))
        (with-env env'))))

(def: #export (function in-poly out-poly)
  (All [i o] (-> (Poly i) (Poly o) (Poly [i o])))
  (do p.Monad<Parser>
    [headT any
     #let [[inputsT outputT] (type.flatten-function (type.un-name headT))]]
    (if (n/> +0 (list.size inputsT))
      (p.and (local inputsT in-poly)
             (local (list outputT) out-poly))
      (p.fail (ex.construct not-function headT)))))

(def: #export (apply poly)
  (All [a] (-> (Poly a) (Poly a)))
  (do p.Monad<Parser>
    [headT any
     #let [[funcT paramsT] (type.flatten-application (type.un-name headT))]]
    (if (n/= +0 (list.size paramsT))
      (p.fail (ex.construct not-application headT))
      (local (#.Cons funcT paramsT) poly))))

(do-template [<name> <test>]
  [(def: #export (<name> expected)
     (-> Type (Poly Any))
     (do p.Monad<Parser>
       [actual any]
       (if (<test> expected actual)
         (wrap [])
         (p.fail (ex.construct types-do-not-match [expected actual])))))]

  [exactly type/=]
  [sub     check.checks?]
  [super   (function.flip check.checks?)]
  )

(def: (adjusted-idx env idx)
  (-> Env Nat Nat)
  (let [env-level (n// +2 (dict.size env))
        parameter-level (n// +2 idx)
        parameter-idx (n/% +2 idx)]
    (|> env-level dec (n/- parameter-level) (n/* +2) (n/+ parameter-idx))))

(def: #export parameter
  (Poly Code)
  (do p.Monad<Parser>
    [env ..env
     headT any]
    (case headT
      (#.Parameter idx)
      (case (dict.get (adjusted-idx env idx) env)
        (#.Some [poly-type poly-code])
        (wrap poly-code)

        #.None
        (p.fail (ex.construct unknown-parameter headT)))

      _
      (p.fail (ex.construct not-parameter headT)))))

(def: #export (parameter! id)
  (-> Nat (Poly Any))
  (do p.Monad<Parser>
    [env ..env
     headT any]
    (case headT
      (#.Parameter idx)
      (if (n/= id (adjusted-idx env idx))
        (wrap [])
        (p.fail (ex.construct wrong-parameter [(#.Parameter id) headT])))

      _
      (p.fail (ex.construct not-parameter headT)))))

(def: #export existential
  (Poly Nat)
  (do p.Monad<Parser>
    [headT any]
    (case headT
      (#.Ex ex-id)
      (wrap ex-id)

      _
      (p.fail (ex.construct not-existential headT)))))

(def: #export named
  (Poly [Name Type])
  (do p.Monad<Parser>
    [inputT any]
    (case inputT
      (#.Named name anonymousT)
      (wrap [name anonymousT])

      _
      (p.fail (ex.construct not-named inputT)))))

(def: #export (recursive poly)
  (All [a] (-> (Poly a) (Poly [Code a])))
  (do p.Monad<Parser>
    [headT any]
    (case (type.un-name headT)
      (#.Apply (#.Named ["lux" "Nothing"] _) (#.UnivQ _ headT'))
      (do @
        [[recT _ output] (|> poly
                             (with-extension .Nothing)
                             (with-extension headT)
                             (local (list headT')))]
        (wrap [recT output]))

      _
      (p.fail (ex.construct not-recursive headT)))))

(def: #export recursive-self
  (Poly Code)
  (do p.Monad<Parser>
    [env ..env
     headT any]
    (case (type.un-name headT)
      (^multi (#.Apply (#.Named ["lux" "Nothing"] _) (#.Parameter funcT-idx))
              (n/= +0 (adjusted-idx env funcT-idx))
              [(dict.get +0 env) (#.Some [self-type self-call])])
      (wrap self-call)

      _
      (p.fail (ex.construct not-recursive headT)))))

(def: #export recursive-call
  (Poly Code)
  (do p.Monad<Parser>
    [env ..env
     [funcT argsT] (apply (p.and any (p.many any)))
     _ (local (list funcT) (..parameter! +0))
     allC (let [allT (list& funcT argsT)]
            (|> allT
                (monad.map @ (function.constant ..parameter))
                (local allT)))]
    (wrap (` ((~+ allC))))))

(def: #export log
  (All [a] (Poly a))
  (do p.Monad<Parser>
    [current any
     #let [_ (log! ($_ text/compose
                       "{" (name/encode (name-of ..log)) "} "
                       (%type current)))]]
    (p.fail "LOGGING")))

## [Syntax]
(syntax: #export (poly: {export csr.export}
                   {name s.local-identifier}
                   body)
  (with-gensyms [g!_ g!type g!output]
    (let [g!name (code.identifier ["" name])]
      (wrap (.list (` (syntax: (~+ (csw.export export)) ((~ g!name) {(~ g!type) s.identifier})
                        (do macro.Monad<Meta>
                          [(~ g!type) (macro.find-type-def (~ g!type))]
                          (case (|> (~ body)
                                    (.function ((~ g!_) (~ g!name)))
                                    p.rec
                                    (do p.Monad<Parser> [])
                                    (..run (~ g!type))
                                    (: (.Either .Text .Code)))
                            (#.Left (~ g!output))
                            (macro.fail (~ g!output))

                            (#.Right (~ g!output))
                            ((~' wrap) (.list (~ g!output))))))))))))

(def: (common-poly-name? poly-func)
  (-> Text Bit)
  (text.contains? "?" poly-func))

(def: (derivation-name poly args)
  (-> Text (List Text) (Maybe Text))
  (if (common-poly-name? poly)
    (#.Some (list/fold (text.replace-once "?") poly args))
    #.None))

(syntax: #export (derived: {export csr.export}
                   {?name (p.maybe s.local-identifier)}
                   {[poly-func poly-args] (s.form (p.and s.identifier (p.many s.identifier)))}
                   {?custom-impl (p.maybe s.any)})
  (do @
    [poly-args (monad.map @ macro.normalize poly-args)
     name (case ?name
            (#.Some name)
            (wrap name)

            (^multi #.None
                    [(derivation-name (product.right poly-func) (list/map product.right poly-args))
                     (#.Some derived-name)])
            (wrap derived-name)

            _
            (p.fail "derived: was given no explicit name, and cannot generate one from given information."))
     #let [impl (case ?custom-impl
                  (#.Some custom-impl)
                  custom-impl

                  #.None
                  (` ((~ (code.identifier poly-func)) (~+ (list/map code.identifier poly-args)))))]]
    (wrap (.list (` (def: (~+ (csw.export export))
                      (~ (code.identifier ["" name]))
                      {#.struct? #1}
                      (~ impl)))))))

## [Derivers]
(def: #export (to-code env type)
  (-> Env Type Code)
  (case type
    (#.Primitive name params)
    (` (#.Primitive (~ (code.text name))
                    (list (~+ (list/map (to-code env) params)))))

    (^template [<tag>]
      (<tag> idx)
      (` (<tag> (~ (code.nat idx)))))
    ([#.Var] [#.Ex])

    (#.Parameter idx)
    (let [idx (adjusted-idx env idx)]
      (if (n/= +0 idx)
        (|> (dict.get idx env) maybe.assume product.left (to-code env))
        (` (.$ (~ (code.nat (dec idx)))))))

    (#.Apply (#.Named ["lux" "Nothing"] _) (#.Parameter idx))
    (let [idx (adjusted-idx env idx)]
      (if (n/= +0 idx)
        (|> (dict.get idx env) maybe.assume product.left (to-code env))
        (undefined)))
    
    (^template [<tag>]
      (<tag> left right)
      (` (<tag> (~ (to-code env left))
                (~ (to-code env right)))))
    ([#.Function] [#.Apply])

    (^template [<tag> <macro> <flattener>]
      (<tag> left right)
      (` (<macro> (~+ (list/map (to-code env) (<flattener> type))))))
    ([#.Sum  | type.flatten-variant]
     [#.Product & type.flatten-tuple])

    (#.Named name sub-type)
    (code.identifier name)

    (^template [<tag>]
      (<tag> scope body)
      (` (<tag> (list (~+ (list/map (to-code env) scope)))
                (~ (to-code env body)))))
    ([#.UnivQ] [#.ExQ])
    ))