open Base_types

(* we store steps as binary chunks *)
module CS = Chunk_stack

module Config = struct
  type storage =
    | No_store
    | In_memory
    | On_disk_at of string

  let pp_storage out = function
    | No_store -> Fmt.string out "no-store"
    | In_memory -> Fmt.string out "in-memory"
    | On_disk_at file -> Fmt.fprintf out "(on-file :at %S)" file

  type t = {
    enabled: bool;
    storage: storage;
  }

  let default = { enabled=true; storage=In_memory }
  let empty = { enabled=false; storage=No_store }

  let pp out (self:t) =
    let { enabled; storage } = self in
    Fmt.fprintf out
      "(@[config@ :enabled %B@ :storage %a@])"
      enabled pp_storage storage

  let enable b self = {self with enabled=b}
  let store_in_memory self = {self with storage=In_memory}
  let store_on_disk_at file self = {self with storage=On_disk_at file}
  let no_store self = {self with storage=No_store}
end

(* where we store steps *)
module Storage = struct
  type t =
    | No_store
    | In_memory of CS.Buf.t
    | On_disk of string * out_channel

  let pp out = function
    | No_store -> Fmt.string out "no-store"
    | In_memory _ -> Fmt.string out "in-memory"
    | On_disk (file,_) -> Fmt.fprintf out "(on-file %S)" file
end

(* a step is just a unique integer ID.
   The actual step is stored in the chunk_stack. *)
type proof_step = Proof_ser.ID.t
type term_id = Proof_ser.ID.t

type lit = Lit.t
type term = Term.t

type t = {
  mutable enabled : bool;
  config: Config.t;
  buf: Buffer.t;
  mutable storage: Storage.t;
  mutable dispose: unit -> unit;
  mutable steps_writer: CS.Writer.t;
  mutable next_id: int;
  map_term: term_id Term.Tbl.t; (* term -> proof ID *)
  map_fun: term_id Fun.Tbl.t;
}
type proof_rule = t -> proof_step

module Step_vec = struct
  type elt=proof_step
  include VecI32
end

let disable (self:t) : unit =
  self.enabled <- false;
  self.storage <- Storage.No_store;
  self.dispose();
  self.steps_writer <- CS.Writer.dummy;
  ()

let nop_ _ = ()

let create ?(config=Config.default) () : t =
  (* acquire resources for logging *)
  let storage, steps_writer, dispose =
    match config.Config.storage with
    | Config.No_store ->
      Storage.No_store, CS.Writer.dummy, nop_

    | Config.In_memory ->
      let buf = CS.Buf.create ~cap:256 () in
      Storage.In_memory buf, CS.Writer.into_buf buf, nop_

    | Config.On_disk_at file ->
      let oc =
        open_out_gen [Open_creat; Open_wronly; Open_trunc; Open_binary] 0o644 file
      in
      let w = CS.Writer.into_channel oc in
      let dispose () = close_out oc in
      Storage.On_disk (file, oc), w, dispose
  in
  { enabled=config.Config.enabled;
    config;
    next_id=1;
    buf=Buffer.create 1_024;
    map_term=Term.Tbl.create 32;
    map_fun=Fun.Tbl.create 32;
    steps_writer; storage; dispose;
  }

let iter_chunks_ (r:CS.Reader.t) k =
  let rec loop () =
    CS.Reader.next r
      ~finish:nop_
      ~yield:(fun b i _len ->
          let step =
            Proof_ser.Bare.of_bytes_exn Proof_ser.Step.decode b ~off:i in
          k step;
          loop ()
        )
  in
  loop ()

let iter_steps_backward (self:t) : Proof_ser.Step.t Iter.t =
  fun yield ->
    begin match self.storage with
      | Storage.No_store -> ()
      | Storage.In_memory buf ->
        let r = CS.Reader.from_buf buf in
        iter_chunks_ r yield
      | Storage.On_disk (file, _oc) ->
        let ic = open_in file in
        let iter = CS.Reader.from_channel_backward ~close_at_end:true ic in
        iter_chunks_ iter yield
    end

let dummy_step : proof_step = Int32.min_int

let[@inline] enabled (self:t) = self.enabled

(* allocate a unique ID to refer to an event in the trace *)
let[@inline] alloc_id (self:t) : Proof_ser.ID.t =
  let n = self.next_id in
  self.next_id <- 1 + self.next_id;
  Int32.of_int n

(* emit a proof step *)
let emit_step_ (self:t) (step:Proof_ser.Step.t) : unit =
  if enabled self then (
    Buffer.clear self.buf;
    Proof_ser.Step.encode self.buf step;
    Chunk_stack.Writer.add_buffer self.steps_writer self.buf;
  )

let emit_fun_ (self:t) (f:Fun.t) : term_id =
  try Fun.Tbl.find self.map_fun f
  with Not_found ->
    let id = alloc_id self in
    Fun.Tbl.add self.map_fun f id;
    let f_name = ID.to_string (Fun.id f) in
    emit_step_ self
      Proof_ser.({ Step.id; view=Fun_decl {Fun_decl.f=f_name}});
    id

let rec emit_term_ (self:t) (t:Term.t) : term_id =
  try Term.Tbl.find self.map_term t
  with Not_found ->
    let view = match Term_cell.map (emit_term_ self) @@ Term.view t with
      | Term_cell.Bool b ->
        Proof_ser.Step_view.Expr_bool {Proof_ser.Expr_bool.b}

      | Term_cell.Ite (a,b,c) ->
        Proof_ser.Step_view.Expr_if {Proof_ser.Expr_if.cond=a; then_=b; else_=c}

      | Term_cell.Not a ->
        Proof_ser.Step_view.Expr_not {Proof_ser.Expr_not.f=a}

      | Term_cell.App_fun (f, arr) ->
        let f = emit_fun_ self f in
        Proof_ser.Step_view.Expr_app {Proof_ser.Expr_app.f; args=(arr:_ IArray.t:> _ array)}

      | Term_cell.Eq (a, b) ->
        Proof_ser.Step_view.Expr_eq {Proof_ser.Expr_eq.lhs=a; rhs=b}

      | LRA _ -> assert false (* TODO *)
    in

    let id = alloc_id self in
    emit_step_ self Proof_ser.({id; view});
    id

let emit_lit_ (self:t) (lit:Lit.t) : term_id =
  let sign = Lit.sign lit in
  let t = emit_term_ self (Lit.term lit) in
  if sign then t else Int32.neg t

let emit_redundant_clause _ ~hyps:_ _ = dummy_step
let emit_input_clause (lits:Lit.t Iter.t) (self:t) =
  if enabled self then (
    let lits = Iter.map (emit_lit_ self) lits |> Iter.to_array in
    let id = alloc_id self in
    emit_step_ self (
      Proof_ser.({Step.id; view=Step_view.Step_input {Step_input.c={Clause.lits}}})
    );
    id
  ) else dummy_step

let define_term _ _ _ = dummy_step
let proof_p1 _ _ (_pr:t) = dummy_step
let lemma_preprocess _ _ ~using:_ (_pr:t) = dummy_step
let lemma_true _ _ = dummy_step
let lemma_cc _ _ = dummy_step
let lemma_rw_clause _ ~using:_ (_pr:t) = dummy_step
let with_defs _ _ (_pr:t) = dummy_step
let del_clause _ _ (_pr:t) = ()
let emit_unsat_core _ (_pr:t) = dummy_step
let emit_unsat _ _ = ()

let lemma_lra _ _ = dummy_step

let lemma_bool_tauto _ _ = dummy_step
let lemma_bool_c _ _ _ = dummy_step
let lemma_bool_equiv _ _ _ = dummy_step
let lemma_ite_true ~ite:_ _ = dummy_step
let lemma_ite_false ~ite:_ _ = dummy_step

let lemma_isa_cstor ~cstor_t:_ _ (_pr:t) = dummy_step
let lemma_select_cstor ~cstor_t:_ _ (_pr:t) = dummy_step
let lemma_isa_split _ _ (_pr:t) = dummy_step
let lemma_isa_sel _ (_pr:t) = dummy_step
let lemma_isa_disj _ _ (_pr:t) = dummy_step
let lemma_cstor_inj _ _ _ (_pr:t) = dummy_step
let lemma_cstor_distinct _ _ (_pr:t) = dummy_step
let lemma_acyclicity _ (_pr:t) = dummy_step