(.module: {#.doc "Basic functionality for working with types."}
  [lux (#- function)
   [control
    [equivalence (#+ Equivalence)]
    [monad (#+ Monad do)]
    ["p" parser]]
   [data
    ["." text ("#/." monoid equivalence)]
    ["." name ("#/." equivalence codec)]
    [number
     ["." nat ("#/." decimal)]]
    ["." maybe]
    [collection
     ["." array]
     ["." list ("#/." functor monoid fold)]]]
   ["." macro
    ["." code]
    ["s" syntax (#+ Syntax syntax:)]]])

(def: (beta-reduce env type)
  (-> (List Type) Type Type)
  (case type
    (#.Primitive name params)
    (#.Primitive name (list/map (beta-reduce env) params))
    
    (^template [<tag>]
      (<tag> left right)
      (<tag> (beta-reduce env left) (beta-reduce env right)))
    ([#.Sum]      [#.Product]
     [#.Function] [#.Apply])
    
    (^template [<tag>]
      (<tag> old-env def)
      (case old-env
        #.Nil
        (<tag> env def)

        _
        (<tag> (list/map (beta-reduce env) old-env) def)))
    ([#.UnivQ]
     [#.ExQ])
    
    (#.Parameter idx)
    (maybe.default (error! (text/compose "Unknown type var: " (nat/encode idx)))
                   (list.nth idx env))
    
    _
    type
    ))

(structure: #export equivalence (Equivalence Type)
  (def: (= x y)
    (case [x y]
      [(#.Primitive xname xparams) (#.Primitive yname yparams)]
      (and (text/= xname yname)
           (n/= (list.size yparams) (list.size xparams))
           (list/fold (.function (_ [x y] prev) (and prev (= x y)))
                      #1
                      (list.zip2 xparams yparams)))

      (^template [<tag>]
        [(<tag> xid) (<tag> yid)]
        (n/= yid xid))
      ([#.Var] [#.Ex] [#.Parameter])

      (^or [(#.Function xleft xright) (#.Function yleft yright)]
           [(#.Apply xleft xright) (#.Apply yleft yright)])
      (and (= xleft yleft)
           (= xright yright))

      [(#.Named xname xtype) (#.Named yname ytype)]
      (and (name/= xname yname)
           (= xtype ytype))

      (^template [<tag>]
        [(<tag> xL xR) (<tag> yL yR)]
        (and (= xL yL) (= xR yR)))
      ([#.Sum] [#.Product])
      
      (^or [(#.UnivQ xenv xbody) (#.UnivQ yenv ybody)]
           [(#.ExQ xenv xbody) (#.ExQ yenv ybody)])
      (and (n/= (list.size yenv) (list.size xenv))
           (= xbody ybody)
           (list/fold (.function (_ [x y] prev) (and prev (= x y)))
                      #1
                      (list.zip2 xenv yenv)))

      _
      #0
      )))

(do-template [<name> <tag>]
  [(def: #export (<name> type)
     (-> Type [Nat Type])
     (loop [num-args 0
            type type]
       (case type
         (<tag> env sub-type)
         (recur (inc num-args) sub-type)

         _
         [num-args type])))]

  [flatten-univ-q #.UnivQ]
  [flatten-ex-q   #.ExQ]
  )

(def: #export (flatten-function type)
  (-> Type [(List Type) Type])
  (case type
    (#.Function in out')
    (let [[ins out] (flatten-function out')]
      [(list& in ins) out])

    _
    [(list) type]))

(def: #export (flatten-application type)
  (-> Type [Type (List Type)])
  (case type
    (#.Apply arg func')
    (let [[func args] (flatten-application func')]
      [func (list/compose args (list arg))])

    _
    [type (list)]))

(do-template [<name> <tag>]
  [(def: #export (<name> type)
     (-> Type (List Type))
     (case type
       (<tag> left right)
       (list& left (<name> right))

       _
       (list type)))]

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

(def: #export (apply params func)
  (-> (List Type) Type (Maybe Type))
  (case params
    #.Nil
    (#.Some func)

    (#.Cons param params')
    (case func
      (^template [<tag>]
        (<tag> env body)
        (|> body
            (beta-reduce (list& func param env))
            (apply params')))
      ([#.UnivQ] [#.ExQ])

      (#.Apply A F)
      (apply (list& A params) F)

      (#.Named name unnamed)
      (apply params unnamed)
      
      _
      #.None)))

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

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

    (^template [<tag>]
      (<tag> left right)
      (` (<tag> (~ (to-code left))
                (~ (to-code right)))))
    ([#.Sum] [#.Product] [#.Function] [#.Apply])

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

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

(def: #export (to-text type)
  (-> Type Text)
  (case type
    (#.Primitive name params)
    (case params
      #.Nil
      ($_ text/compose "(primitive " name ")")

      _
      ($_ text/compose "(primitive " name " " (|> params (list/map to-text) list.reverse (list.interpose " ") (list/fold text/compose "")) ")"))

    (^template [<tag> <open> <close> <flatten>]
      (<tag> _)
      ($_ text/compose <open>
          (|> (<flatten> type)
              (list/map to-text)
              list.reverse
              (list.interpose " ")
              (list/fold text/compose ""))
          <close>))
    ([#.Sum  "(| " ")" flatten-variant]
     [#.Product "["   "]" flatten-tuple])

    (#.Function input output)
    (let [[ins out] (flatten-function type)]
      ($_ text/compose  "(-> "
          (|> ins
              (list/map to-text)
              list.reverse
              (list.interpose " ")
              (list/fold text/compose ""))
          " " (to-text out) ")"))

    (#.Parameter idx)
    (nat/encode idx)

    (#.Var id)
    ($_ text/compose "⌈v:" (nat/encode id) "⌋")

    (#.Ex id)
    ($_ text/compose "⟨e:" (nat/encode id) "⟩")

    (#.Apply param fun)
    (let [[type-func type-args] (flatten-application type)]
      ($_ text/compose  "(" (to-text type-func) " " (|> type-args (list/map to-text) list.reverse (list.interpose " ") (list/fold text/compose "")) ")"))

    (^template [<tag> <desc>]
      (<tag> env body)
      ($_ text/compose "(" <desc> " {" (|> env (list/map to-text) (text.join-with " ")) "} " (to-text body) ")"))
    ([#.UnivQ "All"]
     [#.ExQ "Ex"])

    (#.Named [module name] type)
    ($_ text/compose module "." name)
    ))

(def: #export (un-alias type)
  (-> Type Type)
  (case type
    (#.Named _ (#.Named name type'))
    (un-alias (#.Named name type'))

    _
    type))

(def: #export (un-name type)
  (-> Type Type)
  (case type
    (#.Named name type')
    (un-name type')

    _
    type))

(do-template [<name> <base> <ctor>]
  [(def: #export (<name> types)
     (-> (List Type) Type)
     (case types
       #.Nil
       <base>

       (#.Cons type #.Nil)
       type

       (#.Cons type types')
       (<ctor> type (<name> types'))))]

  [variant Nothing #.Sum]
  [tuple   Any     #.Product]
  )

(def: #export (function inputs output)
  (-> (List Type) Type Type)
  (case inputs
    #.Nil
    output

    (#.Cons input inputs')
    (#.Function input (function inputs' output))))

(def: #export (application params quant)
  (-> (List Type) Type Type)
  (case params
    #.Nil
    quant

    (#.Cons param params')
    (application params' (#.Apply param quant))))

(do-template [<name> <tag>]
  [(def: #export (<name> size body)
     (-> Nat Type Type)
     (case size
       0 body
       _  (|> body (<name> (dec size)) (<tag> (list)))))]

  [univ-q #.UnivQ]
  [ex-q   #.ExQ]
  )

(def: #export (quantified? type)
  (-> Type Bit)
  (case type
    (#.Named [module name] _type)
    (quantified? _type)

    (#.Apply A F)
    (maybe.default #0
                   (do maybe.monad
                     [applied (apply (list A) F)]
                     (wrap (quantified? applied))))
    
    (^or (#.UnivQ _) (#.ExQ _))
    #1

    _
    #0))

(def: #export (array depth elem-type)
  (-> Nat Type Type)
  (case depth
    0 elem-type
    _ (|> elem-type (array (dec depth)) (list) (#.Primitive array.type-name))))

(syntax: #export (:log! {input (p.or s.identifier
                                     s.any)})
  (case input
    (#.Left valueN)
    (do @
      [cursor macro.cursor
       valueT (macro.find-type valueN)
       #let [_ (log! ($_ text/compose
                         ":log!" " @ " (.cursor-description cursor) text.new-line
                         (name/encode valueN) " : " (..to-text valueT) text.new-line))]]
      (wrap (list (code.identifier valueN))))

    (#.Right valueC)
    (macro.with-gensyms [g!value]
      (wrap (list (` (.let [(~ g!value) (~ valueC)]
                       (..:log! (~ g!value)))))))))

(def: type-parameters
  (Syntax (List Text))
  (s.tuple (p.some s.local-identifier)))

(syntax: #export (:cast {type-vars type-parameters}
                        input
                        output
                        {value (p.maybe s.any)})
  (let [casterC (` (: (All [(~+ (list/map code.local-identifier type-vars))]
                        (-> (~ input) (~ output)))
                      (|>> :assume)))]
    (case value
      #.None
      (wrap (list casterC))
      
      (#.Some value)
      (wrap (list (` ((~ casterC) (~ value))))))))

(type: Typed
  {#type Code
   #expression Code})

(def: typed
  (Syntax Typed)
  (s.record (p.and s.any s.any)))

## TODO: Make sure the generated code always gets optimized away.
(syntax: #export (:share {type-vars type-parameters}
                         {exemplar typed}
                         {computation typed})
  (macro.with-gensyms [g!_]
    (let [shareC (` (: (All [(~+ (list/map code.local-identifier type-vars))]
                         (-> (~ (get@ #type exemplar))
                             (~ (get@ #type computation))))
                       (.function ((~ g!_) (~ g!_))
                         ## TODO: this should use : instead of :assume
                         (:assume (~ (get@ #expression computation))))))]
      (wrap (list (` ((~ shareC) (~ (get@ #expression exemplar)))))))))