module String_map = Map.Make (String)

type state = Term.t Label.Map.t list String_map.t
(* relation name -> stored tuples, the whole tree of one state *)

type t = {
  schemas : (Label.t * World.primitive) list String_map.t;
  states : (string, state) Hashtbl.t;  (* digest -> state; append-only *)
  mutable tips : (string * string) list;  (* branch -> digest; the one mutable place *)
}

let forced_label name =
  match Label.of_string name with Ok l -> l | Error _ -> assert false

let relation_l = forced_label "relation"
let tuple_l = forced_label "tuple"
let before_l = forced_label "before"
let after_l = forced_label "after"

(* Canonical serialization: relations in name order (the map gives
   it), tuples in term order, printed in canonical notation. Equal
   states digest equally because terms are canonical by
   construction. *)
let digest_of_state (state : state) =
  let buffer = Buffer.create 256 in
  String_map.iter
    (fun name tuples ->
      Buffer.add_string buffer name;
      let terms =
        List.map (fun tuple -> Term.record (Label.Map.bindings tuple)) tuples
        |> List.filter_map Result.to_option
        |> List.sort Term.compare
      in
      List.iter
        (fun term -> Buffer.add_string buffer (Format.asprintf "|%a" Term.pp term))
        terms;
      Buffer.add_char buffer '\n')
    state;
  Digest.to_hex (Digest.string (Buffer.contents buffer))

let store world state =
  let digest = digest_of_state state in
  if not (Hashtbl.mem world.states digest) then
    Hashtbl.add world.states digest state;
  digest

let create ~schemas =
  let catalog =
    List.fold_left
      (fun acc (name, schema) -> String_map.add name schema acc)
      String_map.empty schemas
  in
  let root : state =
    String_map.map (fun _ -> []) catalog
  in
  let world = { schemas = catalog; states = Hashtbl.create 16; tips = [] } in
  let root_digest = store world root in
  world.tips <- [ ("main", root_digest) ];
  world

let state_at world = function
  | Some (Term.String digest) -> Hashtbl.find_opt world.states digest
  | Some _ -> None
  | None -> (
      match List.assoc_opt "main" world.tips with
      | Some digest -> Hashtbl.find_opt world.states digest
      | None -> None)

let find world ?at name =
  match (String_map.find_opt name world.schemas, state_at world at) with
  | Some schema, Some state ->
      let facts = Option.value ~default:[] (String_map.find_opt name state) in
      let schema = List.map (fun (l, p) -> (l, World.Primitive p)) schema in
      Some (World.Table { schema; facts })
  | _ -> None

(* [Insert]: the transition relation between states (decision 3).
   Ground relation name, ground tuple, ground predecessor state;
   the tuple must match the schema exactly and satisfy its
   primitive domains, or there is no such transition. *)
let insert world bindings =
  let ground label =
    Option.bind (Label.Map.find_opt label bindings) Pattern.to_term
  in
  match (ground relation_l, ground tuple_l, ground before_l) with
  | Some (Term.Reference relation), Some (Term.Record tuple), Some (Term.String before)
    -> (
      match
        (String_map.find_opt relation world.schemas, Hashtbl.find_opt world.states before)
      with
      | Some schema, Some state ->
          let conforms =
            List.length schema = Label.Map.cardinal tuple
            && List.for_all
                 (fun (label, domain) ->
                   match Label.Map.find_opt label tuple with
                   | Some value -> World.member_of_primitive value domain
                   | None -> false)
                 schema
          in
          if not conforms then World.T_fail
          else
            let facts = Option.value ~default:[] (String_map.find_opt relation state) in
            let successor = String_map.add relation (facts @ [ tuple ]) state in
            World.T_state (Term.string (store world successor))
      | _ -> World.T_fail)
  | Some _, Some _, Some _ -> World.T_fail
  | _ ->
      (* Something essential is not ground yet: wait for the chain. *)
      World.T_stuck

let transition world name bindings =
  match name with "Insert" -> Some (insert world bindings) | _ -> None

let tip world branch =
  Option.map Term.string (List.assoc_opt branch world.tips)

(* The single effect: compare-and-swap on the tip (decision 8). *)
let commit world ~branch ~expect ~next =
  match (expect, next, List.assoc_opt branch world.tips) with
  | Term.String expected, Term.String next_digest, Some current
    when String.equal current expected && Hashtbl.mem world.states next_digest ->
      world.tips <-
        (branch, next_digest) :: List.remove_assoc branch world.tips;
      Ok ()
  | _ -> Error `Conflict

(* Linearity (decision 6): within one expression, a state variable
   is consumed at most once as [before:] and produced at most once
   as [after:]. Sequencing comes from the chain, and the chain must
   be a chain. *)
let linear expr =
  let rec applications = function
    | World.Apply (name, bindings) -> [ (name, bindings) ]
    | World.And (a, b) | World.Or (a, b) | World.Not (a, b) | World.Forall (a, b)
      ->
        applications a @ applications b
  in
  let uses label =
    applications expr
    |> List.filter (fun (name, _) -> String.equal name "Insert")
    |> List.filter_map (fun (_, bindings) ->
           match Label.Map.find_opt label bindings with
           | Some (Pattern.Var v) -> Some v
           | _ -> None)
  in
  let duplicate vars =
    let sorted = List.sort String.compare vars in
    let rec scan = function
      | a :: b :: _ when String.equal a b -> Some a
      | _ :: rest -> scan rest
      | [] -> None
    in
    scan sorted
  in
  match (duplicate (uses before_l), duplicate (uses after_l)) with
  | Some v, _ ->
      Error
        (Printf.sprintf
           "state %s consumed twice: a fork inside one expression is not a chain" v)
  | _, Some v -> Error (Printf.sprintf "state %s produced twice" v)
  | None, None -> Ok ()