##  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 [cli #+ program:]
       [compiler #+ Monad<Lux>]
       host
       [lexer]
       [macro]
       [math]
       [pipe]
       [random]
       [regex]
       [test]
       [type]
       (codata [cont]
               [env]
               function
               [io #- run]
               [state]
               (struct [stream #+ Stream "Stream/" Functor<Stream>]))
       (concurrency [actor]
                    [atom]
                    [frp]
                    [promise]
                    [stm])
       (control [applicative]
                [bounded]
                [codec]
                [comonad]
                [effect]
                [enum]
                [eq]
                [fold]
                [functor]
                [hash]
                monad
                [monoid]
                ["_;" number]
                [ord])
       (data [bit]
             [bool]
             [char]
             [error]
             [ident "Ident/" Codec<Text,Ident>]
             [identity]
             [log]
             maybe
             [number #* "Nat/" Codec<Text,Nat>]
             [product]
             [sum]
             [text "Text/" Monoid<Text> Eq<Text>]
             (error [exception])
             (format [json])
             (struct [array]
                     [dict]
                     [list #+ "List/" Monoid<List> Functor<List> Fold<List>]
                     [queue]
                     [set]
                     [stack]
                     [tree]
                     [vector]
                     [zipper])
             text/format)
       (macro [ast]
              [poly]
              ["s" syntax]
              (poly ["poly_;" eq]
                    ["poly_;" functor]
                    ["poly_;" text-encoder])
              (syntax [common]))
       (math [complex]
             [ratio]
             [simple])
       (type [auto]
             [check])
       ))

(def: name-options "abcdefghijklmnopqrstuvwxyz")
(def: name-options-count (text;size name-options))

(def: (id->name id)
  (-> Nat Text)
  (if (n.> name-options-count id)
    (format (id->name (n./ name-options-count id))
            (id->name (n.% name-options-count id)))
    (char;as-text (default #"?" (text;at id name-options)))))

(def: type-var-names
  (Stream Text)
  (Stream/map id->name (stream;iterate n.inc +0)))

(def: (type-arg? id)
  (-> Nat Bool)
  (n.= +1 (n.% +2 id)))

(def: (arg-id level id)
  (-> Nat Nat Nat)
  (n.- (n./ +2 id) level))

(def: (bound->name [type-fun-name type-fun-args] level id)
  (-> [Text (List Text)] Nat Nat Text)
  (if (type-arg? id)
    (let [_arg-id (arg-id level id)]
      (case (list;at _arg-id type-fun-args)
        (#;Some found)
        found

        _
        (|> type-var-names
            (stream;filter (lambda [var-name] (not (list;member? text;Eq<Text> type-fun-args var-name))))
            (stream;at _arg-id))))      
    type-fun-name))

(do-template [<name> <base> <tag>]
  [(def: (<base> level type)
     (-> Nat Type [Nat Type])
     (case type
       (<tag> env type')
       (<base> (n.inc level) type')

       _
       [level type]))

   (def: <name> (<base> +0))]

  [unravel-univ unravel-univ' #;UnivQ]
  [unravel-ex   unravel-ex'   #;ExQ])

(def: (level->args offset level)
  (-> Nat Nat (List Text))
  (if (n.= +0 level)
    (list)
    (|> level
        n.dec
        (list;n.range +0)
        (List/map (|>. (n.+ (n.inc offset)) id->name)))))

(def: (prefix-lines prefix lines)
  (-> Text Text Text)
  (|> lines
      text;split-lines
      (List/map (Text/append prefix))
      ## (list;interpose "\n")
      (text;join-with "")))

(def: (pprint-type-def level type-fun-info tags module sig? type-rec? type)
  (-> Nat [Text (List Text)] (List Ident) Text Bool Bool Type Text)
  (case tags
    (^ (list [_ single-tag]))
    (if sig?
      (format "(: " (pprint-type-def level type-fun-info #;None module sig? type-rec? type) "\n   " single-tag ")")
      (format "{#" single-tag " " (pprint-type-def level type-fun-info #;None module sig? type-rec? type) "}"))

    _
    (case [type-rec? type]
      [_ (#;HostT name params)]
      (case params
        #;Nil
        (format "(host " name ")")

        _
        (format "(host " name " " (|> params (List/map (pprint-type-def level type-fun-info #;None module sig? type-rec?)) (text;join-with " ")) ")"))

      [_ #;VoidT]
      "Void"

      [_ #;UnitT]
      "Unit"

      [_ (#;SumT _)]
      (let [members (type;flatten-variant type)]
        (case tags
          #;Nil
          (format "(| "
                  (|> members
                      (List/map (pprint-type-def level type-fun-info #;None module sig? type-rec?))
                      (text;join-with " "))
                  ")")

          _
          (|> members
              (list;zip2 tags)
              (List/map (lambda [[[_ t-name] type]]
                          (case type
                            #;UnitT
                            (format "#" t-name)

                            (#;ProdT _)
                            (let [types (type;flatten-tuple type)]
                              (format "(#" t-name " "
                                      (|> types
                                          (List/map (pprint-type-def level type-fun-info #;None module sig? type-rec?))
                                          (text;join-with " "))
                                      ")"))

                            _
                            (format "(#" t-name " " (pprint-type-def level type-fun-info #;None module sig? type-rec? type) ")"))))
              (text;join-with "\n"))))

      [_ (#;ProdT _)]
      (let [members (type;flatten-tuple type)]
        (case tags
          #;Nil
          (format "[" (|> members (List/map (pprint-type-def level type-fun-info #;None module sig? type-rec?)) (text;join-with " ")) "]")

          _
          (let [member-docs (|> members
                                (list;zip2 tags)
                                (List/map (lambda [[[_ t-name] type]]
                                            (if sig?
                                              (format "(: " (pprint-type-def level type-fun-info #;None module sig? type-rec? type) "\n   "  t-name ")")
                                              (format "#" t-name " " (pprint-type-def level type-fun-info #;None module sig? type-rec? type)))))
                                (text;join-with "\n "))]
            (if sig?
              member-docs
              (format "{" member-docs "}")))))

      [_ (#;LambdaT input output)]
      (let [[ins out] (type;flatten-function type)]
        (format  "(-> " (|> ins (List/map (pprint-type-def level type-fun-info #;None module sig? type-rec?)) (text;join-with " "))
                 " "
                 (pprint-type-def level type-fun-info #;None module sig? type-rec? out)
                 ")"))

      [_ (#;BoundT idx)]
      (bound->name type-fun-info level idx)

      (^template [<tag> <pre> <post>]
        [_ (<tag> id)]
        (format <pre> (Nat/encode id) <post>))
      ([#;VarT "⌈v:" "⌋"]
       [#;ExT  "⟨e:" "⟩"])
      
      (^template [<tag> <name> <unravel>]
        [_ (<tag> _)]
        (let [[level' body] (<unravel> type)
              args (level->args level level')
              body-doc (pprint-type-def (n.+ level level') type-fun-info tags module sig? type-rec? body)]
          (format "(" <name> " " "[" (text;join-with " " args) "]"
                  (case tags
                    #;Nil
                    (format " " body-doc)

                    _
                    (format "\n" (prefix-lines "  " body-doc)))
                  ")")))
      ([#;UnivQ "All" unravel-univ]
       [#;ExQ   "Ex"  unravel-ex])

      [true (#;AppT (#;BoundT +0) (#;BoundT +1))]
      (product;left type-fun-info)
      
      [_ (#;AppT fun param)]
      (let [[type-fun type-args] (type;flatten-application type)]
        (format  "(" (pprint-type-def level type-fun-info tags module sig? type-rec? type-fun) " " (|> type-args (List/map (pprint-type-def level type-fun-info #;None module sig? type-rec?)) (text;join-with " ")) ")"))
      
      [_ (#;NamedT [_module _name] type)]
      (if (Text/= module _module)
        _name
        (Ident/encode [_module _name]))
      )))

(def: (pprint-type level type-fun-name module type)
  (-> Nat Text Text Type Text)
  (case type
    (#;HostT name params)
    (case params
      #;Nil
      (format "(host " name ")")

      _
      (format "(host " name " " (|> params (List/map (pprint-type level type-fun-name module)) (list;interpose " ") (text;join-with "")) ")"))

    #;VoidT
    "Void"

    (#;SumT _)
    (let [members (type;flatten-variant type)]
      (format "(| " (|> members (List/map (pprint-type level type-fun-name module)) (list;interpose " ") (text;join-with "")) ")"))

    #;UnitT
    "Unit"

    (#;ProdT _)
    (let [members (type;flatten-tuple type)]
      (format "[" (|> members (List/map (pprint-type level type-fun-name module)) (list;interpose " ") (text;join-with "")) "]"))

    (#;LambdaT input output)
    (let [[ins out] (type;flatten-function type)]
      (format  "(-> "
               (|> ins (List/map (pprint-type level type-fun-name module)) (list;interpose " ") (text;join-with ""))
               " "
               (pprint-type level type-fun-name module out)
               ")"))

    (#;BoundT idx)
    (bound->name [type-fun-name (list)] level idx)

    (^template [<tag> <pre> <post>]
      (<tag> id)
      (format <pre> (Nat/encode id) <post>))
    ([#;VarT "⌈" "⌋"]
     [#;ExT  "⟨" "⟩"])

    (^template [<tag> <name> <unravel>]
      (<tag> _)
      (let [[level' body] (<unravel> type)
            args (level->args level level')
            body-doc (pprint-type (n.+ level level') type-fun-name module body)]
        (format "(" <name> " " "[" (|> args (list;interpose " ") (text;join-with "")) "]"
                (format " " body-doc)
                ")")))
    ([#;UnivQ "All" unravel-univ]
     [#;ExQ   "Ex"  unravel-ex])

    (#;AppT fun param)
    (let [[type-fun type-args] (type;flatten-application type)]
      (format  "(" (pprint-type level type-fun-name module type-fun) " " (|> type-args (List/map (pprint-type level type-fun-name module)) (list;interpose " ") (text;join-with "")) ")"))

    (#;NamedT [_module _name] type)
    (if (Text/= module _module)
      _name
      (Ident/encode [_module _name]))
    ))

(type: Markdown
  Text)

(type: DefOrg
  {#types (List [Text Anns Type])
   #macros (List [Text Anns])
   #structs (List [Text Anns Type])
   #values (List [Text Anns Type])})

(def: (lux-module? module-name)
  (-> Text Bool)
  (or (Text/= "lux" module-name)
      (text;starts-with? "lux/" module-name)))

(do-template [<name> <joiner>]
  [(def: (<name> docs)
     (-> (List Markdown) Markdown)
     (text;join-with "\n\n" docs))]

  [join-def-docs     "\n\n"]
  [join-doc-sections "\n\n"]
  )

(def: (cons h t)
  (All [a] (-> a (List a) (List a)))
  (#;Cons h t))

(def: (add-def [name [def-type def-meta def-value]] org)
  (-> [Text Def] DefOrg DefOrg)
  (cond (compiler;type? def-meta)
        (update@ #types (cons [name def-meta (:coerce Type def-value)]) org)
        
        (compiler;macro? def-meta)
        (update@ #macros (cons [name def-meta]) org)

        (compiler;struct? def-meta)
        (update@ #structs (cons [name def-meta def-type]) org)

        ## else
        (update@ #values (cons [name def-meta def-type]) org)))

(def: def-sorter
  (All [r] (-> (List [Text r]) (List [Text r])))
  (list;sort (: (All [r] (-> [Text r] [Text r] Bool))
                (lambda [[n1 _] [n2 _]] (:: text;Ord<Text> < n1 n2)))))

(def: (organize-defs defs)
  (-> (List [Text Def]) DefOrg)
  (let [init {#types (list)
              #macros (list)
              #structs (list)
              #values (list)}]
    (|> (List/fold add-def init defs)
        (update@ #types def-sorter)
        (update@ #macros def-sorter)
        (update@ #structs def-sorter)
        (update@ #values def-sorter))))

(def: (unravel-type-func level type)
  (-> Nat Type Type)
  (if (n.> +0 level)
    (case type
      (#;UnivQ _env _type)
      (unravel-type-func (n.dec level) _type)

      _
      type)
    type))

(def: (unrecurse-type type)
  (-> Type Type)
  (case type
    (#;AppT (#;UnivQ _env _type) _)
    _type

    _
    type))

(def: #export (when! test f x)
  (All [a] (-> Bool (-> a a) a a))
  (if test
    (f x)
    x))

(def: #export (if! test f g x)
  (All [a b] (-> Bool (-> a b) (-> a b) a b))
  (if test
    (f x)
    (g x)))

(def: (doc-type module type def-meta)
  (-> Text Type Anns (Lux Markdown))
  (case type
    (#;NamedT type-name type)
    (do Monad<Lux>
      [tags (compiler;tags-of type-name)
       #let [[_ _name] type-name
             type-rec? (compiler;type-rec? def-meta)
             type-args (compiler;type-args def-meta)
             sig? (compiler;sig? def-meta)
             usage (case type-args
                     #;Nil
                     _name

                     _
                     (format "(" (text;join-with " " (list& _name type-args)) ")"))
             nesting (list;size type-args)]]
      (wrap (format (if sig? "(signature: " "(type: ") (if type-rec? "#rec " "") usage "\n"
                    (|> (pprint-type-def (n.dec nesting)
                                         [_name type-args]
                                         tags module sig? type-rec?
                                         (|> type
                                             (unravel-type-func nesting)
                                             (when! type-rec? unrecurse-type)))
                        text;split-lines
                        (List/map (Text/append "  "))
                        (text;join-with "\n"))
                    ")")))

    _
    (compiler;fail (format "A type definition must always be named! - " (type;to-text type)))))

(def: (doc-types module types)
  (-> Text (List [Text Anns Type]) (Lux Markdown))
  (do Monad<Lux>
    [type-docs (mapM @
                     (: (-> [Text Anns Type] (Lux Markdown))
                        (lambda [[name def-meta type]]
                          (do Monad<Lux>
                            [#let [?doc (compiler;get-doc def-meta)
                                   name (|> name
                                            (text;replace "<" "&lt;")
                                            (text;replace ">" "&gt;"))]
                             type-doc (doc-type module type def-meta)]
                            (wrap (format "### " name "\n"
                                          (case ?doc
                                            (#;Some doc)
                                            (format doc "\n")

                                            _
                                            "")
                                          "```\n" type-doc "\n```")))))
                     types)]
    (|> type-docs
        join-def-docs
        (format "## Types\n")
        wrap)))

(def: (doc-macros module-name names)
  (-> Text (List [Text Anns]) Markdown)
  (|> names
      (List/map (: (-> [Text Anns] Markdown)
                   (lambda [[name def-meta]]
                     (let [name (|> name
                                    (text;replace "<" "&lt;")
                                    (text;replace ">" "&gt;"))]
                       (format "### " name "\n"
                               (default ""
                                 (do Monad<Maybe>
                                   [doc (compiler;get-doc def-meta)]
                                   (wrap (format "```\n" doc "\n```")))))))))
      join-def-docs
      (format "## Macros\n")))

(do-template [<singular> <plural> <header>]
  [(def: <singular>
     (-> Text Type Markdown)
     (pprint-type (n.dec +0) "?"))

   (def: (<plural> module values)
     (-> Text (List [Text Anns Type]) Markdown)
     (|> values
         (List/map (lambda [[name def-meta value-type]]
                      (let [?doc (compiler;get-doc def-meta)
                            usage (case (compiler;func-args def-meta)
                                    #;Nil
                                    name

                                    args
                                    (format "(" (text;join-with " " (list& name args)) ")"))
                            usage (|> usage
                                      (text;replace "<" "&lt;")
                                      (text;replace ">" "&gt;"))]
                        (format "### " usage "\n"
                                (case ?doc
                                  (#;Some doc)
                                  (format "```\n" doc "\n```\n")

                                  _
                                  "")
                                "`" (<singular> module value-type) "`"))))
         join-def-docs
         (format <header>)))]

  [doc-struct doc-structs "## Structs\n"]
  [doc-value  doc-values  "## Values\n"]
  )

(def: (enclose-lines pre+post block)
  (-> [Text Text] Text Text)
  (|> block
      text;split-lines
      (List/map (text;enclose pre+post))
      (text;join-with "\n")))

(def: (doc-module [[module-name module] org])
  (-> [[Text Module] DefOrg] (Lux [Text Markdown]))
  (do Monad<Lux>
    [#let [(^slots [#types #macros #structs #values]) org
           anns (|> module (get@ #;module-anns) compiler;get-doc)]
     types-md (if (list;empty? types)
                (wrap "")
                (doc-types module-name types))
     #let [doc-desc (case anns
                      #;None
                      ""

                      (#;Some doc)
                      (format "\n" (enclose-lines ["> " ""] doc) "\n"))
           doc-body (join-doc-sections (list types-md
                                             (if (list;empty? macros) "" (doc-macros module-name macros))
                                             (if (list;empty? structs) "" (doc-structs module-name structs))
                                             (if (list;empty? values) "" (doc-values module-name values))))]]
    (wrap [module-name
           (format "## " module-name "\n"
                   doc-desc "\n"
                   doc-body)])))

(jvm-import java.io.File
  (new [java.lang.String]))

(jvm-import java.io.PrintWriter
  (new [java.io.File] #io #try)
  (println [java.lang.String] #io void))

(jvm-import java.io.Writer
  (flush [] #io #try void))

(def: (save-docs! [module-name docs])
  (-> [Text Markdown] (IO Unit))
  (do Monad<IO>
    [?target (|> (format (text;replace "/" "_" module-name) ".md")
                 File.new
                 PrintWriter.new)]
    (case ?target
      (#;Left _)
      (wrap [])
      
      (#;Right target)
      (do @
        [_ (PrintWriter.println docs target)
         _ (Writer.flush [] target)]
        (wrap [])))))

(macro: (gen-docs! _)
  (do Monad<Lux>
    [all-modules compiler;modules
     #let [lux-modules (|> all-modules
                           (list;filter (. lux-module? product;left))
                           (list;sort (lambda [[left _] [right _]]
                                        (:: text;Ord<Text> < left right))))]
     lux-exports (mapM @ (. compiler;exports product;left) lux-modules)
     module-docs (mapM @ doc-module
                       (list;zip2 lux-modules
                                  (List/map organize-defs lux-exports)))
     #let [_ (io;run (mapM Monad<IO> save-docs! module-docs))]]
    (wrap (list))))

(gen-docs!)

(program: args
  (io (log! "Done!")))