type error = Non_total_plan of string list

let pp_error fmt (Non_total_plan goals) =
  Format.fprintf fmt
    "non-total plan rejected: no decidable goal among %s (bind more arguments or bound the search)"
    (String.concat ", " goals)

type residual = string * Pattern.t Label.Map.t

let pp_residual fmt (name, bindings) =
  Format.fprintf fmt "@[<hv 1>[%s" name;
  Label.Map.iter
    (fun label pattern ->
      Format.fprintf fmt "@ %a %a" Label.pp label Pattern.pp pattern)
    bindings;
  Format.fprintf fmt "]@]"

type solution = {
  bindings : (string * Pattern.t) list;
  residue : residual list;
}

module Make (W : World.WORLD) = struct
  exception Stuck of string list

  (* A goal is a proposition to establish or a denial to refute.
     Conjunction flattens; [u :not d] contributes the universe
     positively and the denial negatively. *)
  type goal = Positive of World.expr | Negated of World.expr

  (* What attempting a goal can say: run it into branches, or wait.
     A waiting goal is residuatable when it is pure membership that
     decides as soon as its arguments are ground; searches are
     never residuatable. *)
  type attempt =
    | Run of (Pattern.Subst.t * goal list) list
    | Wait of { residuatable : bool }

  let rec flatten = function
    | World.And (a, b) -> flatten a @ flatten b
    | World.Not (universe, denial) -> flatten universe @ [ Negated denial ]
    | expr -> [ Positive expr ]

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

  let goal_name = function
    | Positive (World.Apply (name, _)) -> name
    | Positive (World.Or _) -> ":or"
    | Positive (World.And _) -> ":and"
    | Positive (World.Not _) | Negated _ -> ":not"
    | Positive (World.Forall _) -> "ForAll"

  let unresolved_vars subst expr =
    expr_vars expr
    |> List.filter (fun v ->
           not (Pattern.is_ground (Pattern.Subst.resolve subst (Pattern.Var v))))

  (* Rule variables are renamed apart at every unfolding. The suffix
     character cannot be lexed, so renamed variables are unforgeable
     from source text. *)
  let freshen =
    let counter = ref 0 in
    fun rule ->
      incr counter;
      let stamp = "@" ^ string_of_int !counter in
      let rec rename = function
        | Pattern.Var v -> Pattern.Var (v ^ stamp)
        | Pattern.Record map -> Pattern.Record (Label.Map.map rename map)
        | ground -> ground
      in
      let rec rename_expr = function
        | World.Apply (name, bindings) ->
            World.Apply (name, Label.Map.map rename bindings)
        | World.And (a, b) -> World.And (rename_expr a, rename_expr b)
        | World.Or (a, b) -> World.Or (rename_expr a, rename_expr b)
        | World.Not (a, b) -> World.Not (rename_expr a, rename_expr b)
        | World.Forall (a, b) -> World.Forall (rename_expr a, rename_expr b)
      in
      (Label.Map.map rename rule.World.head, rename_expr rule.World.body)

  (* Unify a goal's bindings against a rule head or a stored tuple.
     Goal labels absent from the head fail the branch; head labels
     absent from the goal are existential. *)
  let unify_bindings subst goal_bindings target_bindings =
    Label.Map.fold
      (fun label pattern acc ->
        Option.bind acc (fun subst ->
            match Label.Map.find_opt label target_bindings with
            | Some target -> Pattern.Subst.unify subst pattern target
            | None -> None))
      goal_bindings (Some subst)

  let proven solutions = List.exists (fun (_, residue) -> residue = []) solutions

  (* Interval propagation over waiting membership constraints. What
     residuation returns inert, this looks at together: bounds with
     an empty intersection fail the branch outright, and an integer
     variable squeezed into a finite range converts to bounded
     generation, the same admissible enumeration Between performs
     explicitly. Sakura's cardinality analysis, inside the engine. *)
  module Propagate = struct
    let forced_label name =
      match Label.of_string name with Ok l -> l | Error _ -> assert false

    let value_l = forced_label "value"
    let left_l = forced_label "left"
    let right_l = forced_label "right"
    let result_l = forced_label "result"

    type bound = { value : Q.t; strict : bool }

    type info = {
      integer : bool;
      stringy : bool;
      lower : bound option;
      upper : bound option;
    }

    let no_info = { integer = false; stringy = false; lower = None; upper = None }

    module M = Map.Make (String)

    let note var f store =
      M.update var
        (fun info -> Some (f (Option.value ~default:no_info info)))
        store

    let stronger_lower a b =
      if Q.gt a.value b.value then a
      else if Q.lt a.value b.value then b
      else { a with strict = a.strict || b.strict }

    let stronger_upper a b =
      if Q.lt a.value b.value then a
      else if Q.gt a.value b.value then b
      else { a with strict = a.strict || b.strict }

    let add_lower bound info =
      {
        info with
        lower =
          Some
            (match info.lower with
            | None -> bound
            | Some existing -> stronger_lower bound existing);
      }

    let add_upper bound info =
      {
        info with
        upper =
          Some
            (match info.upper with
            | None -> bound
            | Some existing -> stronger_upper bound existing);
      }

    let numeric_value = function
      | Pattern.Integer z -> Some (Q.of_bigint z)
      | Pattern.Rational q -> Some q
      | _ -> None

    (* Read one waiting builtin goal into the store. Bindings arrive
       already resolved, so a variable here is genuinely unbound. *)
    let gather store (name, bindings) =
      let value = Label.Map.find_opt value_l bindings
      and left = Label.Map.find_opt left_l bindings
      and right = Label.Map.find_opt right_l bindings in
      match (name, value, left, right) with
      | ("Integer" | "Natural"), Some (Pattern.Var x), _, _ ->
          let store = note x (fun i -> { i with integer = true }) store in
          if String.equal name "Natural" then
            note x (add_lower { value = Q.zero; strict = false }) store
          else store
      | "String", Some (Pattern.Var x), _, _ ->
          note x (fun i -> { i with stringy = true }) store
      | ( ("LessThan" | "LessThanEqual" | "GreaterThan" | "GreaterEqual"),
          _,
          Some left,
          Some right ) -> (
          let strict =
            String.equal name "LessThan" || String.equal name "GreaterThan"
          in
          let less = String.equal name "LessThan" || String.equal name "LessThanEqual" in
          match (left, right) with
          | Pattern.Var x, ground when Option.is_some (numeric_value ground) ->
              let bound = { value = Option.get (numeric_value ground); strict } in
              note x (if less then add_upper bound else add_lower bound) store
          | ground, Pattern.Var x when Option.is_some (numeric_value ground) ->
              let bound = { value = Option.get (numeric_value ground); strict } in
              note x (if less then add_lower bound else add_upper bound) store
          | _ -> store)
      | _ -> store

    (* Integer endpoints of a rational interval. *)
    let lower_int { value; strict } =
      let n = Q.num value and d = Q.den value in
      if strict then Z.add (Z.fdiv n d) Z.one else Z.cdiv n d

    let upper_int { value; strict } =
      let n = Q.num value and d = Q.den value in
      if strict then Z.sub (Z.cdiv n d) Z.one else Z.fdiv n d

    type verdict = Fail | Enumerate of string * Z.t * Z.t | Unknown

    let analyze constraints =
      let store = List.fold_left gather M.empty constraints in
      let contradictory _ info =
        (info.stringy && (info.integer || info.lower <> None || info.upper <> None))
        ||
        match (info.lower, info.upper) with
        | Some lo, Some hi ->
            Q.gt lo.value hi.value
            || (Q.equal lo.value hi.value && (lo.strict || hi.strict))
            || (info.integer && Z.gt (lower_int lo) (upper_int hi))
        | _ -> false
      in
      if M.exists contradictory store then Fail
      else
        match
          M.choose_opt
            (M.filter
               (fun _ info ->
                 info.integer && info.lower <> None && info.upper <> None)
               store)
        with
        | Some (var, info) ->
            Enumerate
              ( var,
                lower_int (Option.get info.lower),
                upper_int (Option.get info.upper) )
        | None -> Unknown
  end

  (* Symbolic discharge (decision 10, rung three, first fragment):
     universals whose quantifier is an infinite domain are decided
     algebraically when quantifier and body both read as kind and
     interval constraints over the one quantified variable. Holds
     means every member is covered; Fails is always backed by a
     concrete counterexample; anything outside the fragment stays
     deferred. Subtype inclusions (a Natural is an Integer is a
     Rational) are decided here, literally, per decision 12. *)
  module Symbolic = struct
    type side = {
      is_int : bool;  (** members are integers *)
      is_str : bool;
      has_domain : bool;  (** some domain goal declared the kind *)
      needs_num : bool;  (** the body demands a numeric member *)
      lower : Propagate.bound option;
      upper : Propagate.bound option;
    }

    let free =
      {
        is_int = false;
        is_str = false;
        has_domain = false;
        needs_num = false;
        lower = None;
        upper = None;
      }

    (* One goal read as a constraint on [var]; None = outside the
       fragment. Goals without variables evaluate outright. *)
    let read subst var side (name, bindings) =
      let value = Label.Map.find_opt Propagate.value_l bindings
      and left = Label.Map.find_opt Propagate.left_l bindings
      and right = Label.Map.find_opt Propagate.right_l bindings in
      let ground_goal () =
        match Builtin.run subst name bindings with
        | Some (Builtin.Solved []) -> Some `Contradiction
        | Some (Builtin.Solved _) -> Some `Trivial
        | _ -> None
      in
      let is_var = function
        | Pattern.Var v -> String.equal v var
        | _ -> false
      in
      match (name, value, left, right) with
      | ("Integer" | "Natural"), Some v, _, _ when is_var v ->
          let side = { side with is_int = true; has_domain = true } in
          if String.equal name "Natural" then
            Some
              (`Side
                {
                  side with
                  lower =
                    Some
                      (match side.lower with
                      | None -> { Propagate.value = Q.zero; strict = false }
                      | Some b ->
                          Propagate.stronger_lower b
                            { Propagate.value = Q.zero; strict = false });
                })
          else Some (`Side side)
      | "Rational", Some v, _, _ when is_var v ->
          Some (`Side { side with has_domain = true })
      | "String", Some v, _, _ when is_var v ->
          Some (`Side { side with is_str = true; has_domain = true })
      | ( ("LessThan" | "LessThanEqual" | "GreaterThan" | "GreaterEqual"),
          _,
          Some l,
          Some r ) -> (
          let strict =
            String.equal name "LessThan" || String.equal name "GreaterThan"
          in
          let less =
            String.equal name "LessThan" || String.equal name "LessThanEqual"
          in
          let with_bound upper bound =
            if upper then
              `Side
                {
                  side with
                  upper =
                    Some
                      (match side.upper with
                      | None -> bound
                      | Some b -> Propagate.stronger_upper b bound);
                }
            else
              `Side
                {
                  side with
                  lower =
                    Some
                      (match side.lower with
                      | None -> bound
                      | Some b -> Propagate.stronger_lower b bound);
                }
          in
          match (l, r) with
          | v, g when is_var v && Option.is_some (Propagate.numeric_value g) ->
              Some
                (with_bound less
                   { Propagate.value = Option.get (Propagate.numeric_value g); strict })
          | g, v when is_var v && Option.is_some (Propagate.numeric_value g) ->
              Some
                (with_bound (not less)
                   { Propagate.value = Option.get (Propagate.numeric_value g); strict })
          | _ -> None)
      | ("Plus" | "Times"), _, Some l, Some r -> (
          (* The algebraic identities of the numeric relations:
             membership statements that hold for every numeric
             member, so a numeric quantifier discharges them and a
             string quantifier refutes them. Anything beyond the
             identities stays outside the fragment. *)
          let result = Label.Map.find_opt Propagate.result_l bindings in
          let is_zero = function
            | Pattern.Integer z -> Z.equal z Z.zero
            | _ -> false
          and is_one = function
            | Pattern.Integer z -> Z.equal z Z.one
            | _ -> false
          in
          let identity =
            match (name, l, r, result) with
            | "Plus", v, unit_, Some v' when is_var v && is_zero unit_ && is_var v' ->
                true
            | "Plus", unit_, v, Some v' when is_var v && is_zero unit_ && is_var v' ->
                true
            | "Times", v, unit_, Some v' when is_var v && is_one unit_ && is_var v' ->
                true
            | "Times", unit_, v, Some v' when is_var v && is_one unit_ && is_var v' ->
                true
            | "Times", v, zero, Some zero' when is_var v && is_zero zero && is_zero zero'
              ->
                true
            | "Times", zero, v, Some zero' when is_var v && is_zero zero && is_zero zero'
              ->
                true
            | _ -> false
          in
          if identity then Some (`Side { side with needs_num = true })
          else if Label.Map.for_all (fun _ p -> Pattern.vars p = []) bindings then
            ground_goal ()
          else None)
      | _, _, _, _ ->
          if
            Label.Map.for_all
              (fun _ p -> Pattern.vars p = [])
              bindings
          then ground_goal ()
          else None

    let read_side subst var goals =
      List.fold_left
        (fun acc goal ->
          match acc with
          | None | Some `Contradiction -> acc
          | Some (`Side side) -> (
              match read subst var side goal with
              | Some (`Side _ as s) -> Some s
              | Some `Trivial -> Some (`Side side)
              | Some `Contradiction -> Some `Contradiction
              | None -> None))
        (Some (`Side free)) goals

    (* Does the quantifier's lower bound guarantee the body's? For
       integer kinds compare tightened endpoints; missing body bound
       is no demand, missing quantifier bound cannot guarantee one. *)
    let lower_entails ~int q b =
      match b with
      | None -> true
      | Some b -> (
          match q with
          | None -> false
          | Some q ->
              if int then Z.geq (Propagate.lower_int q) (Propagate.lower_int b)
              else
                Q.gt q.Propagate.value b.Propagate.value
                || (Q.equal q.Propagate.value b.Propagate.value
                   && ((not b.strict) || q.strict)))

    let upper_entails ~int q b =
      match b with
      | None -> true
      | Some b -> (
          match q with
          | None -> false
          | Some q ->
              if int then Z.leq (Propagate.upper_int q) (Propagate.upper_int b)
              else
                Q.lt q.Propagate.value b.Propagate.value
                || (Q.equal q.Propagate.value b.Propagate.value
                   && ((not b.strict) || q.strict)))

    (* A rational interval that is one integer point. *)
    let single_integer_point q =
      match (q.lower, q.upper) with
      | Some lo, Some hi ->
          (not lo.Propagate.strict)
          && (not hi.Propagate.strict)
          && Q.equal lo.value hi.value
          && Z.equal (Q.den lo.value) Z.one
      | _ -> false

    let entails q b =
      if q.is_str then
        if b.is_int || b.needs_num || b.lower <> None || b.upper <> None then
          `Fails
        else `Holds
      else if b.is_str then `Fails
      else if b.is_int && (not q.is_int) && not (single_integer_point q) then
        `Fails
      else if
        lower_entails ~int:q.is_int q.lower b.lower
        && upper_entails ~int:q.is_int q.upper b.upper
      then `Holds
      else `Fails

    let discharge subst quantifier_goals body_goals =
      let apply_only goals =
        List.fold_left
          (fun acc goal ->
            Option.bind acc (fun goals ->
                match goal with
                | World.Apply (name, bindings) ->
                    Some
                      ((name, Label.Map.map (Pattern.Subst.resolve subst) bindings)
                      :: goals)
                | _ -> None))
          (Some []) goals
      in
      match (apply_only quantifier_goals, apply_only body_goals) with
      | Some quantifier, Some body -> (
          let vars =
            List.concat_map
              (fun (_, bindings) ->
                Label.Map.fold (fun _ p acc -> acc @ Pattern.vars p) bindings [])
              quantifier
            |> List.sort_uniq String.compare
          in
          match vars with
          | [ var ] -> (
              match (read_side subst var quantifier, read_side subst var body) with
              | Some (`Side q), Some (`Side b) ->
                  if q.has_domain then entails q b else `Unknown
              | Some (`Side q), Some `Contradiction ->
                  (* The body is falsified outright; any member of a
                     nonempty quantifier is a counterexample. *)
                  if q.has_domain then `Fails else `Unknown
              | Some `Contradiction, _ ->
                  (* An empty quantifier satisfies any universal. *)
                  `Holds
              | _ -> `Unknown)
          | _ -> `Unknown)
      | _ -> `Unknown
  end


  let rec try_goal world subst goal =
    match goal with
    | Negated denial ->
        if unresolved_vars subst denial = [] then
          match search world subst (flatten denial) with
          | [] -> Run [ (subst, []) ]
          | solutions ->
              (* Ground denials cannot residuate in practice, but if
                 one somehow does, refusing to conclude is the only
                 sound answer. *)
              if proven solutions then Run [] else Wait { residuatable = false }
        else Wait { residuatable = false }
    | Positive (World.Or (a, b)) ->
        Run [ (subst, flatten a); (subst, flatten b) ]
    | Positive ((World.And _ | World.Not _) as expr) -> Run [ (subst, flatten expr) ]
    | Positive (World.Forall (quantifier, body)) -> (
        (* Universal over the quantifier's solutions. Runs when the
           quantifier is decidable and closed; otherwise symbolic
           discharge gets a chance before the universal defers, so
           truths over infinite domains are proven, never
           enumerated. *)
        let rec conjuncts = function
          | World.And (a, b) -> conjuncts a @ conjuncts b
          | expr -> [ expr ]
        in
        let symbolically () =
          match
            Symbolic.discharge subst (conjuncts quantifier) (conjuncts body)
          with
          | `Holds -> Run [ (subst, []) ]
          | `Fails -> Run []
          | `Unknown -> Wait { residuatable = false }
        in
        match search world subst (flatten quantifier) with
        | exception Stuck _ -> symbolically ()
        | quantified ->
            if List.exists (fun (_, residue) -> residue <> []) quantified then
              (* The quantifier residuated: an infinite or
                 underdetermined domain. Enumeration is off the
                 table; proof is not. *)
              symbolically ()
            else
              let body_status (s, _) =
                match search world s (flatten body) with
                | exception Stuck _ -> `Unknown
                | [] -> `Fails
                | solutions -> if proven solutions then `Holds else `Unknown
              in
              let statuses = List.map body_status quantified in
              if List.mem `Fails statuses then Run []
              else if List.mem `Unknown statuses then
                Wait { residuatable = false }
              else Run [ (subst, []) ])
    | Positive (World.Apply (name, bindings)) -> (
        match Builtin.run subst name bindings with
        | Some (Builtin.Solved substs) ->
            Run (List.map (fun s -> (s, [])) substs)
        | Some Builtin.Stuck ->
            (* Pure membership waiting on bindings: the residuatable
               case. Constraining is not searching. *)
            Wait { residuatable = true }
        | None -> (
            match W.find world name with
            | Some (World.Table { facts; _ }) ->
                Run
                  (List.filter_map
                     (fun tuple ->
                       let target = Label.Map.map Pattern.of_term tuple in
                       Option.map
                         (fun s -> (s, []))
                         (unify_bindings subst bindings target))
                     facts)
            | Some (World.Rules rules) ->
                let witness_ground rule =
                  List.for_all
                    (fun label ->
                      match Label.Map.find_opt label bindings with
                      | Some pattern ->
                          Pattern.is_ground (Pattern.Subst.resolve subst pattern)
                      | None -> false)
                    rule.World.witness
                in
                if List.for_all witness_ground rules then
                  Run
                    (List.filter_map
                       (fun rule ->
                         let head, body = freshen rule in
                         Option.map
                           (fun s -> (s, flatten body))
                           (unify_bindings subst bindings head))
                       rules)
                else Wait { residuatable = false }
            | None ->
                (* Unreachable on admitted programs: the loader
                   checks references. An embedded host that lies
                   about its catalog surfaces here as a stuck goal
                   rather than a crash. *)
                Wait { residuatable = false }))

  (* Search yields substitutions with the goals they left
     residuated. A state where every remaining goal is residuatable
     membership is an answer; one with a stalled search is a
     non-total plan. *)
  and search world subst goals =
    match goals with
    | [] -> [ (subst, []) ]
    | _ -> (
        let rec select scanned = function
          | [] -> None
          | goal :: rest -> (
              match try_goal world subst goal with
              | Run branches -> Some (branches, List.rev_append scanned rest)
              | Wait _ -> select (goal :: scanned) rest)
        in
        match select [] goals with
        | Some (branches, remaining) ->
            List.concat_map
              (fun (subst, extra) -> search world subst (extra @ remaining))
              branches
        | None -> (
            (* Everything waits. Propagate over the waiting
               membership constraints before residuating or
               rejecting: they may be jointly unsatisfiable, or may
               pin a variable into a finite range that unblocks the
               rest. *)
            let attempts = List.map (fun g -> (g, try_goal world subst g)) goals in
            let membership =
              List.filter_map
                (fun (goal, attempt) ->
                  match (goal, attempt) with
                  | ( Positive (World.Apply (name, bindings)),
                      Wait { residuatable = true } ) ->
                      Some
                        ( name,
                          Label.Map.map (Pattern.Subst.resolve subst) bindings )
                  | _ -> None)
                attempts
            in
            match Propagate.analyze membership with
            | Propagate.Fail -> []
            | Propagate.Enumerate (var, lo, hi) ->
                let rec enumerate v acc =
                  if Z.gt v hi then acc
                  else
                    let branch =
                      match
                        Pattern.Subst.unify subst (Pattern.Var var)
                          (Pattern.Integer v)
                      with
                      | Some subst -> search world subst goals
                      | None -> []
                    in
                    enumerate (Z.succ v) (acc @ branch)
                in
                enumerate lo []
            | Propagate.Unknown ->
                if
                  List.for_all
                    (fun (_, attempt) ->
                      match attempt with
                      | Wait { residuatable } -> residuatable
                      | Run _ -> false)
                    attempts
                then [ (subst, goals) ]
                else raise (Stuck (List.map goal_name goals))))

  (* Solution variables: the free variables of the query. Whatever
     occurs only under a universal is bound per quantifier solution
     and is nobody's answer. *)
  let rec outer_vars = function
    | World.Apply (_, bindings) ->
        Label.Map.fold (fun _ p acc -> acc @ Pattern.vars p) bindings []
    | World.And (a, b) | World.Or (a, b) | World.Not (a, b) ->
        outer_vars a @ outer_vars b
    | World.Forall _ -> []

  let solve world expr =
    let variables =
      List.fold_left
        (fun acc v -> if List.mem v acc then acc else acc @ [ v ])
        [] (outer_vars expr)
      |> List.filter (fun v -> not (String.contains v '@'))
    in
    match search world Pattern.Subst.empty (flatten expr) with
    | exception Stuck goals -> Error (Non_total_plan goals)
    | outcomes ->
        let read (subst, residue_goals) =
          let bindings =
            List.map
              (fun v -> (v, Pattern.Subst.resolve subst (Pattern.Var v)))
              variables
          in
          let residue =
            List.filter_map
              (function
                | Positive (World.Apply (name, goal_bindings)) ->
                    Some
                      (name, Label.Map.map (Pattern.Subst.resolve subst) goal_bindings)
                | _ -> None)
              residue_goals
          in
          { bindings; residue }
        in
        let equal_residual (n, b) (m, c) =
          String.equal n m && Label.Map.equal Pattern.equal b c
        in
        let equal_solution a b =
          List.for_all2 (fun (_, p) (_, q) -> Pattern.equal p q) a.bindings
            b.bindings
          && List.length a.residue = List.length b.residue
          && List.for_all2 equal_residual a.residue b.residue
        in
        let solutions =
          List.fold_left
            (fun acc outcome ->
              let solution = read outcome in
              if List.exists (equal_solution solution) acc then acc
              else acc @ [ solution ])
            [] outcomes
        in
        Ok solutions
end

module Standalone = Make (World.Standalone)