type witness = {
  case : Syntax.t;
  model : Syntax.model;
  bound : int;
  fences : Syntax.location list;
  actors : int array;
}
type result = { witness : witness option; query : Solver.query }
type phase = Searching of int * Syntax.location list
  | Counterexample of witness | Excluded of int | Refining of int
  | Verified of Syntax.location list * int * int
type repair = {
  case : Syntax.t;
  fences : Syntax.location list;
  clauses : int list list;
  rounds : int;
  bound : int;
  states : int;
}

let decision = function
  | Solver.Unknown reason -> failwith ("solver returned unknown: " ^ reason)
  | answer -> answer

let verify case model fences bound values =
  let rec split count prefix = function
    | rest when count = 0 -> List.rev prefix, rest
    | value :: rest -> split (count - 1) (value :: prefix) rest
    | [] -> failwith "incomplete scheduler model"
  in
  let actors, observations = split bound [] values in
  let actors = Array.of_list (List.map (fun value -> Int64.to_int (Solver.bits value)) actors) in
  let states, _ = Machine.replay case model fences actors in
  let expected = List.concat (Array.to_list (Array.map Machine.observations states)) in
  if List.map Solver.bits observations <> expected then
    failwith "solver model disagrees with independent state-by-state replay";
  actors, states

let witness_sexp (witness : witness) =
  let open Sexp in
  List [Atom "witness"; List [Atom "version"; i 1];
    List [Atom "case"; Syntax.encode witness.case];
    List [Atom "model"; a (Syntax.model_name witness.model)];
    List [Atom "bound"; i witness.bound];
    List (Atom "fences" :: List.map (fun (t, pc) -> List [i t; i pc]) witness.fences);
    List (Atom "actors" :: Array.to_list (Array.map i witness.actors))]

let fields tag = function
  | Sexp.List (Sexp.Atom name :: fields) when name = tag -> fields
  | _ -> invalid_arg ("expected " ^ tag)
let single key fields = match Sexp.field key fields with
  | [value] -> value | _ -> invalid_arg ("invalid field " ^ key)

let replay (witness : witness) =
  Syntax.validate witness.case;
  Syntax.validate {witness.case with candidates=witness.fences};
  if witness.bound < 1 || witness.bound > 64 || Array.length witness.actors <> witness.bound then
    invalid_arg "witness length differs from its bound";
  let states, events = Machine.replay witness.case witness.model witness.fences witness.actors in
  if not (Machine.bad witness.case states.(Array.length states - 1)) then
    invalid_arg "witness does not reach the requested outcome";
  states, events

let load_witness path =
  let fields = fields "witness" (Sexp.one (Store.read path)) in
  if Sexp.integer (single "version" fields) <> 1 then invalid_arg "unsupported witness version";
  let witness = {
    case = Syntax.decode (single "case" fields);
    model = Syntax.model (Sexp.atom (single "model" fields));
    bound = Sexp.integer (single "bound" fields);
    fences = List.map Syntax.location (Sexp.field "fences" fields);
    actors = Array.of_list (List.map Sexp.integer (Sexp.field "actors" fields));
  } in
  ignore (replay witness);
  witness

let save folder result =
  Store.mkdir folder;
  Store.write (Filename.concat folder "query.smt2") (Solver.standalone result.query);
  let path = Filename.concat folder "witness.sexp" in
  match result.witness with
  | Some witness -> Store.write path (Sexp.print (witness_sexp witness) ^ "\n")
  | None -> if Sys.file_exists path then Sys.remove path

let check ?(timeout=300000) ?(tick=fun () -> ()) ?bound ?(fences=[]) ?folder ?fixed case model =
  let bound = Option.value bound ~default:case.Syntax.bound in
  let query = Solver.compile ~timeout ?fixed case model fences bound in
  let witness = match decision (Solver.solve ~timeout ~tick query) with
    | Solver.Unsat -> None
    | Solver.Sat values ->
        let actors, states = verify case model fences bound values in
        if not (Machine.bad case states.(bound)) then failwith "invalid final model state";
        Some {case; model; bound; fences; actors}
    | Solver.Unknown _ -> assert false in
  let result = {witness; query} in
  Option.iter (fun folder -> save folder result) folder;
  result

let blockers (witness : witness) =
  let states, _ = Machine.replay witness.case Syntax.Tso [] witness.actors in
  let result = ref [] in
  Array.iteri (fun i actor ->
    if actor < 2 && states.(i).buffers.(actor) <> [] then
      List.iteri (fun candidate location ->
        if location = (actor, states.(i).pc.(actor)) then result := candidate :: !result)
        witness.case.candidates) witness.actors;
  List.sort_uniq compare !result

let repair_sexp (repair : repair) =
  let open Sexp in
  List [Atom "repair"; List [Atom "case"; Syntax.encode repair.case];
    List [Atom "rounds"; i repair.rounds];
    List [Atom "bound"; i repair.bound]; List [Atom "states"; i repair.states];
    List (Atom "fences" :: List.map (fun (t, pc) -> List [i t; i pc]) repair.fences);
    List (Atom "clauses" :: List.map (fun clause -> List (List.map i clause)) repair.clauses)]

let load_repair path =
  let fields = fields "repair" (Sexp.one (Store.read path)) in
  let repair = {
    case = Syntax.decode (single "case" fields);
    rounds = Sexp.integer (single "rounds" fields);
    bound = Sexp.integer (single "bound" fields);
    states = Sexp.integer (single "states" fields);
    fences = List.map Syntax.location (Sexp.field "fences" fields);
    clauses = List.map (fun clause -> List.map Sexp.integer (Sexp.list clause)) (Sexp.field "clauses" fields);
  } in
  if repair.bound < repair.case.bound || repair.bound > 64 || repair.states < 1 ||
     repair.rounds <> List.length repair.clauses + 1 ||
     List.sort_uniq compare repair.fences <> List.sort compare repair.fences ||
     not (List.for_all (fun location -> List.mem location repair.case.candidates) repair.fences)
  then invalid_arg "invalid repair certificate";
  repair

let repair ?(timeout=300000) ?(tick=fun () -> ()) ?(notify=fun _ -> ()) case folder =
  Store.mkdir folder;
  List.iter (fun name ->
    let path = Filename.concat folder name in
    if Sys.file_exists path then Sys.remove path) ["repair.sexp"; "witness.sexp"];
  let count = List.length case.Syntax.candidates in
  let clauses = ref [] in
  let rounds = ref 0 in
  let search_bound = ref case.bound in
  let rec search budget =
    if budget > count then failwith "fence search exhausted";
    let candidate = Solver.candidate ~count ~budget !clauses in
    Store.write (Filename.concat folder "synthesis.smt2") (Solver.standalone candidate);
    match decision (Solver.solve ~timeout ~tick candidate) with
    | Solver.Unsat -> notify (Excluded budget); search (budget + 1)
    | Solver.Sat values ->
        let selected = List.filter_map Fun.id (List.mapi (fun i value ->
          if Solver.truth value then Some (List.nth case.candidates i) else None) values) in
        notify (Searching (budget, selected));
        let target = Filename.concat folder (Printf.sprintf "round%02d" !rounds) in
        let bounded = check ~timeout ~tick ~bound:!search_bound ~fences:selected ~folder:target case Syntax.Tso in
        let result, closure = match bounded.witness with
          | Some _ -> bounded, None
          | None ->
              let graph = Machine.explore case Syntax.Tso selected in
              match Machine.counterexample graph with
              | None -> bounded, Some graph
              | Some actors ->
                  search_bound := max !search_bound (Array.length actors);
                  notify (Refining !search_bound);
                  let deeper = check ~timeout ~tick ~bound:(Array.length actors)
                    ~fixed:actors ~fences:selected ~folder:target case Syntax.Tso in
                  if deeper.witness = None then failwith "graph counterexample has no SMT execution";
                  deeper, None
        in
        incr rounds;
        begin match result.witness with
        | None ->
            let lower = Solver.candidate ~count ~budget:(List.length selected - 1) !clauses in
            begin match decision (Solver.solve ~timeout ~tick lower) with
            | Solver.Unsat -> () | _ -> failwith "minimum fence certificate is satisfiable"
            end;
            Store.write (Filename.concat folder "minimal.smt2") (Solver.standalone lower);
            let graph = match closure with Some graph -> graph | None -> assert false in
            let report = {case; fences = selected; clauses = !clauses; rounds = !rounds;
              bound= !search_bound; states=Array.length graph.states} in
            Store.write (Filename.concat folder "repair.sexp") (Sexp.print (repair_sexp report) ^ "\n");
            notify (Verified (selected, report.bound, report.states));
            Some report
        | Some witness ->
            if !rounds = 1 then Store.write (Filename.concat folder "witness.sexp")
              (Sexp.print (witness_sexp witness) ^ "\n");
            notify (Counterexample witness);
            let clause = blockers witness in
            if clause = [] then None else begin
              if List.mem clause !clauses then failwith "repeated excluded counterexample";
              clauses := !clauses @ [clause];
              search budget
            end
        end
    | Solver.Unknown _ -> assert false
  in
  search 0

let certify ?(timeout=300000) ?(tick=fun () -> ()) path =
  let report = load_repair path in
  List.iteri (fun i clause ->
    let path = Filename.concat (Filename.dirname path) (Printf.sprintf "round%02d/witness.sexp" i) in
    let witness = load_witness path in
    if witness.case <> report.case || witness.model <> Syntax.Tso || witness.bound > report.bound || blockers witness <> clause then
      invalid_arg "counterexample does not support its fence constraint") report.clauses;
  let lower = Solver.candidate ~count:(List.length report.case.candidates)
    ~budget:(List.length report.fences - 1) report.clauses in
  begin match decision (Solver.solve ~timeout ~tick lower) with
  | Solver.Unsat -> () | _ -> invalid_arg "minimum fence count is unsupported"
  end;
  let graph = Machine.explore report.case Syntax.Tso report.fences in
  if graph.violations <> [] || Array.length graph.states <> report.states then
    invalid_arg "repair state graph does not support the certificate";
  report

let audit ?(timeout=300000) ?(tick=fun () -> ()) path folder =
  let report = certify ~timeout ~tick path in
  let verified = check ~timeout ~tick ~bound:report.bound ~fences:report.fences ~folder report.case Syntax.Tso in
  if verified.witness <> None then invalid_arg "repair still admits the forbidden outcome";
  report

let validate_schedule ?(timeout=300000) case fences actors =
  let bound = Array.length actors in
  let query = Solver.compile ~timeout ~fixed:actors ~require_goal:false case Syntax.Tso fences bound in
  match decision (Solver.solve ~timeout query) with
  | Solver.Sat values -> ignore (verify case Syntax.Tso fences bound values)
  | _ -> failwith "repaired schedule has no matching SMT execution"

let outcomes ?(timeout=300000) ?(tick=fun () -> ()) case model =
  let case = {case with Syntax.goal = Syntax.Outcome {registers=[]; memory=[]; drained=false}} in
  let rec loop found =
    if List.length found = 16 then failwith "outcome enumeration exceeded sixteen distinct register pairs";
    let query = Solver.compile ~timeout ~blocks:found case model [] case.bound in
    match decision (Solver.solve ~timeout ~tick query) with
    | Solver.Unsat -> List.sort_uniq compare found
    | Solver.Sat values ->
        let _, states = verify case model [] case.bound values in
        let final = states.(case.bound) in
        loop ((final.regs.(0).(0), final.regs.(1).(0)) :: found)
    | Solver.Unknown _ -> assert false
  in
  loop []