(.require
 [library
  [lux (.except with)
   [abstract
    ["[0]" monad (.only do)]
    ["[0]" equivalence]]
   [control
    ["[0]" maybe]
    ["[0]" try]
    ["<>" parser (.only)
     ["<[0]>" code (.only Parser)]]]
   [data
    ["[0]" product]
    ["[0]" text (.use "[1]#[0]" equivalence)
     ["%" \\format (.only format)]]
    [collection
     ["[0]" list (.use "[1]#[0]" monad mix)]
     ["[0]" dictionary (.only Dictionary)]]]
   ["[0]" macro (.only)
    ["[0]" code]
    [syntax (.only syntax)]]
   [math
    ["[0]" number (.only)
     ["n" nat]]]
   ["[0]" meta]
   ["[0]" type (.use "[1]#[0]" equivalence)
    ["[0]" check (.only Check)]]]])

(def (type_var id env)
  (-> Nat Type_Context (Meta Type))
  (case (list.example (|>> product.left (n.= id))
                      (the .#var_bindings env))
    {.#Some [_ {.#Some type}]}
    (case type
      {.#Var id'}
      (type_var id' env)

      _
      (at meta.monad in type))

    {.#Some [_ {.#None}]}
    (meta.failure (format "Unbound type-var " (%.nat id)))

    {.#None}
    (meta.failure (format "Unknown type-var " (%.nat id)))
    ))

(def (implicit_type var_name)
  (-> Symbol (Meta Type))
  (do meta.monad
    [raw_type (meta.type var_name)
     compiler meta.compiler_state]
    (case raw_type
      {.#Var id}
      (type_var id (the .#type_context compiler))

      _
      (in raw_type))))

(def (member_type idx sig_type)
  (-> Nat Type (Check Type))
  (case sig_type
    {.#Named _ sig_type'}
    (member_type idx sig_type')

    {.#Apply arg func}
    (case (type.applied (list arg) func)
      {.#None}
      (check.failure (format "Cannot apply type " (%.type func) " to type " (%.type arg)))

      {.#Some sig_type'}
      (member_type idx sig_type'))

    {.#Product left right}
    (if (n.= 0 idx)
      (at check.monad in left)
      (member_type (-- idx) right))

    _
    (if (n.= 0 idx)
      (at check.monad in sig_type)
      (check.failure (format "Cannot find member type " (%.nat idx) " for " (%.type sig_type))))))

(def (member_name member)
  (-> Symbol (Meta Symbol))
  (case member
    ["" simple_name]
    (meta.either (do meta.monad
                   [member (meta.normal member)
                    _ (meta.slot member)]
                   (in member))
                 (do [! meta.monad]
                   [this_module_name meta.current_module_name
                    imp_mods (meta.imported_modules this_module_name)
                    tag_lists (monad.each ! meta.tag_lists imp_mods)
                    .let [tag_lists (|> tag_lists list#conjoint (list#each product.left) list#conjoint)
                          candidates (list.only (|>> product.right (text#= simple_name))
                                                tag_lists)]]
                   (case candidates
                     {.#End}
                     (meta.failure (format "Unknown tag: " (%.symbol member)))
                     
                     {.#Item winner {.#End}}
                     (in winner)
                     
                     _
                     (meta.failure (format "Too many candidate tags: " (%.list %.symbol candidates))))))

    _
    (at meta.monad in member)))

(def (implicit_member member)
  (-> Symbol (Meta [Nat Type]))
  (do meta.monad
    [member (member_name member)
     [idx tag_list sig_type] (meta.slot member)]
    (in [idx sig_type])))

(def (compatible_type? interface candidate)
  (-> Type Type Bit)
  (with_expansions [<found?> (type#= interface candidate)]
    (<| (or <found?>)

        (let [[parameters candidate] (type.flat_univ_q candidate)])
        (or <found?>)

        (let [[inputs candidate] (type.flat_function candidate)])
        (or <found?>)

        (let [[candidate parameters] (type.flat_application candidate)])
        (or <found?>)
        
        (let [candidate (type.de_aliased candidate)])
        <found?>)))

(def (available_definitions sig_type source_module target_module constants aggregate)
  (-> Type Text Text (List [Text Definition]) (-> (List [Symbol Type]) (List [Symbol Type])))
  (list#mix (function (_ [name [exported? def_type def_value]] aggregate)
              (if (and (or (text#= target_module source_module)
                           exported?)
                       (compatible_type? sig_type def_type))
                {.#Item [[source_module name] def_type] aggregate}
                aggregate))
            aggregate
            constants))

(def (local_env sig_type)
  (-> Type (Meta (List [Symbol Type])))
  (do meta.monad
    [local_batches meta.locals
     .let [total_locals (list#mix (function (_ [name type] table)
                                    (try.else table (dictionary.has' name type table)))
                                  (is (Dictionary Text Type)
                                      (dictionary.empty text.hash))
                                  (list#conjoint local_batches))]]
    (in (|> total_locals
            dictionary.entries
            (list.all (function (_ [name type])
                        (if (compatible_type? sig_type type)
                          {.#Some [["" name] type]}
                          {.#None})))))))

(def (local_structs sig_type)
  (-> Type (Meta (List [Symbol Type])))
  (do [! meta.monad]
    [this_module_name meta.current_module_name
     definitions (meta.definitions this_module_name)]
    (in (available_definitions sig_type this_module_name this_module_name definitions {.#End}))))

(def (imported_structs sig_type)
  (-> Type (Meta (List [Symbol Type])))
  (do [! meta.monad]
    [this_module_name meta.current_module_name
     imported_modules (meta.imported_modules this_module_name)
     accessible_definitions (monad.each ! meta.definitions imported_modules)]
    (in (list#mix (function (_ [imported_module definitions] tail)
                    (available_definitions sig_type imported_module this_module_name definitions tail))
                  {.#End}
                  (list.zipped_2 imported_modules accessible_definitions)))))

(def (on_argument arg func)
  (-> Type Type (Check Type))
  (case func
    {.#Named _ func'}
    (on_argument arg func')

    {.#UnivQ _}
    (do check.monad
      [[id var] check.var]
      (|> func
          (type.applied (list var))
          maybe.trusted
          (on_argument arg)))

    {.#Function input output}
    (do check.monad
      [_ (check.check input arg)]
      (in output))

    _
    (check.failure (format "Invalid function type: " (%.type func)))))

(def (concrete_type type)
  (-> Type (Check [(List Nat) Type]))
  (case type
    {.#UnivQ _}
    (do check.monad
      [[id var] check.var
       [ids final_output] (concrete_type (maybe.trusted (type.applied (list var) type)))]
      (in [{.#Item id ids}
           final_output]))
    
    _
    (at check.monad in [(list) type])))

(def (ensure_function_application! member_type input_types expected_output)
  (-> Type (List Type) Type (Check []))
  (do check.monad
    [actual_output (monad.mix check.monad ..on_argument member_type input_types)]
    (check.check expected_output actual_output)))

(type: Instance
  (Rec Instance
    (Record
     [#constructor Symbol
      #dependencies (List Instance)])))

(def (candidate_provision provision context dep alts)
  (-> (-> Lux Type_Context Type (Check Instance))
      Type_Context Type (List [Symbol Type])
      (Meta (List Instance)))
  (do meta.monad
    [compiler meta.compiler_state]
    (case (|> alts
              (list#each (function (_ [alt_name alt_type])
                           (case (check.result context
                                               (do [! check.monad]
                                                 [[tvars alt_type] (concrete_type alt_type)
                                                  .let [[deps alt_type] (type.flat_function alt_type)]
                                                  _ (check.check dep alt_type)
                                                  context' check.context
                                                  =deps (monad.each ! (provision compiler context') deps)]
                                                 (in =deps)))
                             {.#Left error}
                             (list)

                             {.#Right =deps}
                             (list [alt_name =deps]))))
              list#conjoint)
      {.#End}
      (meta.failure (format "No candidates for provisioning: " (%.type dep)))

      found
      (in found))))

(def (provision sig_type compiler context dep)
  (-> Type Lux Type_Context Type (Check Instance))
  (case (meta.result compiler
                     (all meta.either
                          (do meta.monad [alts (..local_env sig_type)] (..candidate_provision (provision sig_type) context dep alts))
                          (do meta.monad [alts (..local_structs sig_type)] (..candidate_provision (provision sig_type) context dep alts))
                          (do meta.monad [alts (..imported_structs sig_type)] (..candidate_provision (provision sig_type) context dep alts))))
    {.#Left error}
    (check.failure error)

    {.#Right candidates}
    (case candidates
      {.#End}
      (check.failure (format "No candidates for provisioning: " (%.type dep)))

      {.#Item winner {.#End}}
      (at check.monad in winner)

      _
      (check.failure (format "Too many candidates for provisioning: " (%.type dep) " --- " (%.list (|>> product.left %.symbol) candidates))))
    ))

(def (candidate_alternatives sig_type member_idx input_types output_type alts)
  (-> Type Nat (List Type) Type (List [Symbol Type]) (Meta (List Instance)))
  (do meta.monad
    [compiler meta.compiler_state
     context meta.type_context]
    (case (|> alts
              (list#each (function (_ [alt_name alt_type])
                           (case (<| (check.result context)
                                     (do [! check.monad]
                                       [[tvars alt_type] (concrete_type alt_type)
                                        .let [[deps alt_type] (type.flat_function alt_type)]
                                        _ (check.check alt_type sig_type)
                                        member_type (member_type member_idx alt_type)
                                        _ (ensure_function_application! member_type input_types output_type)
                                        context' check.context
                                        =deps (monad.each ! (provision sig_type compiler context') deps)]
                                       (in =deps)))
                             {.#Left error}
                             (list)

                             {.#Right =deps}
                             (list [alt_name =deps]))))
              list#conjoint)
      {.#End}
      (meta.failure (format "No alternatives for " (%.type (type.function input_types output_type))))

      found
      (in found))))

(def (alternatives sig_type member_idx input_types output_type)
  (-> Type Nat (List Type) Type (Meta (List Instance)))
  (let [test (candidate_alternatives sig_type member_idx input_types output_type)]
    (all meta.either
         (do meta.monad [alts (..local_env sig_type)] (test alts))
         (do meta.monad [alts (..local_structs sig_type)] (test alts))
         (do meta.monad [alts (..imported_structs sig_type)] (test alts)))))

(def (var? input)
  (-> Code Bit)
  (case input
    [_ {.#Symbol _}]
    #1

    _
    #0))

(def (pair_list [l r])
  (All (_ a) (-> [a a] (List a)))
  (list l r))

(def (instance$ [constructor dependencies])
  (-> Instance Code)
  (case dependencies
    {.#End}
    (code.symbol constructor)

    _
    (` ((~ (code.symbol constructor)) (~+ (list#each instance$ dependencies))))))

(def .public a/an
  (syntax (_ [member <code>.symbol
              args (<>.or (<>.and (<>.some <code>.symbol) <code>.end)
                          (<>.and (<>.some <code>.any) <code>.end))])
    (case args
      {.#Left [args _]}
      (do [! meta.monad]
        [[member_idx sig_type] (..implicit_member member)
         input_types (monad.each ! ..implicit_type args)
         output_type meta.expected_type
         chosen_ones (alternatives sig_type member_idx input_types output_type)]
        (case chosen_ones
          {.#End}
          (meta.failure (format "No implementation could be found for member: " (%.symbol member)))

          {.#Item chosen {.#End}}
          (in (list (` (.at (~ (instance$ chosen))
                            (~ (code.local (product.right member)))
                            (~+ (list#each code.symbol args))))))

          _
          (meta.failure (format "Too many implementations available: "
                                (|> chosen_ones
                                    (list#each (|>> product.left %.symbol))
                                    (text.interposed ", "))
                                " --- for type: " (%.type sig_type)))))

      {.#Right [args _]}
      (do [! meta.monad]
        [labels (|> (macro.symbol "g!parameter")
                    (list.repeated (list.size args))
                    (monad.all !))]
        (in (list (` (let [(~+ (|> args (list.zipped_2 labels) (list#each ..pair_list) list#conjoint))]
                       (..a/an (~ (code.symbol member)) (~+ labels)))))))
      )))

(def .public a ..a/an)
(def .public an ..a/an)

(def (implicit_bindings amount)
  (-> Nat (Meta (List Code)))
  (|> (macro.symbol "g!implicit")
      (list.repeated amount)
      (monad.all meta.monad)))

(def .public with
  (syntax (_ [implementations (<code>.tuple (<>.many <code>.any))
              body <code>.any])
    (do meta.monad
      [g!implicit+ (implicit_bindings (list.size implementations))]
      (in (list (` (let [(~+ (|> (list.zipped_2 g!implicit+ implementations)
                                 (list#each (function (_ [g!implicit implementation])
                                              (list g!implicit implementation)))
                                 list#conjoint))]
                     (~ body))))))))

(def .public implicitly
  (syntax (_ [implementations (<>.many <code>.any)])
    (do meta.monad
      [g!implicit+ (implicit_bindings (list.size implementations))]
      (in (|> (list.zipped_2 g!implicit+ implementations)
              (list#each (function (_ [g!implicit implementation])
                           (` (def .private (~ g!implicit)
                                (~ implementation))))))))))