(.module:
  [library
   [lux #*
    [program (#+ program:)]
    ["." type ("#\." equivalence)]
    ["." debug]
    [abstract
     ["." monad (#+ do)]
     ["." enum]]
    [control
     [pipe (#+ when> new>)]
     ["." try (#+ Try)]
     ["." exception (#+ exception:)]
     ["." io (#+ IO io)]
     ["." function]]
    [data
     ["." maybe]
     ["." product]
     [format
      ["md" markdown (#+ Markdown Span Block)]]
     ["." text ("#\." equivalence)
      ["%" format (#+ format)]
      [encoding
       ["." utf8]]]
     [collection
      ["." sequence (#+ Sequence) ("#\." functor)]
      ["." list ("#\." functor fold)]]]
    [math
     [number
      ["n" nat]]]
    ["." meta
     ["." annotation]]
    [world
     ["." file]]]]
  ## This was added to make sure that all tested modules are picked up
  ## and their documentation is generated.
  [test/lux (#+)])

(def: name_options "abcdefghijklmnopqrstuvwxyz")
(def: name_options_count (text.size name_options))

(def: (parameter_type_name id)
  (-> Nat Text)
  (case (text.nth id ..name_options)
    (#.Some char)
    (text.of_code char)

    #.None
    (format (parameter_type_name (n./ name_options_count id))
            (parameter_type_name (n.% name_options_count id)))))

(def: type_var_names
  (Sequence Text)
  (|> 0 (sequence.iterate inc) (sequence\map parameter_type_name)))

(template [<name> <partition>]
  [(def: (<name> id)
     (-> Nat Bit)
     (<partition> id))]

  [type_func? n.even?]
  [type_arg? n.odd?]
  )

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

(def: (parameter_to_name [type_func_name type_function_arguments] level id)
  (-> [Text (List Text)] Nat Nat Text)
  (if (type_arg? id)
    (let [arg_id (..arg_id level id)]
      (case (list.nth arg_id type_function_arguments)
        (#.Some found)
        found

        _
        (|> type_var_names
            (sequence.only (function (_ var_name)
                             (not (list.member? text.equivalence type_function_arguments var_name))))
            (sequence.nth arg_id))))      
    type_func_name))

(def: (level_to_args offset level)
  (-> Nat Nat (List Text))
  (if (n.= 0 level)
    (list)
    (|> level
        dec
        (enum.range n.enum 0)
        (list\map (|>> (n.+ (inc offset)) parameter_type_name)))))

(def: (prefix_lines prefix lines)
  (-> Text Text Text)
  (|> lines
      (text.split_all_with text.new_line)
      (list\map (|>> (format prefix)))
      (text.join_with text.new_line)))

(def: (pprint_type_definition level type_func_info tags module signature? recursive_type? type)
  (-> Nat [Text (List Text)] (List Name) Text Bit Bit Type Text)
  (case tags
    (^ (list [_ single_tag]))
    (if signature?
      (format "(: " (pprint_type_definition level type_func_info #.None module signature? recursive_type? type) text.new_line "   " single_tag ")")
      (format "{#" single_tag " " (pprint_type_definition level type_func_info #.None module signature? recursive_type? type) "}"))

    _
    (case [recursive_type? type]
      [_ (#.Primitive name params)]
      (case params
        #.Nil
        (format "(primitive " (%.text name) ")")

        _
        (format "(primitive " (%.text name) " " (|> params (list\map (pprint_type_definition level type_func_info #.None module signature? recursive_type?)) (text.join_with " ")) ")"))

      [_ (#.Sum _)]
      (let [members (type.flat_variant type)]
        (case tags
          #.Nil
          (format "(| "
                  (|> members
                      (list\map (pprint_type_definition level type_func_info #.None module signature? recursive_type?))
                      (text.join_with " "))
                  ")")

          _
          (|> members
              (list.zipped/2 tags)
              (list\map (function (_ [[_ t_name] type])
                          (case type
                            (#.Product _)
                            (let [types (type.flat_tuple type)]
                              (format "(#" t_name " "
                                      (|> types
                                          (list\map (pprint_type_definition level type_func_info #.None module signature? recursive_type?))
                                          (text.join_with " "))
                                      ")"))

                            _
                            (format "(#" t_name " " (pprint_type_definition level type_func_info #.None module signature? recursive_type? type) ")"))))
              (text.join_with text.new_line))))

      [_ (#.Product _)]
      (let [members (type.flat_tuple type)]
        (case tags
          #.Nil
          (format "[" (|> members (list\map (pprint_type_definition level type_func_info #.None module signature? recursive_type?)) (text.join_with " ")) "]")

          _
          (let [member_docs (|> members
                                (list.zipped/2 tags)
                                (list\map (function (_ [[_ t_name] type])
                                            (if signature?
                                              (format "(: " (pprint_type_definition level type_func_info #.None module signature? recursive_type? type) text.new_line "   "  t_name ")")
                                              (format "#" t_name " " (pprint_type_definition level type_func_info #.None module signature? recursive_type? type)))))
                                (text.join_with (format text.new_line " ")))]
            (if signature?
              member_docs
              (format "{" member_docs "}")))))

      [_ (#.Function input output)]
      (let [[ins out] (type.flat_function type)]
        (format  "(-> " (|> ins (list\map (pprint_type_definition level type_func_info #.None module signature? recursive_type?)) (text.join_with " "))
                 " "
                 (pprint_type_definition level type_func_info #.None module signature? recursive_type? out)
                 ")"))

      [_ (#.Parameter idx)]
      (parameter_to_name type_func_info level idx)

      (^template [<tag> <pre> <post>]
        [[_ (<tag> id)]
         (format <pre> (%.nat id) <post>)])
      ([#.Var "⌈v:" "⌋"]
       [#.Ex  "⟨e:" "⟩"])

      (^template [<tag> <name> <flatten>]
        [[_ (<tag> _)]
         (let [[level' body] (<flatten> type)
               args (level_to_args level level')
               body_doc (pprint_type_definition (n.+ level level') type_func_info tags module signature? recursive_type? body)]
           (format "(" <name> " " "[" (text.join_with " " args) "]"
                   (case tags
                     #.Nil
                     (format " " body_doc)

                     _
                     (format text.new_line (prefix_lines "  " body_doc)))
                   ")"))])
      ([#.UnivQ "All" type.flat_univ_q]
       [#.ExQ   "Ex"  type.flat_ex_q])

      [true (#.Apply (#.Parameter 1) (#.Parameter 0))]
      (product.left type_func_info)

      [_ (#.Apply param fun)]
      (let [[type_func type_arguments] (type.flat_application type)]
        (format  "(" (pprint_type_definition level type_func_info tags module signature? recursive_type? type_func) " " (|> type_arguments (list\map (pprint_type_definition level type_func_info #.None module signature? recursive_type?)) (text.join_with " ")) ")"))

      [_ (#.Named [_module _name] type)]
      (if (text\= module _module)
        _name
        (%.name [_module _name]))
      )))

(def: (pprint_type level type_func_name module type)
  (-> Nat Text Text Type Text)
  (case type
    (#.Primitive name params)
    (case params
      #.Nil
      (format "(primitive " (%.text name) ")")

      _
      (format "(primitive " (%.text name) " " (|> params (list\map (pprint_type level type_func_name module)) (list.interpose " ") (text.join_with "")) ")"))

    (#.Sum _)
    (let [members (type.flat_variant type)]
      (format "(| " (|> members (list\map (pprint_type level type_func_name module)) (list.interpose " ") (text.join_with "")) ")"))

    (#.Product _)
    (let [members (type.flat_tuple type)]
      (format "[" (|> members (list\map (pprint_type level type_func_name module)) (list.interpose " ") (text.join_with "")) "]"))

    (#.Function input output)
    (let [[ins out] (type.flat_function type)]
      (format  "(-> "
               (|> ins (list\map (pprint_type level type_func_name module)) (list.interpose " ") (text.join_with ""))
               " "
               (pprint_type level type_func_name module out)
               ")"))

    (#.Parameter idx)
    (parameter_to_name [type_func_name (list)] level idx)

    (^template [<tag> <pre> <post>]
      [(<tag> id)
       (format <pre> (%.nat id) <post>)])
    ([#.Var "⌈" "⌋"]
     [#.Ex  "⟨" "⟩"])

    (^template [<tag> <name> <flatten>]
      [(<tag> _)
       (let [[level' body] (<flatten> type)
             args (level_to_args level level')
             body_doc (pprint_type (n.+ level level') type_func_name module body)]
         (format "(" <name> " " "[" (|> args (list.interpose " ") (text.join_with "")) "]"
                 (format " " body_doc)
                 ")"))])
    ([#.UnivQ "All" type.flat_univ_q]
     [#.ExQ   "Ex"  type.flat_ex_q])

    (#.Apply param fun)
    (let [[type_func type_arguments] (type.flat_application type)]
      (format  "(" (pprint_type level type_func_name module type_func) " " (|> type_arguments (list\map (pprint_type level type_func_name module)) (list.interpose " ") (text.join_with "")) ")"))

    (#.Named [_module _name] type)
    (if (text\= module _module)
      _name
      (%.name [_module _name]))
    ))

(type: (Mutation a)
  (-> a a))

(type: Value
  [Text Code Type])

(type: Organization
  {#types (List Value)
   #macros (List [Text Code])
   #implementations (List Value)
   #values (List Value)})

(def: (lux_module? module_name)
  (-> Text Bit)
  (let [prefix (format .prelude_module "/")]
    (or (text\= .prelude_module module_name)
        (text.starts_with? prefix module_name))))

(def: (add_definition [name [exported? def_type def_annotations def_value]] organization)
  (-> [Text Definition] Organization Organization)
  (cond (type\= .Type def_type)
        (update@ #types
                 (: (Mutation (List Value))
                    (|>> (#.Cons [name def_annotations (:as Type def_value)])))
                 organization)

        (type\= .Macro def_type)
        (update@ #macros
                 (: (Mutation (List [Text Code]))
                    (|>> (#.Cons [name def_annotations])))
                 organization)

        (annotation.implementation? def_annotations)
        (update@ #implementations
                 (: (Mutation (List Value))
                    (|>> (#.Cons [name def_annotations def_type])))
                 organization)

        ## else
        (update@ #values
                 (: (Mutation (List Value))
                    (|>> (#.Cons [name def_annotations def_type])))
                 organization)))

(def: name_sort
  (All [r] (-> [Text r] [Text r] Bit))
  (let [text\< (\ text.order <)]
    (function (_ [n1 _] [n2 _])
      (text\< n1 n2))))

(def: (organize_definitions defs)
  (-> (List [Text Definition]) Organization)
  (let [init {#types (list)
              #macros (list)
              #implementations (list)
              #values (list)}]
    (|> (list\fold add_definition init defs)
        (update@ #types (list.sort name_sort))
        (update@ #macros (list.sort name_sort))
        (update@ #implementations (list.sort name_sort))
        (update@ #values (list.sort name_sort)))))

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

      _
      type)
    type))

(def: (unrecurse_type type)
  (-> Type Type)
  (case type
    (#.Apply _ (#.UnivQ _env _type))
    _type

    _
    type))

(exception: #export (anonymous_type_definition {type Type})
  (exception.report
   ["Type" (%.type type)]))

(def: (document_type module type def_annotations)
  (-> Text Type Code (Meta (Markdown Block)))
  (case type
    (#.Named type_name type)
    (do meta.monad
      [tags (meta.tags_of type_name)
       #let [[_ _name] type_name
             recursive_type? (annotation.recursive_type? def_annotations)
             type_arguments (annotation.type_arguments def_annotations)
             signature? (annotation.signature? def_annotations)
             usage (case type_arguments
                     #.Nil
                     _name

                     _
                     (format "(" (text.join_with " " (list& _name type_arguments)) ")"))
             nesting (list.size type_arguments)]]
      (in (md.code (format (if signature? "(interface: " "(type: ")
                           (if recursive_type? "#rec " "")
                           usage text.new_line
                           (|> type
                               (unravel_type_func nesting)
                               (when> [(new> recursive_type? [])] [unrecurse_type])
                               (pprint_type_definition (dec nesting) [_name type_arguments] (maybe.default (list) tags) module signature? recursive_type?)
                               (text.split_all_with text.new_line)
                               (list\map (|>> (format "  ")))
                               (text.join_with text.new_line))
                           ")"))))

    _
    (meta.failure (exception.construct anonymous_type_definition type))))

(def: (document_types module types)
  (-> Text (List Value) (Meta (Markdown Block)))
  (do {! meta.monad}
    [type_docs (monad.map !
                          (: (-> Value (Meta (Markdown Block)))
                             (function (_ [name def_annotations type])
                               (do meta.monad
                                 [#let [?doc (annotation.documentation def_annotations)]
                                  type_code (document_type module type def_annotations)]
                                 (in ($_ md.then
                                         (md.heading/3 name)
                                         (case ?doc
                                           (#.Some doc)
                                           (md.paragraph (md.text doc))

                                           _
                                           md.empty)
                                         type_code)))))
                          types)]
    (in (list\fold (function.flip md.then)
                   (md.heading/2 "Types")
                   type_docs))))

(def: (document_macros module_name names)
  (-> Text (List [Text Code]) (Markdown Block))
  (|> names
      (list\map (: (-> [Text Code] (Markdown Block))
                   (function (_ [name def_annotations])
                     ($_ md.then
                         (md.heading/3 name)
                         (<| (: (Markdown Block))
                             (maybe.default md.empty)
                             (do maybe.monad
                               [documentation (annotation.documentation def_annotations)]
                               (in (md.code documentation))))))))
      (list\fold (function.flip md.then)
                 (md.heading/2 "Macros"))))

(template [<singular> <plural> <header>]
  [(def: (<singular> module type)
     (-> Text Type (Markdown Block))
     (md.code (pprint_type (dec 0) "?" module type)))

   (def: (<plural> module values)
     (-> Text (List Value) (Markdown Block))
     (|> values
         (list\map (function (_ [name def_annotations value_type])
                     (let [?doc (annotation.documentation def_annotations)
                           usage (case (annotation.function_arguments def_annotations)
                                   #.Nil
                                   name

                                   args
                                   (format "(" (text.join_with " " (list& name args)) ")"))]
                       ($_ md.then
                           (md.heading/3 usage)
                           (case ?doc
                             (#.Some doc)
                             (md.code doc)

                             _
                             md.empty)
                           (<singular> module value_type)))))
         (list\fold (function.flip md.then)
                    (md.heading/2 <header>))))]

  [document_implementation document_implementations "Implementations"]
  [document_value     document_values     "Values"]
  )

(def: (enclose_lines pre+post block)
  (-> [Text Text] Text Text)
  (|> block
      (text.split_all_with text.new_line)
      (list\map (text.enclose pre+post))
      (text.join_with text.new_line)))

(def: (document_module [[module_name module] organization])
  (-> [[Text Module] Organization] (Meta [Text (Markdown Block)]))
  (do meta.monad
    [#let [(^slots [#types #macros #implementations #values]) organization
           annotations (|> module
                           (get@ #.module_annotations)
                           (maybe.default (' {}))
                           annotation.documentation)
           description (case annotations
                         (#.Some doc_text)
                         (md.quote (md.paragraph (md.text doc_text)))

                         #.None
                         md.empty)
           empty_block (: (Markdown Block) md.empty)]
     types_documentation (if (list.empty? types)
                           (in empty_block)
                           (document_types module_name types))
     #let [documentation ($_ md.then
                             types_documentation
                             (if (list.empty? macros) empty_block (document_macros module_name macros))
                             (if (list.empty? implementations) empty_block (document_implementations module_name implementations))
                             (if (list.empty? values) empty_block (document_values module_name values)))]]
    (in [module_name
         ($_ md.then
             (md.heading/1 module_name)
             description
             documentation)])))

(exception: #export (io_error {error Text})
  error)

(def: (save_documentation! [module_name documentation])
  (-> [Text (Markdown Block)] (IO Any))
  (let [path (format (text.replace_all "/" "_" module_name) ".md")]
    (do io.monad
      [outcome (\ file.default write (\ utf8.codec encode (md.markdown documentation)) path)]
      (in (case outcome
            (#try.Failure error)
            (debug.log! (exception.construct io_error error))

            (#try.Success _)
            [])))))

(macro: (gen_documentation! _)
  (do {! meta.monad}
    [all_modules meta.modules
     #let [lux_modules (|> all_modules
                           (list.only (function.compose lux_module? product.left))
                           (list.sort name_sort))]
     lux_exports (monad.map ! (function.compose meta.exports product.left)
                            lux_modules)
     module_documentation (|> (list\map organize_definitions lux_exports)
                              (list.zipped/2 lux_modules)
                              (monad.map ! document_module))
     #let [_ (io.run (monad.map io.monad save_documentation! module_documentation))]]
    (in (list))))

(gen_documentation!)

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