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 list;
}

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

type transition_result = T_state of Term.t | T_fail | T_stuck

module type WORLD = sig
  type t

  val find : t -> ?at:Term.t -> string -> def option

  val transition :
    t -> string -> Pattern.t Label.Map.t -> transition_result 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_component 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_component relations ->
      Format.fprintf fmt
        "recursion among %s rejected: no lexicographic structural decrease covers every call"
        (String.concat ", " relations)
  | 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.Ref name -> Ok (Pattern.Reference name)
  | 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.Ref _ | 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 ?at:_ name = String_map.find_opt name world
  let transition _ _ _ = None

  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,
     generalized to call-graph components. For every strongly
     connected component of the call graph among rule-defined
     relations, a lexicographic measure of shared labels must
     decrease on every intra-component call: at some position the
     call argument sits strictly inside the head's, and at every
     earlier position they are syntactically equal. The measure is
     stored on each recursive rule and its positions must be ground
     at unfolding, which turns the proof into a runtime termination
     measure. *)

  type call = { head : Pattern.t Label.Map.t; args : Pattern.t Label.Map.t }

  let decreases c label =
    match (Label.Map.find_opt label c.head, Label.Map.find_opt label c.args) with
    | Some head_arg, Some call_arg ->
        Pattern.strictly_contains head_arg ~sub:call_arg
    | _ -> false

  let stays_equal c label =
    match (Label.Map.find_opt label c.head, Label.Map.find_opt label c.args) with
    | Some head_arg, Some call_arg -> Pattern.equal head_arg call_arg
    | _ -> false

  (* Greedy lexicographic search with backtracking: pick a label
     that decreases or stays equal on every remaining call and
     decreases on at least one; the calls it covers drop out, the
     rest must be covered by later positions. *)
  let rec lex_measure calls candidates =
    match calls with
    | [] -> Some []
    | _ ->
        List.find_map
          (fun label ->
            if
              List.for_all
                (fun c -> decreases c label || stays_equal c label)
                calls
              && List.exists (fun c -> decreases c label) calls
            then
              let remaining =
                List.filter (fun c -> not (decreases c label)) calls
              in
              Option.map
                (fun rest -> label :: rest)
                (lex_measure remaining
                   (List.filter (fun l -> not (Label.equal l label)) candidates))
            else None)
          candidates

  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 = { rule_name; head; body; witness = [] } 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
    (* Component analysis: group rule-defined relations into
       strongly connected components of the call graph, find a
       lexicographic measure for each recursive component, and stamp
       it on the recursive rules as their runtime witness. *)
    let rule_defined name =
      match String_map.find_opt name world with
      | Some (Rules _) -> true
      | _ -> false
    in
    let edges relation =
      match String_map.find_opt relation world with
      | Some (Rules rules) ->
          List.concat_map
            (fun rule -> List.map fst (applications rule.body))
            rules
          |> List.filter rule_defined
          |> List.sort_uniq String.compare
      | _ -> []
    in
    let reachable start =
      let rec go seen frontier =
        match frontier with
        | [] -> seen
        | name :: rest ->
            if List.mem name seen then go seen rest
            else go (name :: seen) (edges name @ rest)
      in
      go [] (edges start)
    in
    let* world =
      let relations =
        String_map.fold
          (fun name def acc ->
            match def with Rules _ -> name :: acc | Table _ -> acc)
          world []
      in
      let component relation =
        let reach = reachable relation in
        relation
        :: List.filter
             (fun other ->
               (not (String.equal other relation))
               && List.mem other reach
               && List.mem relation (reachable other))
             relations
      in
      List.fold_left
        (fun acc relation ->
          let* world = acc in
          let members = component relation in
          (* Process each component once, from its least member. *)
          if not (String.equal relation (List.fold_left min relation members))
          then Ok world
          else
            let member name = List.mem name members in
            let calls_of (rule : rule) =
              applications rule.body
              |> List.filter_map (fun (target, args) ->
                     if member target then Some { head = rule.head; args }
                     else None)
            in
            let rules_of name =
              match String_map.find_opt name world with
              | Some (Rules rules) -> rules
              | _ -> []
            in
            let calls = List.concat_map calls_of (List.concat_map rules_of members) in
            if calls = [] then Ok world
            else
              let candidates =
                match calls with
                | [] -> []
                | first :: rest ->
                    Label.Map.fold
                      (fun label _ acc ->
                        if
                          List.for_all
                            (fun c ->
                              Label.Map.mem label c.head
                              && Label.Map.mem label c.args)
                            rest
                          && Label.Map.mem label first.args
                        then label :: acc
                        else acc)
                      first.head []
              in
              match lex_measure calls candidates with
              | None -> Error (Non_total_component members)
              | Some measure
                when
                  (* The search is an untrusted heuristic; the
                     measure it found is validated call by call with
                     the lexicographic checker extracted from Rocq,
                     whose soundness and well-foundedness are
                     proven. *)
                  not
                    (List.for_all
                       (fun c ->
                         let project bindings =
                           List.map
                             (fun label ->
                               Pattern.to_checker (Label.Map.find label bindings))
                             measure
                         in
                         Decrease_extracted.lex_ltb Label.equal
                           Pattern.checker_leaf_equal (project c.args)
                           (project c.head))
                       calls) ->
                  Error (Non_total_component members)
              | Some measure ->
                  Ok
                    (List.fold_left
                       (fun world name ->
                         let stamped =
                           List.map
                             (fun rule ->
                               if calls_of rule = [] then rule
                               else { rule with witness = measure })
                             (rules_of name)
                         in
                         String_map.add name (Rules stamped) world)
                       world members))
        (Ok world) relations
    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