open Cmdliner

let log_src = Logs.Src.create "sql"

module Log = (val Logs.src_log log_src : Logs.LOG)

(* -- write subcommand -- *)

let open_or_create ~sw path =
  try Sqlite.open_ ~sw path with Failure _ | Eio.Io _ -> Sqlite.v ~sw path

let write_entry t (k, v) =
  Log.info (fun m -> m "put %s=%s" k v);
  Sqlite.put t k v

let write () db pairs =
  Eio_main.run @@ fun env ->
  let cwd = Eio.Stdenv.cwd env in
  Eio.Switch.run @@ fun sw ->
  let path = Eio.Path.(cwd / db) in
  let t = open_or_create ~sw path in
  List.iter (write_entry t) pairs;
  Sqlite.close t;
  Log.info (fun m -> m "wrote %d entries to %s" (List.length pairs) db)

let err_kv_pair s = Error (`Msg (Fmt.str "expected KEY=VALUE, got %S" s))

let kv_pair =
  let parse s =
    match String.split_on_char '=' s with
    | [ k; v ] -> Ok (k, v)
    | _ -> err_kv_pair s
  in
  let pp ppf (k, v) = Fmt.pf ppf "%s=%s" k v in
  Arg.conv (parse, pp)

let write_cmd =
  let db =
    Arg.(
      required
      & pos 0 (some string) None
      & info [] ~docv:"DB" ~doc:"Database file path.")
  in
  let pairs =
    Arg.(
      value & pos_right 0 kv_pair []
      & info [] ~docv:"KEY=VALUE" ~doc:"Key-value pairs to write.")
  in
  let info =
    Cmd.info "write" ~doc:"Write key-value pairs to a database."
      ~man:
        [
          `S Manpage.s_description;
          `P "Write key-value pairs to a database, creating it if needed.";
          `S Manpage.s_examples;
          `P "$(iname) mydb.db hello=world foo=bar";
        ]
  in
  Cmd.v info Term.(const write $ Vlog.setup "sql" $ db $ pairs)

(* -- read subcommand -- *)

let pp_value_pipe ppf = function
  | Sqlite.Vnull -> Fmt.string ppf "NULL"
  | Sqlite.Vint n -> Fmt.pf ppf "%Ld" n
  | Sqlite.Vfloat f -> Fmt.pf ppf "%f" f
  | Sqlite.Vblob s -> Fmt.pf ppf "BLOB(%d)" (String.length s)
  | Sqlite.Vtext s -> Fmt.pf ppf "%S" s

let read () db table_name =
  Eio_main.run @@ fun env ->
  let cwd = Eio.Stdenv.cwd env in
  Eio.Switch.run @@ fun sw ->
  let path = Eio.Path.(cwd / db) in
  let t = Sqlite.open_ ~sw path in
  (match table_name with
  | None ->
      (* Original KV behavior *)
      if Tty.is_tty () then begin
        let rows = ref [] in
        Sqlite.iter t ~f:(fun k v ->
            rows := [ Tty.Span.text k; Tty.Span.text v ] :: !rows);
        let rows = List.rev !rows in
        let table =
          Tty.Table.(
            of_rows ~border:Tty.Border.rounded
              [ column "key"; column "value" ]
              rows)
        in
        Tty.Table.pp Format.std_formatter table
      end
      else Sqlite.iter t ~f:(fun k v -> Fmt.pr "%s|%s@." k v)
  | Some name ->
      (* Generic table read *)
      if Tty.is_tty () then begin
        let schemas = Sqlite.tables t in
        let schema = List.find (fun s -> s.Sqlite.tbl_name = name) schemas in
        let cols =
          List.map
            (fun c -> Tty.Table.column c.Sqlite.col_name)
            schema.Sqlite.columns
        in
        let rows = ref [] in
        Sqlite.iter_table t name ~f:(fun _rowid values ->
            let row =
              List.map
                (fun v -> Tty.Span.text (Fmt.str "%a" pp_value_pipe v))
                values
            in
            rows := row :: !rows);
        let rows = List.rev !rows in
        let table = Tty.Table.(of_rows ~border:Tty.Border.rounded cols rows) in
        Tty.Table.pp Format.std_formatter table
      end
      else
        Sqlite.iter_table t name ~f:(fun _rowid values ->
            let strs =
              List.map (fun v -> Fmt.str "%a" pp_value_pipe v) values
            in
            Fmt.pr "%s@." (String.concat "|" strs)));
  Sqlite.close t

let read_cmd =
  let db =
    Arg.(
      required
      & pos 0 (some string) None
      & info [] ~docv:"DB" ~doc:"Database file path.")
  in
  let table_name =
    Arg.(
      value
      & opt (some string) None
      & info [ "t"; "table" ] ~docv:"TABLE"
          ~doc:"Table to read. Defaults to the kv table.")
  in
  let info =
    Cmd.info "read" ~doc:"Read all entries from a database."
      ~man:
        [
          `S Manpage.s_description;
          `P "Dump all key-value pairs from a SQLite database.";
          `P
            "When stdout is a terminal, output is rendered as a table. \
             Otherwise, entries are printed as pipe-separated values.";
          `P
            "Use $(b,--table) to read a specific table. Without it, reads the \
             default $(b,kv) table.";
          `S Manpage.s_examples;
          `P "$(iname) mydb.db";
          `P "$(iname) mydb.db -t users";
        ]
  in
  Cmd.v info Term.(const read $ Vlog.setup "sql" $ db $ table_name)

(* -- delete subcommand -- *)

let delete_entry t k =
  Log.info (fun m -> m "delete %s" k);
  Sqlite.delete t k

let delete () db keys =
  Eio_main.run @@ fun env ->
  let cwd = Eio.Stdenv.cwd env in
  Eio.Switch.run @@ fun sw ->
  let path = Eio.Path.(cwd / db) in
  let t = Sqlite.open_ ~sw path in
  List.iter (delete_entry t) keys;
  Sqlite.close t;
  Log.info (fun m -> m "deleted %d entries from %s" (List.length keys) db)

let delete_cmd =
  let db =
    Arg.(
      required
      & pos 0 (some string) None
      & info [] ~docv:"DB" ~doc:"Database file path.")
  in
  let keys =
    Arg.(
      value & pos_right 0 string [] & info [] ~docv:"KEY" ~doc:"Keys to delete.")
  in
  let info =
    Cmd.info "delete" ~doc:"Delete entries from a database."
      ~man:
        [
          `S Manpage.s_description;
          `P "Remove the given keys from a SQLite database.";
          `S Manpage.s_examples;
          `P "$(iname) mydb.db key1 key2";
        ]
  in
  Cmd.v info Term.(const delete $ Vlog.setup "sql" $ db $ keys)

(* -- tables subcommand -- *)

let pp_column (c : Sqlite.column) =
  let base =
    if c.col_affinity = "" then c.col_name
    else c.col_name ^ " " ^ c.col_affinity
  in
  if c.col_is_rowid_alias then base ^ " PRIMARY KEY" else base

let pp_schema (s : Sqlite.schema) =
  let cols = List.map pp_column s.columns in
  Fmt.pr "%s (%s)@." s.tbl_name (String.concat ", " cols)

let tables () db =
  Eio_main.run @@ fun env ->
  let cwd = Eio.Stdenv.cwd env in
  Eio.Switch.run @@ fun sw ->
  let path = Eio.Path.(cwd / db) in
  let t = Sqlite.open_ ~sw path in
  List.iter pp_schema (Sqlite.tables t);
  Sqlite.close t

let tables_cmd =
  let db =
    Arg.(
      required
      & pos 0 (some string) None
      & info [] ~docv:"DB" ~doc:"Database file path.")
  in
  let info =
    Cmd.info "tables" ~doc:"List tables in a database."
      ~man:
        [
          `S Manpage.s_description;
          `P "List all tables and their schemas in a SQLite database.";
          `S Manpage.s_examples;
          `P "$(iname) mydb.db";
        ]
  in
  Cmd.v info Term.(const tables $ Vlog.setup "sql" $ db)

(* -- main -- *)

let main_cmd =
  let info =
    Cmd.info "sql" ~version:"%%VERSION%%"
      ~doc:"SQLite-compatible key-value store."
      ~man:
        [
          `S Manpage.s_description;
          `P
            "$(tname) creates and reads SQLite-compatible databases using a \
             pure OCaml B-tree implementation. Files produced by $(tname) can \
             be read by the $(b,sqlite3) CLI and vice versa.";
        ]
  in
  Cmd.group info [ write_cmd; read_cmd; delete_cmd; tables_cmd ]

let () = exit (Cmd.eval main_cmd)