(**************************************************************************)
(*                                                                        *)
(*    Copyright (c) 2014 - 2018.                                          *)
(*    Dynamic Ledger Solutions, Inc. <contact@tezos.com>                  *)
(*                                                                        *)
(*    All rights reserved. No warranty, explicit or implicit, provided.   *)
(*                                                                        *)
(**************************************************************************)

open Rresult

type t = {
  dir : Lmdb.t ;
  parent : (Lmdb.rw Lmdb.txn * Lmdb.db * Lmdb.rw Lmdb.cursor) Lwt.key ;
}

type key = string list
type value = MBytes.t

type error += Unknown of string list

let () =
  Error_monad.register_error_kind
    `Permanent
    ~id:"raw_store.unknown"
    ~title:"Missing key in store"
    ~description:"Missing key in store"
    ~pp:(fun ppf keys ->
        Format.fprintf ppf
          "Missing key in store: %s"
          (String.concat "/" keys))
    Data_encoding.(obj1 (req "key" (list string)))
    (function Unknown keys -> Some keys | _ -> None)
    (fun keys -> Unknown keys)

let concat = String.concat "/"
let split = String.split_on_char '/'

let lwt_fail_error err =
  Lwt.fail_with (Lmdb.string_of_error err)

let of_result = function
  | Ok res -> Lwt.return res
  | Error err -> lwt_fail_error err

let (>>=?) v f =
  match v with
  | Error err -> lwt_fail_error err
  | Ok v -> f v

let init ?mapsize path =
  if not (Sys.file_exists path) then Unix.mkdir path 0o755 ;
  match Lmdb.opendir ?mapsize ~flags:[NoTLS] path 0o644 with
  | Ok dir -> return { dir ; parent = Lwt.new_key () }
  | Error err -> failwith "%a" Lmdb.pp_error err

let close { dir } = Lmdb.closedir dir

let known { dir ; parent } key =
  begin match Lwt.get parent with
    | Some (txn, db, _cursor) -> Lmdb.mem txn db (concat key)
    | None ->
        Lmdb.with_ro_db dir ~f:begin fun txn db ->
          Lmdb.mem txn db (concat key)
        end
  end |> of_result

let read_opt { dir ; parent } key =
  begin match Lwt.get parent with
    | Some (txn, db, _cursor) -> Lmdb.get txn db (concat key) >>| MBytes.copy
    | None ->
        Lmdb.with_ro_db dir ~f:begin fun txn db ->
          Lmdb.get txn db (concat key) >>| MBytes.copy
        end
  end |> function
  | Ok v -> Lwt.return_some v
  | Error KeyNotFound -> Lwt.return_none
  | Error err -> lwt_fail_error err

let read { dir ; parent } key =
  begin match Lwt.get parent with
    | Some (txn, db, _cursor) -> Lmdb.get txn db (concat key) >>| MBytes.copy
    | None ->
        Lmdb.with_ro_db dir ~f:begin fun txn db ->
          Lmdb.get txn db (concat key) >>| MBytes.copy
        end
  end |> function
  | Ok v -> return v
  | Error _err -> fail (Unknown key)

let read_exn { dir ; parent } key =
  begin match Lwt.get parent with
    | Some (txn, db, _cursor) -> Lmdb.get txn db (concat key) >>| MBytes.copy
    | None ->
        Lmdb.with_ro_db dir ~f:begin fun txn db ->
          Lmdb.get txn db (concat key) >>| MBytes.copy
        end
  end |> of_result

let store { dir ; parent } k v =
  begin match Lwt.get parent with
    | Some (txn, db, _cursor) -> Lmdb.put txn db (concat k) v
    | None ->
        Lmdb.with_rw_db dir ~f:begin fun txn db ->
          Lmdb.put txn db (concat k) v
        end
  end |> of_result

let remove { dir ; parent } k =
  let remove txn db =
    match Lmdb.del txn db (concat k) with
    | Ok () -> Ok ()
    | Error KeyNotFound -> Ok ()
    | Error err -> Error err in
  begin match Lwt.get parent with
    | Some (txn, db, _cursor) -> remove txn db
    | None -> Lmdb.with_rw_db dir ~f:remove
  end |> of_result

let is_prefix s s' =
  String.(length s <= length s' && compare s (sub s' 0 (length s)) = 0)

let known_dir { dir ; parent } k =
  let k = concat k in
  let cursor_fun cursor =
    Lmdb.cursor_at cursor k >>= fun () ->
    Lmdb.cursor_get cursor >>| fun (first_k, _v) ->
    (is_prefix k (MBytes.to_string first_k))
  in
  begin match Lwt.get parent with
    | Some (txn, db, _cursor) ->
        Lmdb.with_cursor txn db ~f:cursor_fun
    | None ->
        Lmdb.with_ro_db dir ~f:begin fun txn db ->
          Lmdb.with_cursor txn db ~f:cursor_fun
        end
  end |> of_result

let remove_dir { dir ; parent } k =
  let k = concat k in
  let cursor_fun cursor =
    Lmdb.cursor_at cursor k >>= fun () ->
    Lmdb.cursor_iter cursor ~f:begin fun (kk, _v) ->
      let kk_string = MBytes.to_string kk in
      if is_prefix k kk_string then begin
        Lmdb.cursor_del cursor
      end
      else Error KeyNotFound
    end in
  begin match Lwt.get parent with
    | Some (txn, db, _cursor) ->
        Lmdb.with_cursor txn db ~f:cursor_fun
    | None ->
        Lmdb.with_rw_db dir ~f:begin fun txn db ->
          Lmdb.with_cursor txn db ~f:cursor_fun
        end
  end |> function
  | Error KeyNotFound
  | Ok () -> Lwt.return_unit
  | Error err -> lwt_fail_error err

let list_equal l1 l2 len =
  if len < 0 || len > List.length l1 || len > List.length l2
  then invalid_arg "list_compare: invalid len" ;
  let rec inner l1 l2 len =
    match len, l1, l2 with
    | 0, _, _ -> true
    | _, [], _
    | _, _, [] -> false
    | _, h1 :: t1, h2 :: t2 ->
        if h1 <> h2 then false
        else inner t1 t2 (pred len)
  in
  inner l1 l2 len

let is_child ~parent ~child =
  let plen = List.length parent in
  let clen = List.length child in
  clen > plen && list_equal parent child plen

let list_sub l pos len =
  if len < 0 || pos < 0 || pos + len > List.length l then
    invalid_arg "list_sub" ;
  let rec inner (acc, n) = function
    | [] -> List.rev acc
    | h :: t ->
        if n = 0 then List.rev acc
        else inner (h :: acc, pred n) t in
  inner ([], len) l

let with_rw_cursor_lwt ?sync ?metasync ?flags ?name { dir ; parent } ~f =
  let local_parent =
    match Lwt.get parent with
    | None -> None
    | Some (txn, _db, _cursor) -> Some txn in
  Lmdb.create_rw_txn
    ?sync ?metasync ?parent:local_parent dir >>=? fun txn ->
  Lmdb.opendb ?flags ?name txn >>=? fun db ->
  Lmdb.opencursor txn db >>=? fun cursor ->
  Lwt.with_value parent (Some (txn, db, cursor)) begin fun () ->
    Lwt.try_bind (fun () -> f cursor)
      begin fun res ->
        Lmdb.cursor_close cursor ;
        Lmdb.commit_txn txn >>=? fun () ->
        Lwt.return res
      end
      begin fun exn ->
        Lmdb.cursor_close cursor ;
        Lmdb.abort_txn txn ;
        Lwt.fail exn
      end
  end

let cursor_next_lwt cursor acc f =
  match Lmdb.cursor_next cursor with
  | Error KeyNotFound -> acc
  | Error err -> lwt_fail_error err
  | Ok () -> Lwt.bind acc f

let fold t k ~init ~f =
  let base_len = List.length k in
  let rec inner ht cursor acc =
    Lmdb.cursor_get cursor >>=? fun (kk, _v) ->
    let kk = MBytes.to_string kk in
    let kk_split = split kk in
    match is_child ~child:kk_split ~parent:k with
    | false -> Lwt.return acc
    | true ->
        let cur_len = List.length kk_split in
        if cur_len = succ base_len then begin
          cursor_next_lwt cursor (f (`Key kk_split) acc) (inner ht cursor)
        end
        else begin
          let dir = list_sub kk_split 0 (succ base_len) in
          if Hashtbl.mem ht dir then
            cursor_next_lwt cursor (Lwt.return acc) (inner ht cursor)
          else begin
            Hashtbl.add ht dir () ;
            cursor_next_lwt cursor (f (`Dir dir) acc) (inner ht cursor)
          end
        end in
  with_rw_cursor_lwt t ~f:begin fun cursor ->
    match Lmdb.cursor_at cursor (concat k) with
    | Error KeyNotFound -> Lwt.return init
    | Error err -> lwt_fail_error err
    | Ok () ->
        let ht = Hashtbl.create 31 in
        inner ht cursor init
  end

let fold_keys t k ~init ~f =
  with_rw_cursor_lwt t ~f:begin fun cursor ->
    match Lmdb.cursor_at cursor (concat k) with
    | Error KeyNotFound -> Lwt.return init
    | Error err -> lwt_fail_error err
    | Ok () ->
        let rec inner acc =
          Lmdb.cursor_get cursor >>=? fun (kk, _v) ->
          let kk = MBytes.to_string kk in
          let kk_split = split kk in
          match is_child ~child:kk_split ~parent:k with
          | false -> Lwt.return acc
          | true -> cursor_next_lwt cursor (f kk_split acc) inner
        in inner init
  end

let keys t =
  fold_keys t ~init:[] ~f:(fun k acc -> Lwt.return (k :: acc))