(.module: {#.doc "Functions for extracting information from the state of the compiler."}
  [lux #*
   [control
    [functor (#+ Functor)]
    [apply (#+ Apply)]
    ["." monad (#+ do Monad)]]
   [data
    ["." product]
    [name ("name/." Codec<Text,Name> Equivalence<Name>)]
    ["." maybe]
    ["e" error]
    ["." number ("nat/." Codec<Text,Nat>)]
    ["." text ("text/." Monoid<Text> Equivalence<Text>)]
    [collection
     ["." list ("list/." Monoid<List> Monad<List>)]]]]
  [/
   ["." code]])

## (type: (Meta a)
##   (-> Lux (e.Error [Lux a])))

(structure: #export _ (Functor Meta)
  (def: (map f fa)
    (function (_ compiler)
      (case (fa compiler)
        (#e.Error msg)
        (#e.Error msg)

        (#e.Success [compiler' a])
        (#e.Success [compiler' (f a)])))))

(structure: #export _ (Apply Meta)
  (def: functor Functor<Meta>)

  (def: (apply ff fa)
    (function (_ compiler)
      (case (ff compiler)
        (#e.Success [compiler' f])
        (case (fa compiler')
          (#e.Success [compiler'' a])
          (#e.Success [compiler'' (f a)])

          (#e.Error msg)
          (#e.Error msg))

        (#e.Error msg)
        (#e.Error msg)))))

(structure: #export _ (Monad Meta)
  (def: functor Functor<Meta>)

  (def: (wrap x)
    (function (_ compiler)
      (#e.Success [compiler x])))
  
  (def: (join mma)
    (function (_ compiler)
      (case (mma compiler)
        (#e.Error msg)
        (#e.Error msg)

        (#e.Success [compiler' ma])
        (ma compiler')))))

(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] (-> Lux (Meta a) (e.Error [Lux a])))
  (action compiler))

(def: #export (run compiler action)
  (All [a] (-> Lux (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 #0."}
  (-> Text Bit (Meta Any))
  (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 (_ compiler)
    (case (get name (get@ #.modules compiler))
      (#.Some module)
      (#e.Success [compiler module])

      _
      (#e.Error ($_ text/compose "Unknown module: " name)))))

(def: #export current-module-name
  (Meta Text)
  (function (_ compiler)
    (case (get@ #.current-module compiler)
      (#.Some current-module)
      (#e.Success [compiler current-module])

      _
      (#e.Error "No current module.")
      )))

(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."}
  (-> Name Code (Maybe Code))
  (case anns
    [_ (#.Record anns)]
    (loop [anns anns]
      (case anns
        (#.Cons [key value] anns')
        (case key
          [_ (#.Tag tag')]
          (if (name/= tag tag')
            (#.Some value)
            (recur anns'))

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

    _
    #.None))

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

       _
       #.None))]

  [get-bit-ann        #.Bit        Bit]
  [get-int-ann        #.Int        Int]
  [get-frac-ann       #.Frac       Frac]
  [get-text-ann       #.Text       Text]
  [get-identifier-ann #.Identifier Name]
  [get-tag-ann        #.Tag        Name]
  [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 (name-of #.doc) anns))

(def: #export (flag-set? flag-name anns)
  {#.doc "Finds out whether an annotation-as-a-flag is set (has value '#1')."}
  (-> Name Code Bit)
  (maybe.default #0 (get-bit-ann flag-name anns)))

(do-template [<name> <tag> <desc>]
  [(def: #export <name>
     {#.doc (code.text ($_ text/compose "Checks whether a definition is " <desc> "."))}
     (-> Code Bit)
     (flag-set? (name-of <tag>)))]

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

(def: #export (aliased? annotations)
  (-> Code Bit)
  (case (get-identifier-ann (name-of #.alias) annotations)
    (#.Some _)
    #1

    #.None
    #0))

(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 (name-of <tag>) anns)
                       args (parse-tuple _args)]
                      (monad.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 Definition) (|> (: Module $module) (get@ #.definitions) (get name)))]
    (if (and (macro? def-anns)
             (or (export? def-anns)
                 (text/= module this-module)))
      (#.Some (:coerce Macro def-value))
      (case (get-identifier-ann (name-of #.alias) def-anns)
        (#.Some [r-module r-name])
        (find-macro' modules this-module r-module r-name)

        _
        #.None))))

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

          Otherwise, returns the name as-is."}
  (-> Name (Meta Name))
  (case name
    ["" name]
    (do Monad<Meta>
      [module-name current-module-name]
      (wrap [module-name name]))

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

(def: #export (find-macro full-name)
  (-> Name (Meta (Maybe Macro)))
  (do Monad<Meta>
    [[module name] (normalize full-name)
     this-module current-module-name]
    (: (Meta (Maybe Macro))
       (function (_ compiler)
         (#e.Success [compiler (find-macro' (get@ #.modules compiler) this-module module name)])))))

(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 [[_ (#.Identifier name)] args]))]
    (do Monad<Meta>
      [?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 [[_ (#.Identifier name)] args]))]
    (do Monad<Meta>
      [?macro (find-macro name)]
      (case ?macro
        (#.Some macro)
        (do Monad<Meta>
          [expansion (macro args)
           expansion' (monad.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 [[_ (#.Identifier name)] args]))]
    (do Monad<Meta>
      [?macro (find-macro name)]
      (case ?macro
        (#.Some macro)
        (do Monad<Meta>
          [expansion (macro args)
           expansion' (monad.map Monad<Meta> expand-all expansion)]
          (wrap (list/join expansion')))
        
        #.None
        (do Monad<Meta>
          [parts' (monad.map Monad<Meta> expand-all (list& (code.identifier name) args))]
          (wrap (list (code.form (list/join parts')))))))

    [_ (#.Form (#.Cons [harg targs]))]
    (do Monad<Meta>
      [harg+ (expand-all harg)
       targs+ (monad.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' (monad.map Monad<Meta> expand-all members)]
      (wrap (list (code.tuple (list/join members')))))

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

(def: #export count
  (Meta Nat)
  (function (_ compiler)
    (#e.Success [(update@ #.seed inc compiler)
                 (get@ #.seed compiler)])))

(def: #export (gensym prefix)
  {#.doc "Generates a unique name 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 (_ compiler)
    (#e.Success [(update@ #.seed inc compiler)
                 (|> compiler
                     (get@ #.seed)
                     (:: number.Codec<Text,Nat> encode)
                     ($_ text/compose "__gensym__" prefix)
                     [""] code.identifier)])))

(def: (get-local-identifier ast)
  (-> Code (Meta Text))
  (case ast
    [_ (#.Identifier [_ name])]
    (:: Monad<Meta> wrap name)

    _
    (fail (text/compose "Code is not a local identifier: " (code.to-text ast)))))

(macro: #export (with-gensyms tokens)
  {#.doc (doc "Creates new identifiers 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 identifiers)] body))
    (do Monad<Meta>
      [identifier-names (monad.map @ get-local-identifier identifiers)
       #let [identifier-defs (list/join (list/map (: (-> Text (List Code))
                                                     (function (_ name) (list (code.identifier ["" name]) (` (gensym (~ (code.text name)))))))
                                                  identifier-names))]]
      (wrap (list (` ((~! do) (~! Monad<Meta>)
                      [(~+ identifier-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 Bit))
  (function (_ compiler)
    (#e.Success [compiler (case (get module (get@ #.modules compiler))
                            (#.Some _)
                            #1
                            
                            #.None
                            #0)])))

(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: (find-type-var idx bindings)
  (-> Nat (List [Nat (Maybe Type)]) (Maybe Type))
  (case bindings
    #.Nil
    #.None
    
    (#.Cons [var bound] bindings')
    (if (n/= idx var)
      bound
      (find-type-var idx bindings'))))

(def: (clean-type type)
  (-> Type (Meta Type))
  (case type
    (#.Var var)
    (function (_ compiler)
      (case (|> compiler
                (get@ [#.type-context #.var-bindings])
                (find-type-var var))
        (^or #.None (#.Some (#.Var _)))
        (#e.Success [compiler type])

        (#.Some type')
        (#e.Success [compiler type'])))

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

(def: #export (find-var-type name)
  {#.doc "Looks-up the type of a local variable somewhere in the environment."}
  (-> Text (Meta Type))
  (function (_ compiler)
    (let [test (: (-> [Text [Type Any]] Bit)
                  (|>> product.left (text/= name)))]
      (case (do maybe.Monad<Maybe>
              [scope (list.find (function (_ env)
                                  (or (list.any? test (: (List [Text [Type Any]])
                                                         (get@ [#.locals #.mappings] env)))
                                      (list.any? test (: (List [Text [Type Any]])
                                                         (get@ [#.captured #.mappings] env)))))
                                (get@ #.scopes compiler))
               [_ [type _]] (try-both (list.find test)
                                      (: (List [Text [Type Any]])
                                         (get@ [#.locals #.mappings] scope))
                                      (: (List [Text [Type Any]])
                                         (get@ [#.captured #.mappings] scope)))]
              (wrap type))
        (#.Some var-type)
        ((clean-type var-type) compiler)

        #.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)."}
  (-> Name (Meta Definition))
  (do Monad<Meta>
    [name (normalize name)]
    (function (_ compiler)
      (case (: (Maybe Definition)
               (do maybe.Monad<Maybe>
                 [#let [[v-prefix v-name] name]
                  (^slots [#.definitions]) (get v-prefix (get@ #.modules compiler))]
                 (get v-name definitions)))
        (#.Some definition)
        (#e.Success [compiler definition])

        _
        (let [current-module (|> compiler (get@ #.current-module) (maybe.default "???"))]
          (#e.Error ($_ text/compose
                        "Unknown definition: " (name/encode name) "\n"
                        "    Current module: " current-module "\n"
                        (case (get current-module (get@ #.modules compiler))
                          (#.Some this-module)
                          ($_ text/compose
                              "           Imports: " (|> this-module (get@ #.imports) (text.join-with ", ")) "\n"
                              "           Aliases: " (|> this-module (get@ #.module-aliases) (list/map (function (_ [alias real]) ($_ text/compose alias " => " real))) (text.join-with ", ")) "\n")

                          _
                          "")
                        " All Known modules: " (|> compiler (get@ #.modules) (list/map product.left) (text.join-with ", ")) "\n")))))))

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

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

      _
      (find-def-type name))))

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

(def: #export (definitions module-name)
  {#.doc "The entire list of definitions in a module (including the non-exported/private ones)."}
  (-> Text (Meta (List [Text Definition])))
  (function (_ compiler)
    (case (get module-name (get@ #.modules compiler))
      #.None          (#e.Error ($_ text/compose "Unknown module: " module-name))
      (#.Some module) (#e.Success [compiler (get@ #.definitions module)])
      )))

(def: #export (exports module-name)
  {#.doc "All the exported definitions in a module."}
  (-> Text (Meta (List [Text Definition])))
  (do Monad<Meta>
    [definitions (definitions module-name)]
    (wrap (list.filter (function (_ [name [def-type def-anns def-value]])
                         (export? def-anns))
                       definitions))))

(def: #export modules
  {#.doc "All the available modules (including the current one)."}
  (Meta (List [Text Module]))
  (function (_ compiler)
    (|> compiler
        (get@ #.modules)
        [compiler]
        #e.Success)))

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

      _
      (wrap #.None))))

(def: #export cursor
  {#.doc "The cursor of the current expression being analyzed."}
  (Meta Cursor)
  (function (_ compiler)
    (#e.Success [compiler (get@ #.cursor compiler)])))

(def: #export expected-type
  {#.doc "The expected type of the current expression being analyzed."}
  (Meta Type)
  (function (_ compiler)
    (case (get@ #.expected compiler)
      (#.Some type)
      (#e.Success [compiler 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 (imported-by? import module)
  (-> Text Text (Meta Bit))
  (do Monad<Meta>
    [(^slots [#.imports]) (find-module module)]
    (wrap (list.any? (text/= import) imports))))

(def: #export (imported? import)
  (-> Text (Meta Bit))
  (let [(^open ".") Monad<Meta>]
    (|> current-module-name
        (map find-module) join
        (map (|>> (get@ #.imports) (list.any? (text/= import)))))))

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

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

(def: #export (tag-lists module)
  {#.doc "All the tag-lists defined in a module, with their associated types."}
  (-> Text (Meta (List [(List Name) 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 (_ compiler)
    (case (list.inits (get@ #.scopes compiler))
      #.None
      (#e.Error "No local environment")

      (#.Some scopes)
      (#e.Success [compiler
                   (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."}
  (-> Name (Meta Name))
  (do Monad<Meta>
    [[_ def-anns _] (find-def def-name)]
    (case (get-identifier-ann (name-of #.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 Lux)
  (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 (: (Maybe [Bit Code])
              (case tokens
                (^ (list [_ (#.Tag ["" "omit"])]
                         token))
                (#.Some [#1 token])

                (^ (list token))
                (#.Some [#0 token])

                _
                #.None))
       (#.Some [omit? token])
       (do Monad<Meta>
         [cursor ..cursor
          output (<func> token)
          #let [_ (log! ($_ text/compose <desc> " @ " (.cursor-description cursor)))
                _ (list/map (|>> code.to-text log!)
                            output)
                _ (log! "")]]
         (wrap (if omit?
                 (list)
                 output)))

       #.None
       (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!"]
  )