(.module: {#.doc "Functions for extracting information from the state of the compiler."}
  [lux #*
   [control
    [functor (#+ Functor)]
    [apply (#+ Apply)]
    ["." monad (#+ Monad do)]]
   [data
    ["." product]
    ["." name ("#/." codec equivalence)]
    ["." maybe]
    ["." error (#+ Error)]
    [number
     ["." nat ("#/." decimal)]]
    ["." text ("#/." monoid equivalence)]
    [collection
     ["." list ("#/." monoid monad)]]]]
  [/
   ["." code]])

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

(structure: #export functor (Functor Meta)
  (def: (map f fa)
    (function (_ compiler)
      (case (fa compiler)
        (#error.Failure msg)
        (#error.Failure msg)

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

(structure: #export apply (Apply Meta)
  (def: &functor ..functor)

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

          (#error.Failure msg)
          (#error.Failure msg))

        (#error.Failure msg)
        (#error.Failure msg)))))

(structure: #export monad (Monad Meta)
  (def: &functor ..functor)

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

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

(def: #export (run compiler action)
  (All [a] (-> Lux (Meta a) (Error a)))
  (case (action compiler)
    (#error.Failure error)
    (#error.Failure error)

    (#error.Success [_ output])
    (#error.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)
      (#error.Failure error)
      (right compiler)

      (#error.Success [compiler' output])
      (#error.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
      (#error.Success [compiler []])
      (#error.Failure message))))

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

(def: #export (find-module name)
  (-> Text (Meta Module))
  (function (_ compiler)
    (case (get name (get@ #.modules compiler))
      (#.Some module)
      (#error.Success [compiler module])

      _
      (#error.Failure ($_ text/compose "Unknown module: " name)))))

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

      _
      (#error.Failure "No current module.")
      )))

(def: #export current-module
  (Meta Module)
  (do ..monad
    [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-documentation 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"]
  [structure?      #.struct?   "a structure"]
  [recursive-type? #.type-rec? "a recursive type"]
  [signature?      #.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
                      [_args (get-ann (name-of <tag>) anns)
                       args (parse-tuple _args)]
                      (monad.map @ parse-text args))))]

  [function-arguments #.func-args "Looks up the arguments of a function."]
  [type-arguments     #.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
    [$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 (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
      [module-name current-module-name]
      (wrap [module-name name]))

    _
    (:: ..monad wrap name)))

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

(def: #export (expand-once syntax)
  {#.doc (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
      [?macro (find-macro name)]
      (case ?macro
        (#.Some macro)
        (macro args)
        
        #.None
        (:: ..monad wrap (list syntax))))

    _
    (:: ..monad wrap (list syntax))))

(def: #export (expand syntax)
  {#.doc (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
      [?macro (find-macro name)]
      (case ?macro
        (#.Some macro)
        (do ..monad
          [expansion (macro args)
           expansion' (monad.map ..monad expand expansion)]
          (wrap (list/join expansion')))
        
        #.None
        (:: ..monad wrap (list syntax))))

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

    [_ (#.Form (#.Cons [harg targs]))]
    (do ..monad
      [harg+ (expand-all harg)
       targs+ (monad.map ..monad expand-all targs)]
      (wrap (list (code.form (list/compose harg+ (list/join (: (List (List Code)) targs+)))))))

    [_ (#.Tuple members)]
    (do ..monad
      [members' (monad.map ..monad expand-all members)]
      (wrap (list (code.tuple (list/join members')))))

    _
    (:: ..monad wrap (list syntax))))

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

(def: #export (gensym prefix)
  {#.doc (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)
    (#error.Success [(update@ #.seed inc compiler)
                     (|> compiler
                         (get@ #.seed)
                         (:: nat.decimal encode)
                         ($_ text/compose "__gensym__" prefix)
                         [""] code.identifier)])))

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

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

(def: #export wrong-syntax-error
  (-> Name Text)
  (|>> name/encode
       (text/compose "Wrong syntax for ")))

(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
      [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)
                      [(~+ identifier-defs)]
                      (~ body))))))

    _
    (fail (..wrong-syntax-error (name-of ..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
    [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)
    (#error.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 _)))
        (#error.Success [compiler type])

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

    _
    (:: ..monad 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
              [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
        (#error.Failure ($_ 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
    [name (normalize name)]
    (function (_ compiler)
      (case (: (Maybe Definition)
               (do maybe.monad
                 [#let [[v-prefix v-name] name]
                  (^slots [#.definitions]) (get v-prefix (get@ #.modules compiler))]
                 (get v-name definitions)))
        (#.Some definition)
        (#error.Success [compiler definition])

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

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

(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
    [[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
    [#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
    [[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          (#error.Failure ($_ text/compose "Unknown module: " module-name))
      (#.Some module) (#error.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
    [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]
        #error.Success)))

(def: #export (tags-of type-name)
  {#.doc "All the tags associated with a type definition."}
  (-> Name (Meta (Maybe (List Name))))
  (do ..monad
    [#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)
    (#error.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)
      (#error.Success [compiler type])

      #.None
      (#error.Failure "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
    [(^slots [#.imports]) (find-module module-name)]
    (wrap imports)))

(def: #export (imported-by? import module)
  (-> Text Text (Meta Bit))
  (do ..monad
    [(^slots [#.imports]) (find-module module)]
    (wrap (list.any? (text/= import) imports))))

(def: #export (imported? import)
  (-> Text (Meta Bit))
  (let [(^open ".") ..monad]
    (|> 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
    [#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
    [=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
      (#error.Failure "No local environment")

      (#.Some scopes)
      (#error.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
    [[_ 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)
    (#error.Success [compiler compiler])))

(def: #export type-context
  (Meta Type-Context)
  (function (_ compiler)
    (#error.Success [compiler (get@ #.type-context compiler)])))

(do-template [<macro> <func>]
  [(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
         [cursor ..cursor
          output (<func> token)
          #let [_ (log! ($_ text/compose (name/encode (name-of <macro>)) " @ " (.cursor-description cursor)))
                _ (list/map (|>> code.to-text log!)
                            output)
                _ (log! "")]]
         (wrap (if omit?
                 (list)
                 output)))

       #.None
       (fail (..wrong-syntax-error (name-of <macro>)))))]

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