##  Copyright (c) Eduardo Julian. All rights reserved.
##  This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
##  If a copy of the MPL was not distributed with this file,
##  You can obtain one at http://mozilla.org/MPL/2.0/.

(;module:
  lux
  (lux (control eq
                monad)
       (data [text "Text/" Monoid<Text> Eq<Text>]
             [ident "Ident/" Eq<Ident>]
             [number "Nat/" Codec<Text,Nat>]
             maybe
             (struct [list #+ "List/" Monad<List> Monoid<List> Fold<List>]))
       (macro [ast])
       ))

## [Utils]
(def: (beta-reduce env type)
  (-> (List Type) Type Type)
  (case type
    (#;HostT name params)
    (#;HostT name (List/map (beta-reduce env) params))
    
    (^template [<tag>]
      (<tag> left right)
      (<tag> (beta-reduce env left) (beta-reduce env right)))
    ([#;SumT] [#;ProdT])
    
    (^template [<tag>]
      (<tag> left right)
      (<tag> (beta-reduce env left) (beta-reduce env right)))
    ([#;LambdaT]
     [#;AppT])

    (^template [<tag>]
      (<tag> old-env def)
      (case old-env
        #;Nil
        (<tag> env def)

        _
        type))
    ([#;UnivQ]
     [#;ExQ])
    
    (#;BoundT idx)
    (default type (list;at idx env))
    
    (#;NamedT name type)
    (beta-reduce env type)

    _
    type
    ))

## [Structures]
(struct: #export _ (Eq Type)
  (def: (= x y)
    (case [x y]
      [(#;HostT xname xparams) (#;HostT yname yparams)]
      (and (Text/= xname yname)
           (n.= (list;size yparams) (list;size xparams))
           (List/fold (lambda [[x y] prev] (and prev (= x y)))
                      true
                      (list;zip2 xparams yparams)))

      (^template [<tag>]
        [<tag> <tag>]
        true)
      ([#;VoidT] [#;UnitT])
      
      (^template [<tag>]
        [(<tag> xid) (<tag> yid)]
        (n.= yid xid))
      ([#;VarT] [#;ExT] [#;BoundT])

      (^or [(#;LambdaT xleft xright) (#;LambdaT yleft yright)]
           [(#;AppT xleft xright) (#;AppT yleft yright)])
      (and (= xleft yleft)
           (= xright yright))

      [(#;NamedT xname xtype) (#;NamedT yname ytype)]
      (and (Ident/= xname yname)
           (= xtype ytype))

      (^template [<tag>]
        [(<tag> xL xR) (<tag> yL yR)]
        (and (= xL yL) (= xR yR)))
      ([#;SumT] [#;ProdT])
      
      (^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 (lambda [[x y] prev] (and prev (= x y)))
                      true
                      (list;zip2 xenv yenv)))

      _
      false
      )))

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

         _
         [num-args type])))]

  [flatten-univq #;UnivQ]
  [flatten-exq   #;ExQ]
  )

(def: #export (flatten-function type)
  (-> Type [(List Type) Type])
  (case type
    (#;LambdaT 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
    (#;AppT left' right)
    (let [[left rights] (flatten-application left')]
      [left (List/append rights (list right))])

    _
    [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 #;SumT]
  [flatten-tuple   #;ProdT]
  )

(def: #export (apply-type type-fun param)
  (-> Type Type (Maybe Type))
  (case type-fun
    (^template [<tag>]
      (<tag> env body)
      (#;Some (beta-reduce (list& type-fun param env) body)))
    ([#;UnivQ] [#;ExQ])

    (#;AppT F A)
    (do Monad<Maybe>
      [type-fn* (apply-type F A)]
      (apply-type type-fn* param))

    (#;NamedT name type)
    (apply-type type param)
    
    _
    #;None))

(def: #export (to-ast type)
  (-> Type AST)
  (case type
    (#;HostT name params)
    (` (#;HostT (~ (ast;text name))
                (list (~@ (List/map to-ast params)))))

    (^template [<tag>]
      <tag>
      (` <tag>))
    ([#;VoidT] [#;UnitT])

    (^template [<tag>]
      (<tag> idx)
      (` (<tag> (~ (ast;nat idx)))))
    ([#;VarT] [#;ExT] [#;BoundT])

    (^template [<tag>]
      (<tag> left right)
      (` (<tag> (~ (to-ast left))
                (~ (to-ast right)))))
    ([#;LambdaT] [#;AppT])

    (^template [<tag> <macro> <flattener>]
      (<tag> left right)
      (` (<macro> (~@ (List/map to-ast (<flattener> type))))))
    ([#;SumT  | flatten-variant]
     [#;ProdT & flatten-tuple])

    (#;NamedT name sub-type)
    (ast;symbol name)

    (^template [<tag>]
      (<tag> env body)
      (` (<tag> (list (~@ (List/map to-ast env)))
                (~ (to-ast body)))))
    ([#;UnivQ] [#;ExQ])
    ))

(def: #export (to-text type)
  (-> Type Text)
  (case type
    (#;HostT name params)
    (case params
      #;Nil
      ($_ Text/append "(host " name ")")

      _
      ($_ Text/append "(host " name " " (|> params (List/map to-text) list;reverse (list;interpose " ") (List/fold Text/append "")) ")"))

    #;VoidT
    "Void"
    
    #;UnitT
    "Unit"

    (^template [<tag> <open> <close> <flatten>]
      (<tag> _)
      ($_ Text/append <open>
          (|> (<flatten> type)
              (List/map to-text)
              list;reverse
              (list;interpose " ")
              (List/fold Text/append ""))
          <close>))
    ([#;SumT  "(| " ")" flatten-variant]
     [#;ProdT "["   "]" flatten-tuple])

    (#;LambdaT input output)
    (let [[ins out] (flatten-function type)]
      ($_ Text/append  "(-> "
          (|> ins
              (List/map to-text)
              list;reverse
              (list;interpose " ")
              (List/fold Text/append ""))
          " " (to-text out) ")"))

    (#;BoundT idx)
    (Nat/encode idx)

    (#;VarT id)
    ($_ Text/append "⌈v:" (Nat/encode id) "⌋")

    (#;ExT id)
    ($_ Text/append "⟨e:" (Nat/encode id) "⟩")

    (#;AppT fun param)
    (let [[type-fun type-args] (flatten-application type)]
      ($_ Text/append  "(" (to-text type-fun) " " (|> type-args (List/map to-text) list;reverse (list;interpose " ") (List/fold Text/append "")) ")"))

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

    (#;NamedT [module name] type)
    ($_ Text/append module ";" name)
    ))

(def: #export (un-alias type)
  (-> Type Type)
  (case type
    (#;NamedT _ (#;NamedT ident type'))
    (un-alias (#;NamedT ident type'))

    _
    type))

(def: #export (un-name type)
  (-> Type Type)
  (case type
    (#;NamedT ident 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 Void #;SumT]
  [tuple   Unit #;ProdT]
  )

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

    (#;Cons input inputs')
    (#;LambdaT input (function inputs' output))))

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

    (#;Cons param params')
    (application (#;AppT quant param) params')))

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

  [univq #;UnivQ]
  [exq   #;ExQ]
  )