(;module: {#;doc "Functions for extracting information from the state of the compiler."}
  lux
  (lux (macro [ast])
       (control functor
                applicative
                monad)
       (data (coll [list #* "List/" Monoid<List> Monad<List>])
             [number]
             [text "Text/" Monoid<Text> Eq<Text>]
             [product]
             [ident "Ident/" Codec<Text,Ident>]
             maybe
             ["E" error #- fail])))

## (type: (Lux a)
##   (-> Compiler (Error [Compiler a])))

(struct: #export _ (Functor Lux)
  (def: (map f fa)
    (function [state]
      (case (fa state)
        (#;Left msg)
        (#;Left msg)

        (#;Right [state' a])
        (#;Right [state' (f a)])))))

(struct: #export _ (Applicative Lux)
  (def: functor Functor<Lux>)

  (def: (wrap x)
    (function [state]
      (#;Right [state x])))

  (def: (apply ff fa)
    (function [state]
      (case (ff state)
        (#;Right [state' f])
        (case (fa state')
          (#;Right [state'' a])
          (#;Right [state'' (f a)])

          (#;Left msg)
          (#;Left msg))

        (#;Left msg)
        (#;Left msg)))))

(struct: #export _ (Monad Lux)
  (def: applicative Applicative<Lux>)
  
  (def: (join mma)
    (function [state]
      (case (mma state)
        (#;Left msg)
        (#;Left msg)

        (#;Right [state' ma])
        (ma state')))))

(def: (get k plist)
  (All [a]
    (-> Text (List [Text a]) (Maybe a)))
  (case plist
    #;Nil
    #;None

    (#;Cons [k' v] plist')
    (if (Text/= k k')
      (#;Some v)
      (get k plist'))))

(def: #export (run' compiler action)
  (All [a] (-> Compiler (Lux a) (Error [Compiler a])))
  (action compiler))

(def: #export (run compiler action)
  (All [a] (-> Compiler (Lux a) (Error a)))
  (case (action compiler)
    (#;Left error)
    (#;Left error)

    (#;Right [_ output])
    (#;Right output)))

(def: #export (either left right)
  {#;doc "Pick whichever computation succeeds."}
  (All [a] (-> (Lux a) (Lux a) (Lux a)))
  (function [compiler]
    (case (left compiler)
      (#;Left error)
      (right compiler)

      (#;Right [compiler' output])
      (#;Right [compiler' output]))))

(def: #export (assert message test)
  {#;doc "Fails with the given message if the test is false."}
  (-> Text Bool (Lux Unit))
  (function [compiler]
    (if test
      (#;Right [compiler []])
      (#;Left message))))

(def: #export (fail msg)
  {#;doc "Fails with the given message."}
  (All [a]
    (-> Text (Lux a)))
  (function [_]
    (#;Left msg)))

(def: #export (find-module name)
  (-> Text (Lux Module))
  (function [state]
    (case (get name (get@ #;modules state))
      (#;Some module)
      (#;Right [state module])

      _
      (#;Left ($_ Text/append "Unknown module: " name)))))

(def: #export current-module-name
  (Lux Text)
  (function [state]
    (case (list;last (get@ #;scopes state))
      (#;Some scope)
      (case (get@ #;name scope)
        (#;Cons m-name #;Nil)
        (#;Right [state m-name])

        _
        (#;Left "Improper name for scope."))

      _
      (#;Left "Empty environment!")
      )))

(def: #export current-module
  (Lux Module)
  (do Monad<Lux>
    [this-module-name current-module-name]
    (find-module this-module-name)))

(def: #export (get-ann tag anns)
  {#;doc "Looks-up a particular annotation's value within the set of annotations."}
  (-> Ident Anns (Maybe Ann-Value))
  (let [[p n] tag]
    (case anns
      (#;Cons [[p' n'] dmv] anns')
      (if (and (Text/= p p')
               (Text/= n n'))
        (#;Some dmv)
        (get-ann tag anns'))
      
      #;Nil
      #;None)))

(do-template [<name> <tag> <type>]
  [(def: #export (<name> tag anns)
     (-> Ident Anns (Maybe <type>))
     (case (get-ann tag anns)
       (#;Some (<tag> value))
       (#;Some value)

       _
       #;None))]

  [get-bool-ann  #;BoolA  Bool]
  [get-int-ann   #;IntA   Int]
  [get-real-ann  #;RealA  Real]
  [get-char-ann  #;CharA  Char]
  [get-text-ann  #;TextA  Text]
  [get-ident-ann #;IdentA Ident]
  [get-list-ann  #;ListA  (List Ann-Value)]
  [get-dict-ann  #;DictA  (List [Text Ann-Value])]
  )

(def: #export (get-doc anns)
  {#;doc "Looks-up a definition's documentation."}
  (-> Anns (Maybe Text))
  (get-text-ann ["lux" "doc"] anns))

(def: #export (flag-set? flag-name anns)
  {#;doc "Finds out whether an annotation-as-a-flag is set (has value 'true')."}
  (-> Ident Anns Bool)
  (case (get-ann flag-name anns)
    (#;Some (#;BoolA true))
    true

    _
    false))

(do-template [<name> <tag> <desc>]
  [(def: #export <name>
     {#;doc (#;TextA ($_ Text/append "Checks whether a definition is " <desc> "."))}
     (-> Anns Bool)
     (flag-set? (ident-for <tag>)))]

  [export?   #;export?   "exported"]
  [hidden?   #;hidden?   "hidden"]
  [macro?    #;macro?    "a macro"]
  [type?     #;type?     "a type"]
  [struct?   #;struct?   "a structure"]
  [type-rec? #;type-rec? "a recursive type"]
  [sig?      #;sig?      "a signature"]
  )

(do-template [<name> <tag> <type>]
  [(def: (<name> dmv)
     (-> Ann-Value (Maybe <type>))
     (case dmv
       (<tag> actual-value)
       (#;Some actual-value)

       _
       #;None))]

  [try-mlist #;ListA (List Ann-Value)]
  [try-mtext #;TextA Text]
  )

(do-template [<name> <tag> <desc>]
  [(def: #export (<name> anns)
     {#;doc (#;TextA ($_ Text/append "Looks up the arguments of a " <desc> "."))}
     (-> Anns (List Text))
     (default (list)
       (do Monad<Maybe>
         [_args (get-ann (ident-for <tag>) anns)
          args (try-mlist _args)]
         (mapM @ try-mtext args))))]

  [func-args #;func-args "function"]
  [type-args #;type-args "parameterized type"]
  )

(def: (find-macro' modules this-module module name)
  (-> (List [Text Module]) Text Text Text
      (Maybe Macro))
  (do Monad<Maybe>
    [$module (get module modules)
     [def-type def-anns def-value] (: (Maybe Def) (|> (: Module $module) (get@ #;defs) (get name)))]
    (if (and (macro? def-anns)
             (or (export? def-anns) (Text/= module this-module)))
      (#;Some (:! Macro def-value))
      (case (get-ann ["lux" "alias"] def-anns)
        (#;Some (#;IdentA [r-module r-name]))
        (find-macro' modules this-module r-module r-name)

        _
        #;None))))

(def: #export (find-macro ident)
  (-> Ident (Lux (Maybe Macro)))
  (do Monad<Lux>
    [this-module current-module-name]
    (let [[module name] ident]
      (: (Lux (Maybe Macro))
         (function [state]
           (#;Right [state (find-macro' (get@ #;modules state) this-module module name)]))))))

(def: #export (normalize ident)
  {#;doc "If given an identifier without a module prefix, gives it the current module's name as prefix.

          Otherwise, returns the identifier as-is."}
  (-> Ident (Lux Ident))
  (case ident
    ["" name]
    (do Monad<Lux>
      [module-name current-module-name]
      (wrap [module-name name]))

    _
    (:: Monad<Lux> wrap ident)))

(def: #export (macro-expand-once syntax)
  {#;doc "Given code that requires applying a macro, does it once and returns the result.

          Otherwise, returns the code as-is."}
  (-> AST (Lux (List AST)))
  (case syntax
    [_ (#;FormS (#;Cons [[_ (#;SymbolS macro-name)] args]))]
    (do Monad<Lux>
      [macro-name' (normalize macro-name)
       ?macro (find-macro macro-name')]
      (case ?macro
        (#;Some macro)
        (macro args)
        
        #;None
        (:: Monad<Lux> wrap (list syntax))))

    _
    (:: Monad<Lux> wrap (list syntax))))

(def: #export (macro-expand syntax)
  {#;doc "Given code that requires applying a macro, expands repeatedly until no more direct macro-calls are left.

          Otherwise, returns the code as-is."}
  (-> AST (Lux (List AST)))
  (case syntax
    [_ (#;FormS (#;Cons [[_ (#;SymbolS macro-name)] args]))]
    (do Monad<Lux>
      [macro-name' (normalize macro-name)
       ?macro (find-macro macro-name')]
      (case ?macro
        (#;Some macro)
        (do Monad<Lux>
          [expansion (macro args)
           expansion' (mapM Monad<Lux> macro-expand expansion)]
          (wrap (:: Monad<List> join expansion')))
        
        #;None
        (:: Monad<Lux> wrap (list syntax))))

    _
    (:: Monad<Lux> wrap (list syntax))))

(def: #export (macro-expand-all syntax)
  {#;doc "Expands all macro-calls everywhere recursively, until only primitive/base code remains."}
  (-> AST (Lux (List AST)))
  (case syntax
    [_ (#;FormS (#;Cons [[_ (#;SymbolS macro-name)] args]))]
    (do Monad<Lux>
      [macro-name' (normalize macro-name)
       ?macro (find-macro macro-name')]
      (case ?macro
        (#;Some macro)
        (do Monad<Lux>
          [expansion (macro args)
           expansion' (mapM Monad<Lux> macro-expand-all expansion)]
          (wrap (:: Monad<List> join expansion')))
        
        #;None
        (do Monad<Lux>
          [parts' (mapM Monad<Lux> macro-expand-all (list& (ast;symbol macro-name) args))]
          (wrap (list (ast;form (:: Monad<List> join parts')))))))

    [_ (#;FormS (#;Cons [harg targs]))]
    (do Monad<Lux>
      [harg+ (macro-expand-all harg)
       targs+ (mapM Monad<Lux> macro-expand-all targs)]
      (wrap (list (ast;form (List/append harg+ (:: Monad<List> join (: (List (List AST)) targs+)))))))

    [_ (#;TupleS members)]
    (do Monad<Lux>
      [members' (mapM Monad<Lux> macro-expand-all members)]
      (wrap (list (ast;tuple (:: Monad<List> join members')))))

    _
    (:: Monad<Lux> wrap (list syntax))))

(def: #export (gensym prefix)
  {#;doc "Generates a unique identifier as an AST node (ready to be used in code templates).

          A prefix can be given (or just be empty text \"\") to better identify the code for debugging purposes."}
  (-> Text (Lux AST))
  (function [state]
    (#;Right [(update@ #;seed n.inc state)
              (ast;symbol ["" ($_ Text/append "__gensym__" prefix (:: number;Codec<Text,Nat> encode (get@ #;seed state)))])])))

(def: (get-local-symbol ast)
  (-> AST (Lux Text))
  (case ast
    [_ (#;SymbolS [_ name])]
    (:: Monad<Lux> wrap name)

    _
    (fail (Text/append "AST is not a local symbol: " (ast;to-text ast)))))

(macro: #export (with-gensyms tokens)
  {#;doc (doc "Creates new symbols and offers them to the body expression."
              (syntax: #export (synchronized lock body)
                (with-gensyms [g!lock g!body g!_]
                  (wrap (list (` (let [(~ g!lock) (~ lock)
                                       (~ g!_) (;_jvm_monitorenter (~ g!lock))
                                       (~ g!body) (~ body)
                                       (~ g!_) (;_jvm_monitorexit (~ g!lock))]
                                   (~ g!body)))))
                  )))}
  (case tokens
    (^ (list [_ (#;TupleS symbols)] body))
    (do Monad<Lux>
      [symbol-names (mapM @ get-local-symbol symbols)
       #let [symbol-defs (List/join (List/map (: (-> Text (List AST))
                                                 (function [name] (list (ast;symbol ["" name]) (` (gensym (~ (ast;text name)))))))
                                              symbol-names))]]
      (wrap (list (` (do Monad<Lux>
                       [(~@ symbol-defs)]
                       (~ body))))))

    _
    (fail "Wrong syntax for with-gensyms")))

(def: #export (macro-expand-1 token)
  {#;doc "Works just like macro-expand, except that it ensures that the output is a single AST token."}
  (-> AST (Lux AST))
  (do Monad<Lux>
    [token+ (macro-expand token)]
    (case token+
      (^ (list token'))
      (wrap token')

      _
      (fail "Macro expanded to more than 1 element."))))

(def: #export (module-exists? module)
  (-> Text (Lux Bool))
  (function [state]
    (#;Right [state (case (get module (get@ #;modules state))
                      (#;Some _)
                      true
                      
                      #;None
                      false)])))

(def: (try-both f x1 x2)
  (All [a b]
    (-> (-> a (Maybe b)) a a (Maybe b)))
  (case (f x1)
    #;None     (f x2)
    (#;Some y) (#;Some y)))

(def: #export (find-var-type name)
  {#;doc "Looks-up the type of a local variable somewhere in the environment."}
  (-> Text (Lux Type))
  (function [state]
    (let [test (: (-> [Text (Meta [Type Cursor] Void)] Bool)
                  (|>. product;left (Text/= name)))]
      (case (do Monad<Maybe>
              [scope (find (function [env]
                             (or (any? test (:! (List [Text (Meta [Type Cursor] Void)])
                                                (get@ [#;locals #;mappings] env)))
                                 (any? test (:! (List [Text (Meta [Type Cursor] Void)])
                                                (get@ [#;closure #;mappings] env)))))
                           (get@ #;scopes state))
               [_ [[type _] _]] (try-both (find test)
                                          (:! (List [Text (Meta [Type Cursor] Void)])
                                              (get@ [#;locals #;mappings] scope))
                                          (:! (List [Text (Meta [Type Cursor] Void)])
                                              (get@ [#;closure #;mappings] scope)))]
              (wrap type))
        (#;Some var-type)
        (#;Right [state var-type])

        #;None
        (#;Left ($_ Text/append "Unknown variable: " name))))))

(def: #export (find-def name)
  {#;doc "Looks-up a definition's whole data in the available modules (including the current one)."}
  (-> Ident (Lux Def))
  (function [state]
    (case (: (Maybe Def)
             (do Monad<Maybe>
               [#let [[v-prefix v-name] name]
                (^slots [#;defs]) (get v-prefix (get@ #;modules state))]
               (get v-name defs)))
      (#;Some _anns)
      (#;Right [state _anns])

      _
      (#;Left ($_ Text/append "Unknown definition: " (Ident/encode name))))))

(def: #export (find-def-type name)
  {#;doc "Looks-up a definition's type in the available modules (including the current one)."}
  (-> Ident (Lux Type))
  (do Monad<Lux>
    [[def-type def-data def-value] (find-def name)]
    (wrap def-type)))

(def: #export (find-type name)
  {#;doc "Looks-up the type of either a local variable or a definition."}
  (-> Ident (Lux Type))
  (do Monad<Lux>
    [#let [[_ _name] name]]
    (either (find-var-type _name)
            (do @
              [name (normalize name)]
              (find-def-type name)))))

(def: #export (find-type-def name)
  {#;doc "Finds the value of a type definition (such as Int, Top or Compiler)."}
  (-> Ident (Lux Type))
  (do Monad<Lux>
    [[def-type def-data def-value] (find-def name)]
    (wrap (:! Type def-value))))

(def: #export (defs module-name)
  {#;doc "The entire list of definitions in a module (including the unexported/private ones)."}
  (-> Text (Lux (List [Text Def])))
  (function [state]
    (case (get module-name (get@ #;modules state))
      #;None          (#;Left ($_ Text/append "Unknown module: " module-name))
      (#;Some module) (#;Right [state (get@ #;defs module)])
      )))

(def: #export (exports module-name)
  {#;doc "All the exported definitions in a module."}
  (-> Text (Lux (List [Text Def])))
  (do Monad<Lux>
    [defs (defs module-name)]
    (wrap (filter (function [[name [def-type def-anns def-value]]]
                    (and (export? def-anns)
                         (not (hidden? def-anns))))
                  defs))))

(def: #export modules
  {#;doc "All the available modules (including the current one)."}
  (Lux (List [Text Module]))
  (function [state]
    (|> state
        (get@ #;modules)
        [state]
        #;Right)))

(def: #export (tags-of type-name)
  {#;doc "All the tags associated with a type definition."}
  (-> Ident (Lux (List Ident)))
  (do Monad<Lux>
    [#let [[module name] type-name]
     module (find-module module)]
    (case (get name (get@ #;types module))
      (#;Some [tags _])
      (wrap tags)

      _
      (wrap (list)))))

(def: #export cursor
  {#;doc "The cursor of the current expression being analyzed."}
  (Lux Cursor)
  (function [state]
    (#;Right [state (get@ #;cursor state)])))

(def: #export expected-type
  {#;doc "The expected type of the current expression being analyzed."}
  (Lux Type)
  (function [state]
    (case (get@ #;expected state)
      (#;Some type)
      (#;Right [state type])

      #;None
      (#;Left "Not expecting any type."))))

(def: #export (imported-modules module-name)
  {#;doc "All the modules imported by a specified module."}
  (-> Text (Lux (List Text)))
  (do Monad<Lux>
    [(^slots [#;imports]) (find-module module-name)]
    (wrap imports)))

(def: #export (resolve-tag tag)
  {#;doc "Given a tag, finds out what is its index, its related tag-list and it's associated type."}
  (-> Ident (Lux [Nat (List Ident) Type]))
  (do Monad<Lux>
    [#let [[module name] tag]
     =module (find-module module)
     this-module-name current-module-name]
    (case (get name (get@ #;tags =module))
      (#;Some [idx tag-list exported? type])
      (if (or exported?
              (Text/= this-module-name module))
        (wrap [idx tag-list type])
        (fail ($_ Text/append "Cannot access tag: " (Ident/encode tag) " from module " this-module-name)))

      _
      (fail ($_ Text/append "Unknown tag: " (Ident/encode tag))))))

(def: #export (tag-lists module)
  {#;doc "All the tag-lists defined in a module, with their associated types."}
  (-> Text (Lux (List [(List Ident) Type])))
  (do Monad<Lux>
    [=module (find-module module)
     this-module-name current-module-name]
    (wrap (|> (get@ #;types =module)
              (list;filter (function [[type-name [tag-list exported? type]]]
                             (or exported?
                                 (Text/= this-module-name module))))
              (List/map (function [[type-name [tag-list exported? type]]]
                          [tag-list type]))))))

(def: #export locals
  {#;doc "All the local variables currently in scope, separated in different scopes."}
  (Lux (List (List [Text Type])))
  (function [state]
    (case (list;inits (get@ #;scopes state))
      #;None
      (#;Left "No local environment")

      (#;Some scopes)
      (#;Right [state
                (List/map (|>. (get@ [#;locals #;mappings])
                               (:! (List [Text (Meta [Type Cursor] Void)]))
                               (List/map (function [[name [[type cursor] analysis]]]
                                           [name type])))
                          scopes)]))))

(def: #export (un-alias def-name)
  {#;doc "Given an aliased definition's name, returns the original definition being referenced."}
  (-> Ident (Lux Ident))
  (do Monad<Lux>
    [def-name (normalize def-name)
     [_ def-anns _] (find-def def-name)]
    (case (get-ann (ident-for #;alias) def-anns)
      (#;Some (#;IdentA real-def-name))
      (wrap real-def-name)

      _
      (wrap def-name))))

(def: #export get-compiler
  {#;doc "Obtains the current state of the compiler."}
  (Lux Compiler)
  (function [compiler]
    (#;Right [compiler compiler])))

(do-template [<macro> <func> <desc>]
  [(macro: #export (<macro> tokens)
     {#;doc (doc "Performs a macro-expansion and logs the resulting ASTs."
                 "You can either use the resulting ASTs, or omit them."
                 "By omitting them, this macro produces nothing (just like the lux;comment macro)."
                 (<macro> #omit
                          (def: (foo bar baz)
                            (-> Int Int Int)
                            (i.+ bar baz))))}
     (case tokens
       (^ (list [_ (#;TagS ["" "omit"])]
                token))
       (do Monad<Lux>
         [output (<func> token)
          #let [_ (List/map (. log! ast;to-text)
                            output)]]
         (wrap (list)))

       (^ (list token))
       (do Monad<Lux>
         [output (<func> token)
          #let [_ (List/map (. log! ast;to-text)
                            output)]]
         (wrap output))

       _
       (fail ($_ Text/append "Wrong syntax for " <desc> "."))))]

  [log-expand      macro-expand      "log-macro-expand"]
  [log-expand-all  macro-expand-all  "log-macro-expand-all"]
  [log-expand-once macro-expand-once "log-macro-expand-once"]
  )