type expr =
  | Apply of string * Pattern.t Label.Map.t
  | And of expr * expr
  | Or of expr * expr
  | Not of expr * expr
  | Forall of expr * expr

type rule = {
  rule_name : string;
  head : Pattern.t Label.Map.t;
  body : expr;
  witness : Label.t option;
}

type primitive = D_integer | D_natural | D_rational | D_string

type domain =
  | Primitive of primitive
  | Criterion of { subject : string; body : expr }

type def =
  | Table of { schema : (Label.t * domain) list; facts : Term.t Label.Map.t list }
  | Rules of rule list

module type WORLD = sig
  type t

  val find : t -> string -> def option
end

type error =
  | Invalid_label of string
  | Duplicate_label of string
  | Zero_denominator
  | Lowercase_name of string
  | Term_expected of string
  | Proposition_expected of string
  | Malformed_directive of string
  | Unknown_domain of string
  | Unknown_relation of string
  | Already_defined of string
  | Fact_not_ground of string
  | Schema_mismatch of { relation : string; detail : string }
  | Domain_violation of { relation : string; label : string }
  | Non_total_rule of string
  | Mutual_recursion of string list
  | Criterion_arity of { relation : string; label : string }

type pending = {
  pending_relation : string;
  pending_label : Label.t;
  pending_body : expr;
}

let pp_error fmt = function
  | Invalid_label l -> Format.fprintf fmt "invalid label: %s" l
  | Duplicate_label l -> Format.fprintf fmt "duplicate label: %s:" l
  | Zero_denominator -> Format.fprintf fmt "rational literal with zero denominator"
  | Lowercase_name n ->
      Format.fprintf fmt "lowercase name %s cannot be a variable or relation here" n
  | Term_expected d -> Format.fprintf fmt "a term was expected: %s" d
  | Proposition_expected d -> Format.fprintf fmt "a proposition was expected: %s" d
  | Malformed_directive d -> Format.fprintf fmt "malformed directive: %s" d
  | Unknown_domain d -> Format.fprintf fmt "unknown domain: %s" d
  | Unknown_relation r -> Format.fprintf fmt "unknown relation: %s" r
  | Already_defined r -> Format.fprintf fmt "relation already defined: %s" r
  | Fact_not_ground r -> Format.fprintf fmt "fact for %s is not ground" r
  | Schema_mismatch { relation; detail } ->
      Format.fprintf fmt "schema mismatch on %s: %s" relation detail
  | Domain_violation { relation; label } ->
      Format.fprintf fmt "domain violation on %s at %s:" relation label
  | Non_total_rule name ->
      Format.fprintf fmt
        "rule %S rejected: no structurally decreasing argument in its recursive call"
        name
  | Mutual_recursion cycle ->
      Format.fprintf fmt "mutual recursion rejected: %s"
        (String.concat " -> " cycle)
  | Criterion_arity { relation; label } ->
      Format.fprintf fmt
        "criterion on %s at %s: must have exactly one free variable, the subject"
        relation label

let ( let* ) = Result.bind

let is_uppercase s = String.length s > 0 && s.[0] >= 'A' && s.[0] <= 'Z'

let label_of_string s =
  match Label.of_string s with
  | Ok l -> Ok l
  | Error (`Invalid_label l) -> Error (Invalid_label l)

let bindings_to_map ~of_node bindings =
  List.fold_left
    (fun acc (name, node) ->
      let* map = acc in
      let* label = label_of_string name in
      if Label.Map.mem label map then Error (Duplicate_label name)
      else
        let* value = of_node node in
        Ok (Label.Map.add label value map))
    (Ok Label.Map.empty) bindings

(* Term position: uppercase names are variables. Blocks and
   combinators are propositions and cannot appear here. *)
let rec pattern_of_node (node : Ast.node) : (Pattern.t, error) result =
  match node with
  | Ast.Name n ->
      if is_uppercase n then Ok (Pattern.Var n) else Error (Lowercase_name n)
  | Ast.Int z -> Ok (Pattern.Integer z)
  | Ast.Rat (num, den) -> (
      match Term.rational ~num ~den with
      | Ok t -> Ok (Pattern.of_term t)
      | Error `Zero_denominator -> Error Zero_denominator)
  | Ast.Str s -> Ok (Pattern.String s)
  | Ast.Record fields ->
      let* map = bindings_to_map ~of_node:pattern_of_node fields in
      Ok (Pattern.Record map)
  | Ast.Seq (items, tail) ->
      let* tail =
        match tail with
        | None -> Ok Pattern.Nil
        | Some node -> pattern_of_node node
      in
      List.fold_right
        (fun item acc ->
          let* acc = acc in
          let* head = pattern_of_node item in
          Ok (Pattern.cons ~head ~tail:acc))
        items (Ok tail)
  | Ast.Block (head, _) -> Error (Term_expected ("[" ^ head ^ " ...]"))
  | Ast.And _ -> Error (Term_expected ":and")
  | Ast.Or _ -> Error (Term_expected ":or")
  | Ast.Not _ -> Error (Term_expected ":not")

(* Renaming a quantified variable apart. The stamp cannot be lexed,
   so quantifier scopes are unforgeable from source text. *)
let quantifier_counter = ref 0

let rec rename_var ~from ~into = function
  | Pattern.Var v when String.equal v from -> Pattern.Var into
  | Pattern.Record map -> Pattern.Record (Label.Map.map (rename_var ~from ~into) map)
  | other -> other

let rec rename_var_expr ~from ~into = function
  | Apply (name, bindings) ->
      Apply (name, Label.Map.map (rename_var ~from ~into) bindings)
  | And (a, b) -> And (rename_var_expr ~from ~into a, rename_var_expr ~from ~into b)
  | Or (a, b) -> Or (rename_var_expr ~from ~into a, rename_var_expr ~from ~into b)
  | Not (a, b) -> Not (rename_var_expr ~from ~into a, rename_var_expr ~from ~into b)
  | Forall (a, b) ->
      Forall (rename_var_expr ~from ~into a, rename_var_expr ~from ~into b)

(* Proposition position: blocks apply relations, combinators
   compose. Terms cannot appear here. Exists and ForAll are the two
   higher-order forms; their variable is renamed apart, and Exists
   lowers to a conjunction because query variables are existential
   already. *)
let rec expr_of_node (node : Ast.node) : (expr, error) result =
  match node with
  | Ast.Block (("Exists" | "ForAll") as form, bindings) -> (
      match
        ( List.assoc_opt "variable" bindings,
          List.assoc_opt "quantifier" bindings,
          List.assoc_opt "body" bindings )
      with
      | Some (Ast.Name variable), Some quantifier, Some body
        when is_uppercase variable && List.length bindings = 3 ->
          let* quantifier = expr_of_node quantifier in
          let* body = expr_of_node body in
          incr quantifier_counter;
          let fresh = variable ^ "@q" ^ string_of_int !quantifier_counter in
          let quantifier = rename_var_expr ~from:variable ~into:fresh quantifier
          and body = rename_var_expr ~from:variable ~into:fresh body in
          if String.equal form "Exists" then Ok (And (quantifier, body))
          else Ok (Forall (quantifier, body))
      | _ ->
          Error
            (Malformed_directive
               (form ^ " variable: Name quantifier: [...] body: [...]")))
  | Ast.Block (head, bindings) ->
      let* map = bindings_to_map ~of_node:pattern_of_node bindings in
      Ok (Apply (head, map))
  | Ast.And (a, b) ->
      let* a = expr_of_node a in
      let* b = expr_of_node b in
      Ok (And (a, b))
  | Ast.Or (a, b) ->
      let* a = expr_of_node a in
      let* b = expr_of_node b in
      Ok (Or (a, b))
  | Ast.Not (a, b) ->
      let* a = expr_of_node a in
      let* b = expr_of_node b in
      Ok (Not (a, b))
  | Ast.Name n -> Error (Proposition_expected n)
  | Ast.Int _ | Ast.Rat _ | Ast.Str _ | Ast.Record _ | Ast.Seq _ ->
      Error (Proposition_expected "a literal")

let rec expr_free_vars = function
  | Apply (_, bindings) ->
      Label.Map.fold (fun _ p acc -> acc @ Pattern.vars p) bindings []
  | And (a, b) | Or (a, b) | Not (a, b) | Forall (a, b) ->
      expr_free_vars a @ expr_free_vars b

let primitive_of_name = function
  | "Integer" -> Some D_integer
  | "Natural" -> Some D_natural
  | "Rational" -> Some D_rational
  | "String" -> Some D_string
  | _ -> None

(* Membership of a ground term in a primitive domain. Note the
   inclusions of decision 12: a Natural is an Integer, an Integer is
   a Rational, all literally. *)
let member_of_primitive (term : Term.t) = function
  | D_integer -> ( match term with Term.Integer _ -> true | _ -> false)
  | D_natural -> (
      match term with Term.Integer z -> Z.geq z Z.zero | _ -> false)
  | D_rational -> (
      match term with Term.Integer _ | Term.Rational _ -> true | _ -> false)
  | D_string -> ( match term with Term.String _ -> true | _ -> false)

let rec applications = function
  | Apply (name, bindings) -> [ (name, bindings) ]
  | And (a, b) | Or (a, b) | Not (a, b) | Forall (a, b) ->
      applications a @ applications b

module String_map = Map.Make (String)

module Standalone = struct
  type t = def String_map.t

  let find world name = String_map.find_opt name world

  let schema_of_node ~relation = function
    | Ast.Seq (entries, None) ->
        List.fold_left
          (fun acc entry ->
            let* schema = acc in
            match entry with
            | Ast.Block (attr, [ ("domain", domain_node) ]) -> (
                let* label = label_of_string attr in
                match domain_node with
                | Ast.Name name when Option.is_some (primitive_of_name name) ->
                    let primitive = Option.get (primitive_of_name name) in
                    Ok ((label, Primitive primitive) :: schema)
                | Ast.Name other -> Error (Unknown_domain other)
                | criterion ->
                    (* A membership criterion: any relational
                       expression with exactly one free variable,
                       the subject. *)
                    let* body = expr_of_node criterion in
                    let subjects =
                      List.sort_uniq String.compare (expr_free_vars body)
                    in
                    (match subjects with
                    | [ subject ] ->
                        Ok ((label, Criterion { subject; body }) :: schema)
                    | _ ->
                        Error
                          (Criterion_arity
                             { relation; label = Label.to_string label })))
            | _ ->
                Error
                  (Malformed_directive "schema entries are [attr domain: ...]"))
          (Ok []) entries
        |> Result.map List.rev
    | _ -> Error (Malformed_directive "schema: takes #( ... ) of entries")

  (* The structural decrease check of decision 10, rung one. Every
     self-recursive application must have an argument strictly
     contained in the head's argument at the same label; that label
     is kept as the runtime witness. *)
  let admit_rule ~relation rule_name head body =
    let recursive =
      applications body |> List.filter (fun (name, _) -> String.equal name relation)
    in
    match recursive with
    | [] -> Ok { rule_name; head; body; witness = None }
    | _ ->
        let witness_for (_, call_bindings) =
          Label.Map.fold
            (fun label head_arg found ->
              match found with
              | Some _ -> found
              | None -> (
                  match Label.Map.find_opt label call_bindings with
                  | Some call_arg
                    when Pattern.strictly_contains head_arg ~sub:call_arg ->
                      Some label
                  | _ -> None))
            head None
        in
        let witnesses = List.map witness_for recursive in
        if List.for_all Option.is_some witnesses then
          (* All recursive calls decrease; any common witness works
             at runtime, and per-call witnesses agree by label here
             because they are all drawn from the same head. *)
          Ok { rule_name; head; body; witness = List.nth witnesses 0 }
        else Error (Non_total_rule rule_name)

  let load statements =
    (* First pass: declare tables, collect raw material. *)
    let* world, raw_facts, raw_rules, raw_queries =
      List.fold_left
        (fun acc statement ->
          let* world, facts, rules, queries = acc in
          match (statement : Ast.statement) with
          | Ast.Directive ("Relation", [ ("named", Ast.Name name); ("schema", schema) ])
            ->
              if String_map.mem name world then Error (Already_defined name)
              else
                let* schema = schema_of_node ~relation:name schema in
                Ok
                  ( String_map.add name (Table { schema; facts = [] }) world,
                    facts,
                    rules,
                    queries )
          | Ast.Directive ("Relation", _) ->
              Error (Malformed_directive "Relation named: Name schema: #( ... )")
          | Ast.Directive
              ("Rule", [ ("name", Ast.Str rule_name); ("head", head); ("body", body) ])
            ->
              Ok (world, facts, (rule_name, head, body) :: rules, queries)
          | Ast.Directive ("Rule", _) ->
              Error (Malformed_directive "Rule name: \"...\" head: [...] body: ...")
          | Ast.Directive ("Query", [ ("expression", body) ]) ->
              Ok (world, facts, rules, body :: queries)
          | Ast.Directive ("Query", _) ->
              Error (Malformed_directive "Query expression: ...")
          | Ast.Directive (name, _) -> Error (Malformed_directive name)
          | Ast.Assert node -> Ok (world, node :: facts, rules, queries))
        (Ok (String_map.empty, [], [], []))
        statements
    in
    let raw_facts = List.rev raw_facts
    and raw_rules = List.rev raw_rules
    and raw_queries = List.rev raw_queries in
    (* Second pass: compile rules, group them by head relation, admit
       each with the totality check. *)
    let* compiled_rules =
      List.fold_left
        (fun acc (rule_name, head_node, body_node) ->
          let* grouped = acc in
          let* head_expr = expr_of_node head_node in
          let* relation, head =
            match head_expr with
            | Apply (name, bindings) -> Ok (name, bindings)
            | _ -> Error (Malformed_directive "a rule head is a single [...] block")
          in
          let* body = expr_of_node body_node in
          let* rule = admit_rule ~relation rule_name head body in
          let existing =
            Option.value ~default:[] (String_map.find_opt relation grouped)
          in
          Ok (String_map.add relation (existing @ [ rule ]) grouped))
        (Ok String_map.empty) raw_rules
    in
    let* world =
      String_map.fold
        (fun relation rules acc ->
          let* world = acc in
          if String_map.mem relation world then Error (Already_defined relation)
          else Ok (String_map.add relation (Rules rules) world))
        compiled_rules (Ok world)
    in
    (* Third pass: admit facts against schemas and domains. This is
       integrity control at loading, the same check membership always
       is. Primitive domains decide here; criteria need the solver,
       so they are handed to Program as pending checks with the
       subject already bound to the attribute value. *)
    let equal_l = Result.get_ok (label_of_string "left")
    and right_l = Result.get_ok (label_of_string "right") in
    let bind_subject subject value body =
      And
        ( Apply
            ( "Equal",
              Label.Map.(
                empty
                |> add equal_l (Pattern.Var subject)
                |> add right_l (Pattern.of_term value)) ),
          body )
    in
    let* world, pendings =
      List.fold_left
        (fun acc node ->
          let* world, pendings = acc in
          let* expr = expr_of_node node in
          let* relation, bindings =
            match expr with
            | Apply (name, bindings) -> Ok (name, bindings)
            | _ -> Error (Malformed_directive "a fact is a single [...] block")
          in
          match String_map.find_opt relation world with
          | Some (Table { schema; facts }) ->
              let* tuple =
                Label.Map.fold
                  (fun label pattern acc ->
                    let* tuple = acc in
                    match Pattern.to_term pattern with
                    | Some term -> Ok (Label.Map.add label term tuple)
                    | None -> Error (Fact_not_ground relation))
                  bindings (Ok Label.Map.empty)
              in
              let* () =
                if
                  List.length schema = Label.Map.cardinal tuple
                  && List.for_all
                       (fun (label, _) -> Label.Map.mem label tuple)
                       schema
                then Ok ()
                else
                  Error
                    (Schema_mismatch
                       { relation; detail = "fact labels must match the schema" })
              in
              let* new_pendings =
                List.fold_left
                  (fun acc (label, domain) ->
                    let* pendings = acc in
                    let value = Label.Map.find label tuple in
                    match domain with
                    | Primitive primitive ->
                        if member_of_primitive value primitive then Ok pendings
                        else
                          Error
                            (Domain_violation
                               { relation; label = Label.to_string label })
                    | Criterion { subject; body } ->
                        Ok
                          ({
                             pending_relation = relation;
                             pending_label = label;
                             pending_body = bind_subject subject value body;
                           }
                          :: pendings))
                  (Ok []) schema
              in
              Ok
                ( String_map.add relation
                    (Table { schema; facts = facts @ [ tuple ] })
                    world,
                  pendings @ List.rev new_pendings )
          | Some (Rules _) ->
              Error
                (Schema_mismatch
                   { relation; detail = "facts cannot be asserted on a rule-defined relation" })
          | None -> Error (Unknown_relation relation))
        (Ok (world, []))
        raw_facts
    in
    (* Fourth pass: reference checking and the mutual recursion cut.
       Self recursion is admitted with a witness; longer cycles are
       rejected outright. *)
    let known name =
      Builtin.is_builtin name || String_map.mem name world
    in
    let check_refs body =
      List.fold_left
        (fun acc (name, _) ->
          let* () = acc in
          if known name then Ok () else Error (Unknown_relation name))
        (Ok ()) (applications body)
    in
    let* () =
      String_map.fold
        (fun _ def acc ->
          let* () = acc in
          match def with
          | Table { schema; _ } ->
              List.fold_left
                (fun acc (_, domain) ->
                  let* () = acc in
                  match domain with
                  | Primitive _ -> Ok ()
                  | Criterion { body; _ } -> check_refs body)
                (Ok ()) schema
          | Rules rules ->
              List.fold_left
                (fun acc rule ->
                  let* () = acc in
                  check_refs rule.body)
                (Ok ()) rules)
        world (Ok ())
    in
    let dependencies relation =
      match String_map.find_opt relation world with
      | Some (Rules rules) ->
          List.concat_map
            (fun rule ->
              applications rule.body
              |> List.filter_map (fun (name, _) ->
                     if
                       (not (String.equal name relation))
                       && (match String_map.find_opt name world with
                          | Some (Rules _) -> true
                          | _ -> false)
                     then Some name
                     else None))
            rules
      | _ -> []
    in
    let* () =
      let rec visit path relation =
        if List.mem relation path then
          Error (Mutual_recursion (List.rev (relation :: path)))
        else
          List.fold_left
            (fun acc dep ->
              let* () = acc in
              visit (relation :: path) dep)
            (Ok ())
            (dependencies relation)
      in
      String_map.fold
        (fun relation def acc ->
          let* () = acc in
          match def with Rules _ -> visit [] relation | Table _ -> Ok ())
        world (Ok ())
    in
    (* Queries compile last, against the full catalog. *)
    let* queries =
      List.fold_left
        (fun acc node ->
          let* queries = acc in
          let* expr = expr_of_node node in
          let* () = check_refs expr in
          Ok (expr :: queries))
        (Ok []) raw_queries
      |> Result.map List.rev
    in
    Ok (world, queries, pendings)
end