type value = Bits of int64 | Truth of bool
type answer = Sat of value list | Unsat | Unknown of string
type query = { source : string; probes : (string * string) list }

let strategy = "(check-sat-using (then simplify solve-eqs bit-blast sat))"
let bits = function Bits value -> value | _ -> invalid_arg "expected bitvector"
let truth = function Truth value -> value | _ -> invalid_arg "expected boolean"
let index value = Printf.sprintf "#x%02x" value
let word value = Printf.sprintf "#x%016Lx" value
let all = function [] -> "true" | values -> "(and " ^ String.concat " " values ^ ")"
let any = function [] -> "false" | values -> "(or " ^ String.concat " " values ^ ")"
let equal_state left right = all (List.map (fun field ->
  Printf.sprintf "(= (%s %s) (%s %s))" field left field right)
  ["memory"; "registers"; "counters"; "heads"; "tails"; "addresses"; "values"; "positions"; "critical"])

let command source =
  if String.starts_with ~prefix:"(set-logic QF_BV)" source then strategy else "(check-sat)"

let connection ~timeout ~tick action =
  let input_read, input_write = Unix.pipe ~cloexec:true () in
  let output_read, output_write =
    try Unix.pipe ~cloexec:true () with error ->
      Unix.close input_read;
      Unix.close input_write;
      raise error in
  let executable = Option.value (Sys.getenv_opt "Z3") ~default:"z3" in
  let pid =
    try Unix.create_process executable [|executable; "-in"; "-smt2"|]
      input_read output_write output_write
    with error ->
      List.iter Unix.close [input_read; input_write; output_read; output_write];
      raise error in
  Unix.close input_read;
  Unix.close output_write;
  let input = Unix.out_channel_of_descr input_write in
  let pending = ref "" in
  let deadline = Unix.gettimeofday () +. float_of_int timeout /. 1000. +. 30. in
  let chunk = Bytes.create 8192 in
  let read () =
    if Unix.gettimeofday () > deadline then failwith "solver transport timed out";
    tick ();
    let ready, _, _ = Unix.select [output_read] [] [] 0.1 in
    if ready = [] then None else
    let size = Unix.read output_read chunk 0 (Bytes.length chunk) in
    if size = 0 then Some "" else Some (Bytes.sub_string chunk 0 size)
  in
  let rec line () = match String.index_opt !pending '\n' with
    | Some ending ->
        let result = String.sub !pending 0 ending in
        pending := String.sub !pending (ending + 1) (String.length !pending - ending - 1);
        result
    | None -> match read () with
        | None -> line ()
        | Some "" -> failwith ("solver closed its output: " ^ !pending)
        | Some text -> pending := !pending ^ text; line ()
  in
  let finish () =
    let buffer = Buffer.create 4096 in
    Buffer.add_string buffer !pending;
    pending := "";
    let rec loop () = match read () with
      | None -> loop ()
      | Some "" -> Buffer.contents buffer
      | Some text -> Buffer.add_string buffer text; loop ()
    in
    loop ()
  in
  let reaped = ref false in
  Fun.protect ~finally:(fun () ->
    close_out_noerr input;
    Unix.close output_read;
    if not !reaped then begin
      (try Unix.kill pid Sys.sigkill with Unix.Unix_error (Unix.ESRCH, _, _) -> ());
      ignore (Unix.waitpid [] pid)
    end) (fun () ->
      let result = action input line finish in
      let _, status = Unix.waitpid [] pid in
      reaped := true;
      if status <> Unix.WEXITED 0 then failwith "solver process failed";
      result)

let script ?(timeout=300000) ?(tick=fun () -> ()) source =
  connection ~timeout ~tick (fun input _ finish ->
    output_string input (source ^ "\n(exit)\n");
    flush input;
    finish ())

let solve ?(timeout=300000) ?(tick=fun () -> ()) query =
  connection ~timeout ~tick (fun input line finish ->
    output_string input query.source;
    List.iteri (fun i (kind, term) ->
      Printf.fprintf input "(define-fun q%d () %s %s)\n" i kind term) query.probes;
    output_string input (command query.source ^ "\n");
    flush input;
    let status = String.trim (line ()) in
    begin match status with
    | "sat" when query.probes <> [] ->
        let names = List.init (List.length query.probes) (Printf.sprintf "q%d") in
        Printf.fprintf input "(get-value (%s))\n" (String.concat " " names)
    | "unknown" -> output_string input "(get-info :reason-unknown)\n"
    | "sat" | "unsat" -> ()
    | _ -> failwith ("invalid solver response: " ^ status)
    end;
    output_string input "(exit)\n";
    flush input;
    let response = finish () in
    match status with
    | "unsat" -> Unsat
    | "unknown" -> Unknown response
    | _ when query.probes = [] -> Sat []
    | _ ->
        let entries = Sexp.list (Sexp.one response) in
        let table = Hashtbl.create (List.length entries) in
        List.iter (function
          | Sexp.List [Sexp.Atom name; value] ->
              let value = match value with
                | Sexp.Atom "true" -> Truth true
                | Sexp.Atom "false" -> Truth false
                | value -> Bits (Sexp.word value) in
              Hashtbl.add table name value
          | _ -> failwith "invalid model entry") entries;
        Sat (List.init (List.length query.probes) (fun i -> Hashtbl.find table (Printf.sprintf "q%d" i))))

let instruction state t =
  let open Syntax in
  let call name values = Printf.sprintf "(%s-step %s %s%s)" name state (index t)
    (if values = [] then "" else " " ^ String.concat " " values) in
  function
  | Store (a, v) -> [Printf.sprintf "(bvult (pending %s %s) #x08)" state (index t)],
      (fun sc -> call "store" [index a; word v; string_of_bool sc])
  | Load (a, r) -> [], (fun _ -> call "load" [index a; index r])
  | Fence -> [Printf.sprintf "(= (pending %s %s) #x00)" state (index t)],
      (fun _ -> call "fence" [])
  | Move (r, v) -> [], (fun _ -> call "move" [index r; word v])
  | Branch (r, v, pc) -> [], (fun _ -> call "branch" [index r; word v; index pc])
  | Enter -> [], (fun _ -> call "enter" [])
  | Leave -> [], (fun _ -> call "leave" [])
  | Exchange (a, r, v) -> [Printf.sprintf "(= (pending %s %s) #x00)" state (index t)],
      (fun _ -> call "exchange" [index a; index r; word v])

let property (case : Syntax.t) state =
  match case.goal with
  | Syntax.Overlap -> Printf.sprintf "(overlap %s)" state
  | Syntax.Outcome goal ->
      let finished = List.init 2 (fun t -> Printf.sprintf "(= (pc %s %s) %s)"
        state (index t) (index (Array.length case.threads.(t)))) in
      let registers = List.map (fun (t, r, v) -> Printf.sprintf "(= (reg %s %s %s) %s)"
        state (index t) (index r) (word v)) goal.registers in
      let memory = List.map (fun (a, v) -> Printf.sprintf "(= (at (memory %s) %s) %s)"
        state (index a) (word v)) goal.memory in
      let drained = if not goal.drained then [] else List.init 2 (fun t ->
        Printf.sprintf "(= (pending %s %s) #x00)" state (index t)) in
      all (finished @ registers @ memory @ drained)

let observations (case : Syntax.t) state =
  List.init 2 (fun t -> "Index", Printf.sprintf "(pc %s %s)" state (index t)) @
  List.init 8 (fun i -> "Word", Printf.sprintf "(reg %s %s %s)" state (index (i / 4)) (index (i mod 4))) @
  List.init (Array.length case.names) (fun a -> "Word", Printf.sprintf "(at (memory %s) %s)" state (index a)) @
  List.init 2 (fun t -> "Index", Printf.sprintf "(pending %s %s)" state (index t)) @
  List.init 2 (fun t -> "Index", Printf.sprintf "(at (critical %s) %s)" state (index t))

let compile ?(timeout=300000) ?fixed ?(require_goal=true) ?(blocks=[]) case model fences bound =
  Syntax.validate case;
  if bound < 1 || bound > 64 then invalid_arg "step bound outside 1..64";
  List.iter (fun (t, pc) ->
    if t < 0 || t > 1 || pc < 0 || pc >= Array.length case.Syntax.threads.(t)
    then invalid_arg "invalid fence location") fences;
  Option.iter (fun actors ->
    if Array.length actors <> bound || Array.exists (fun a -> a < 0 || a > 4) actors
    then invalid_arg "invalid fixed schedule") fixed;
  let buffer = Buffer.create 65536 in
  Buffer.add_string buffer (Layout.generate ~timeout);
  Buffer.add_string buffer (Store.read (Filename.concat (Store.root ()) "smt/machine.smt2"));
  let add line = Buffer.add_string buffer (line ^ "\n") in
  add ("(declare-const s0 State)\n(assert " ^ equal_state "s0" "initial" ^ ")");
  for step = 0 to bound - 1 do
    let state = Printf.sprintf "s%d" step in
    let next = Printf.sprintf "s%d" (step + 1) in
    let actor = Printf.sprintf "a%d" step in
    add (Printf.sprintf "(declare-const %s State)\n(declare-const %s Index)" next actor);
    let alternatives = List.concat (List.init 2 (fun t ->
      let instructions = Array.to_list (Array.mapi (fun pc ins ->
        let guards, call = instruction state t ins in
        let guards = if List.mem (t, pc) fences then
          Printf.sprintf "(= (pending %s %s) #x00)" state (index t) :: guards else guards in
        all (Printf.sprintf "(= (pc %s %s) %s)" state (index t) (index pc) ::
          equal_state next (call (model = Syntax.Sc)) :: guards)) case.threads.(t)) in
      [all [Printf.sprintf "(= %s %s)" actor (index t); any instructions];
       all [Printf.sprintf "(= %s %s)" actor (index (t + 2));
         Printf.sprintf "(bvugt (pending %s %s) #x00)" state (index t);
         equal_state next (Printf.sprintf "(drain %s %s)" state (index t))]])) in
    add ("(assert " ^ any (alternatives @ [all ["(= " ^ actor ^ " #x04)"; equal_state next state]]) ^ ")");
    if step > 0 && fixed = None then
      add (Printf.sprintf "(assert (=> (= %s #x04) (= a%d #x04)))" actor (step - 1));
    Option.iter (fun actors -> add (Printf.sprintf "(assert (= %s %s))" actor (index actors.(step)))) fixed
  done;
  let final = Printf.sprintf "s%d" bound in
  if require_goal then add ("(assert " ^ property case final ^ ")");
  List.iter (fun (left, right) -> add ("(assert " ^ any [
    Printf.sprintf "(distinct (reg %s #x00 #x00) %s)" final (word left);
    Printf.sprintf "(distinct (reg %s #x01 #x00) %s)" final (word right)] ^ ")")) blocks;
  let probes = List.init bound (fun i -> "Index", Printf.sprintf "a%d" i) @
    List.concat (List.init (bound + 1) (fun i -> observations case (Printf.sprintf "s%d" i))) in
  {source = Buffer.contents buffer; probes}

let candidate ~count ~budget clauses =
  if count < 0 then invalid_arg "negative fence candidate count";
  List.iter (List.iter (fun index ->
    if index < 0 || index >= count then invalid_arg "fence constraint index outside candidates")) clauses;
  let lines = ref [Store.read (Filename.concat (Store.root ()) "smt/fences.smt2")] in
  let add line = lines := line :: !lines in
  for i = 0 to count - 1 do add (Printf.sprintf "(declare-const f%d Bool)" i) done;
  let weights = List.init count (Printf.sprintf "(weight f%d)") in
  add (Printf.sprintf "(assert (<= (+ 0 %s) %d))" (String.concat " " weights) budget);
  List.iter (fun clause -> add ("(assert " ^ any (List.map (Printf.sprintf "f%d") clause) ^ ")")) clauses;
  {source = String.concat "\n" (List.rev !lines) ^ "\n";
   probes = List.init count (fun i -> "Bool", Printf.sprintf "f%d" i)}

let standalone query = query.source ^ command query.source ^ "\n"