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

  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 only when
           the quantifier is decidable and closed; an undetermined
           quantifier or body defers the universal, and an infinite
           one therefore never runs. *)
        match search world subst (flatten quantifier) with
        | exception Stuck _ -> Wait { residuatable = false }
        | quantified ->
            if List.exists (fun (_, residue) -> residue <> []) quantified then
              Wait { residuatable = false }
            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 =
                  match rule.World.witness with
                  | None -> true
                  | Some label -> (
                      match Label.Map.find_opt label bindings with
                      | Some pattern ->
                          Pattern.is_ground (Pattern.Subst.resolve subst pattern)
                      | None -> false)
                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 ->
            if
              List.for_all
                (fun goal ->
                  match try_goal world subst goal with
                  | Wait { residuatable } -> residuatable
                  | Run _ -> false)
                goals
            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)