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

## (type: (Meta a)
##   (-> Compiler (E;Error [Compiler a])))

(struct: #export _ (F;Functor Meta)
  (def: (map f fa)
    (function [state]
      (case (fa state)
        (#E;Error msg)
        (#E;Error msg)

        (#E;Success [state' a])
        (#E;Success [state' (f a)])))))

(struct: #export _ (A;Applicative Meta)
  (def: functor Functor<Meta>)

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

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

          (#E;Error msg)
          (#E;Error msg))

        (#E;Error msg)
        (#E;Error msg)))))

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

        (#E;Success [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 (Meta a) (E;Error [Compiler a])))
  (action compiler))

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

    (#E;Success [_ output])
    (#E;Success output)))

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

      (#E;Success [compiler' output])
      (#E;Success [compiler' output]))))

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

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

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

      _
      (#E;Error ($_ text/compose "Unknown module: " name)))))

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

        _
        (#E;Error "Improper name for scope."))

      _
      (#E;Error "Empty environment!")
      )))

(def: #export current-module
  (Meta Module)
  (do Monad<Meta>
    [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 Code (Maybe Code))
  (case anns
    [_ (#;Record anns)]
    (loop [anns anns]
      (case anns
        (#;Cons [key value] anns')
        (case key
          [_ (#;Tag tag')]
          (if (ident/= tag tag')
            (#;Some value)
            (recur anns'))

          _
          (recur anns'))
        
        #;Nil
        #;None))

    _
    #;None))

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

       _
       #;None))]

  [get-bool-ann   #;Bool   Bool]
  [get-int-ann    #;Int    Int]
  [get-frac-ann   #;Frac   Frac]
  [get-text-ann   #;Text   Text]
  [get-symbol-ann #;Symbol Ident]
  [get-tag-ann    #;Tag    Ident]
  [get-form-ann   #;Form   (List Code)]
  [get-tuple-ann  #;Tuple  (List Code)]
  [get-record-ann #;Record (List [Code Code])]
  )

(def: #export (get-doc anns)
  {#;doc "Looks-up a definition's documentation."}
  (-> Code (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 Code Bool)
  (maybe;default false (get-bool-ann flag-name anns)))

(do-template [<name> <tag> <desc>]
  [(def: #export <name>
     {#;doc (code;text ($_ text/compose "Checks whether a definition is " <desc> "."))}
     (-> Code 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> input)
     (-> Code (Maybe <type>))
     (case input
       [_ (<tag> actual-value)]
       (#;Some actual-value)

       _
       #;None))]

  [parse-tuple #;Tuple (List Code)]
  [parse-text  #;Text  Text]
  )

(do-template [<name> <tag> <desc>]
  [(def: #export (<name> anns)
     {#;doc <desc>}
     (-> Code (List Text))
     (maybe;default (list)
                    (do maybe;Monad<Maybe>
                      [_args (get-ann (ident-for <tag>) anns)
                       args (parse-tuple _args)]
                      (M;map @ parse-text args))))]

  [func-args     #;func-args "Looks up the arguments of a function."]
  [type-args     #;type-args "Looks up the arguments of a parameterized type."]
  [declared-tags #;tags      "Looks up the tags of a tagged (variant or record) type."]
  )

(def: (find-macro' modules this-module module name)
  (-> (List [Text Module]) Text Text Text
      (Maybe Macro))
  (do maybe;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-symbol-ann ["lux" "alias"] def-anns)
        (#;Some [r-module r-name])
        (find-macro' modules this-module r-module r-name)

        _
        #;None))))

(def: #export (find-macro ident)
  (-> Ident (Meta (Maybe Macro)))
  (do Monad<Meta>
    [this-module current-module-name]
    (let [[module name] ident]
      (: (Meta (Maybe Macro))
         (function [state]
           (#E;Success [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 (Meta Ident))
  (case ident
    ["" name]
    (do Monad<Meta>
      [module-name current-module-name]
      (wrap [module-name name]))

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

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

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

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

(def: #export (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."}
  (-> Code (Meta (List Code)))
  (case syntax
    [_ (#;Form (#;Cons [[_ (#;Symbol name)] args]))]
    (do Monad<Meta>
      [name' (normalize name)
       ?macro (find-macro name')]
      (case ?macro
        (#;Some macro)
        (do Monad<Meta>
          [expansion (macro args)
           expansion' (M;map Monad<Meta> expand expansion)]
          (wrap (list/join expansion')))
        
        #;None
        (:: Monad<Meta> wrap (list syntax))))

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

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

    [_ (#;Form (#;Cons [harg targs]))]
    (do Monad<Meta>
      [harg+ (expand-all harg)
       targs+ (M;map Monad<Meta> expand-all targs)]
      (wrap (list (code;form (list/compose harg+ (list/join (: (List (List Code)) targs+)))))))

    [_ (#;Tuple members)]
    (do Monad<Meta>
      [members' (M;map Monad<Meta> expand-all members)]
      (wrap (list (code;tuple (list/join members')))))

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

(def: #export (gensym prefix)
  {#;doc "Generates a unique identifier as an Code 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 (Meta Code))
  (function [state]
    (#E;Success [(update@ #;seed n.inc state)
                 (code;symbol ["" ($_ text/compose "__gensym__" prefix (:: number;Codec<Text,Nat> encode (get@ #;seed state)))])])))

(def: (get-local-symbol ast)
  (-> Code (Meta Text))
  (case ast
    [_ (#;Symbol [_ name])]
    (:: Monad<Meta> wrap name)

    _
    (fail (text/compose "Code is not a local symbol: " (code;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 [_ (#;Tuple symbols)] body))
    (do Monad<Meta>
      [symbol-names (M;map @ get-local-symbol symbols)
       #let [symbol-defs (list/join (list/map (: (-> Text (List Code))
                                                 (function [name] (list (code;symbol ["" name]) (` (gensym (~ (code;text name)))))))
                                              symbol-names))]]
      (wrap (list (` (do Monad<Meta>
                       [(~@ symbol-defs)]
                       (~ body))))))

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

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

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

(def: #export (module-exists? module)
  (-> Text (Meta Bool))
  (function [state]
    (#E;Success [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 (Meta Type))
  (function [state]
    (let [test (: (-> [Text [Type Top]] Bool)
                  (|>. product;left (text/= name)))]
      (case (do maybe;Monad<Maybe>
              [scope (list;find (function [env]
                                  (or (list;any? test (: (List [Text [Type Top]])
                                                         (get@ [#;locals #;mappings] env)))
                                      (list;any? test (: (List [Text [Type Top]])
                                                         (get@ [#;captured #;mappings] env)))))
                                (get@ #;scopes state))
               [_ [type _]] (try-both (list;find test)
                                      (: (List [Text [Type Top]])
                                         (get@ [#;locals #;mappings] scope))
                                      (: (List [Text [Type Top]])
                                         (get@ [#;captured #;mappings] scope)))]
              (wrap type))
        (#;Some var-type)
        (#E;Success [state var-type])

        #;None
        (#E;Error ($_ text/compose "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 (Meta Def))
  (function [state]
    (case (: (Maybe Def)
             (do maybe;Monad<Maybe>
               [#let [[v-prefix v-name] name]
                (^slots [#;defs]) (get v-prefix (get@ #;modules state))]
               (get v-name defs)))
      (#;Some _anns)
      (#E;Success [state _anns])

      _
      (#E;Error ($_ text/compose "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 (Meta Type))
  (do Monad<Meta>
    [[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 (Meta Type))
  (do Monad<Meta>
    [#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 (Meta Type))
  (do Monad<Meta>
    [[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 (Meta (List [Text Def])))
  (function [state]
    (case (get module-name (get@ #;modules state))
      #;None          (#E;Error ($_ text/compose "Unknown module: " module-name))
      (#;Some module) (#E;Success [state (get@ #;defs module)])
      )))

(def: #export (exports module-name)
  {#;doc "All the exported definitions in a module."}
  (-> Text (Meta (List [Text Def])))
  (do Monad<Meta>
    [defs (defs module-name)]
    (wrap (list;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)."}
  (Meta (List [Text Module]))
  (function [state]
    (|> state
        (get@ #;modules)
        [state]
        #E;Success)))

(def: #export (tags-of type-name)
  {#;doc "All the tags associated with a type definition."}
  (-> Ident (Meta (List Ident)))
  (do Monad<Meta>
    [#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."}
  (Meta Cursor)
  (function [state]
    (#E;Success [state (get@ #;cursor state)])))

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

      #;None
      (#E;Error "Not expecting any type."))))

(def: #export (imported-modules module-name)
  {#;doc "All the modules imported by a specified module."}
  (-> Text (Meta (List Text)))
  (do Monad<Meta>
    [(^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 (Meta [Nat (List Ident) Type]))
  (do Monad<Meta>
    [#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/compose "Cannot access tag: " (ident/encode tag) " from module " this-module-name)))

      _
      (fail ($_ text/compose "Unknown tag: " (ident/encode tag))))))

(def: #export (tag-lists module)
  {#;doc "All the tag-lists defined in a module, with their associated types."}
  (-> Text (Meta (List [(List Ident) Type])))
  (do Monad<Meta>
    [=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."}
  (Meta (List (List [Text Type])))
  (function [state]
    (case (list;inits (get@ #;scopes state))
      #;None
      (#E;Error "No local environment")

      (#;Some scopes)
      (#E;Success [state
                   (list/map (|>. (get@ [#;locals #;mappings])
                                  (list/map (function [[name [type _]]]
                                              [name type])))
                             scopes)]))))

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

      _
      (wrap def-name))))

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

(def: #export type-context
  (Meta Type-Context)
  (function [compiler]
    (#E;Success [compiler (get@ #;type-context compiler)])))

(do-template [<macro> <func> <desc>]
  [(macro: #export (<macro> tokens)
     {#;doc (doc "Performs a macro-expansion and logs the resulting code."
                 "You can either use the resulting code, 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 [_ (#;Tag ["" "omit"])]
                token))
       (do Monad<Meta>
         [output (<func> token)
          #let [_ (list/map (. log! code;to-text)
                            output)]]
         (wrap (list)))

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

       _
       (fail ($_ text/compose "Wrong syntax for " <desc> "."))))]

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