(*---------------------------------------------------------------------------
   Copyright (c) 2025 Thomas Gazagnaire. All rights reserved.
   SPDX-License-Identifier: MIT
  ---------------------------------------------------------------------------*)

(** Fuzz tests for the pure OCaml B-tree backed key-value store. *)

open Crowbar

let truncate ?(max_len = 4096) s =
  if String.length s > max_len then String.sub s 0 max_len else s

(* Core KV operations *)

(** Roundtrip - put then get must return same value. *)
let test_roundtrip key value =
  let key = truncate key in
  let value = truncate value in
  let db = Sqlite.in_memory () in
  Sqlite.put db key value;
  let result = Sqlite.find db key in
  check_eq ~pp:Format.pp_print_string ~eq:( = ) (Option.get result) value

(** Delete removes key. *)
let test_delete_removes key value =
  let key = truncate key in
  let value = truncate value in
  let db = Sqlite.in_memory () in
  Sqlite.put db key value;
  Sqlite.delete db key;
  check (Option.is_none (Sqlite.find db key))

(** mem consistent with get. *)
let test_mem_consistent key value =
  let key = truncate key in
  let value = truncate value in
  let db = Sqlite.in_memory () in
  Sqlite.put db key value;
  check_eq ~pp:Format.pp_print_bool (Sqlite.mem db key)
    (Option.is_some (Sqlite.find db key))

(** Overwrite replaces value - last put wins. *)
let test_overwrite key value1 value2 =
  let key = truncate key in
  let value1 = truncate value1 in
  let value2 = truncate value2 in
  let db = Sqlite.in_memory () in
  Sqlite.put db key value1;
  Sqlite.put db key value2;
  check_eq ~pp:Format.pp_print_string ~eq:( = )
    (Option.get (Sqlite.find db key))
    value2

(* Table operations *)

(** Table isolation - same key in different tables must be independent. *)
let test_table_isolation key value1 value2 =
  let key = truncate key in
  let value1 = truncate value1 in
  let value2 = truncate value2 in
  let db = Sqlite.in_memory () in
  let t1 = Sqlite.Table.create db ~name:"table1" in
  let t2 = Sqlite.Table.create db ~name:"table2" in
  Sqlite.Table.put t1 key value1;
  Sqlite.Table.put t2 key value2;
  check_eq ~pp:Format.pp_print_string ~eq:( = )
    (Option.get (Sqlite.Table.find t1 key))
    value1;
  check_eq ~pp:Format.pp_print_string ~eq:( = )
    (Option.get (Sqlite.Table.find t2 key))
    value2

(** Table roundtrip. *)
let test_table_roundtrip key value =
  let key = truncate key in
  let value = truncate value in
  let db = Sqlite.in_memory () in
  let t = Sqlite.Table.create db ~name:"test" in
  Sqlite.Table.put t key value;
  check_eq ~pp:Format.pp_print_string ~eq:( = )
    (Option.get (Sqlite.Table.find t key))
    value

(* Crash safety *)

(** Put must not crash on arbitrary binary data. *)
let test_put_crash key value =
  let key = truncate key in
  let value = truncate value in
  let db = Sqlite.in_memory () in
  try Sqlite.put db key value with _exn -> ()

(** Get must not crash on arbitrary key. *)
let test_get_crash key =
  let key = truncate key in
  let db = Sqlite.in_memory () in
  try ignore (Sqlite.find db key) with _exn -> ()

(** Delete must not crash on arbitrary key. *)
let test_delete_crash key =
  let key = truncate key in
  let db = Sqlite.in_memory () in
  try Sqlite.delete db key with _exn -> ()

(** Mem must not crash on arbitrary key. *)
let test_mem_crash key =
  let key = truncate key in
  let db = Sqlite.in_memory () in
  try ignore (Sqlite.mem db key) with _exn -> ()

(* Boundary conditions *)

(** Empty key must work. *)
let test_empty_key value =
  let value = truncate value in
  let db = Sqlite.in_memory () in
  Sqlite.put db "" value;
  check_eq ~pp:Format.pp_print_string ~eq:( = )
    (Option.get (Sqlite.find db ""))
    value

(** Empty value must work. *)
let test_empty_value key =
  let key = truncate key in
  let db = Sqlite.in_memory () in
  Sqlite.put db key "";
  check_eq ~pp:Format.pp_print_string ~eq:( = )
    (Option.get (Sqlite.find db key))
    ""

(** Both empty must work. *)
let test_both_empty () =
  let db = Sqlite.in_memory () in
  Sqlite.put db "" "";
  check_eq ~pp:Format.pp_print_string ~eq:( = )
    (Option.get (Sqlite.find db ""))
    ""

(* Sequence operations *)

(** Multiple puts to same key must always have last value. *)
let test_multiple_puts key values =
  let key = truncate key in
  let values = List.map truncate values in
  if values = [] then check true
  else begin
    let db = Sqlite.in_memory () in
    List.iter (fun v -> Sqlite.put db key v) values;
    let last = List.hd (List.rev values) in
    check_eq ~pp:Format.pp_print_string ~eq:( = )
      (Option.get (Sqlite.find db key))
      last
  end

(** Put then delete then put must have second value. *)
let test_put_delete_put key value1 value2 =
  let key = truncate key in
  let value1 = truncate value1 in
  let value2 = truncate value2 in
  let db = Sqlite.in_memory () in
  Sqlite.put db key value1;
  Sqlite.delete db key;
  Sqlite.put db key value2;
  check_eq ~pp:Format.pp_print_string ~eq:( = )
    (Option.get (Sqlite.find db key))
    value2

(* Register all tests *)

let suite =
  ( "sqlite",
    [
      test_case "roundtrip" [ bytes; bytes ] test_roundtrip;
      test_case "delete removes" [ bytes; bytes ] test_delete_removes;
      test_case "mem consistent" [ bytes; bytes ] test_mem_consistent;
      test_case "overwrite" [ bytes; bytes; bytes ] test_overwrite;
      test_case "table isolation" [ bytes; bytes; bytes ] test_table_isolation;
      test_case "table roundtrip" [ bytes; bytes ] test_table_roundtrip;
      test_case "put crash safety" [ bytes; bytes ] test_put_crash;
      test_case "get crash safety" [ bytes ] test_get_crash;
      test_case "delete crash safety" [ bytes ] test_delete_crash;
      test_case "mem crash safety" [ bytes ] test_mem_crash;
      test_case "empty key" [ bytes ] test_empty_key;
      test_case "empty value" [ bytes ] test_empty_value;
      test_case "both empty" [ const () ] test_both_empty;
      test_case "multiple puts" [ bytes; list bytes ] test_multiple_puts;
      test_case "put delete put" [ bytes; bytes; bytes ] test_put_delete_put;
    ] )