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

open Tezos_protocol_environment_memory

(** Context creation *)

let create_block2 ctxt =
  Context.set ctxt ["a"; "b"] (MBytes.of_string "Novembre") >>= fun ctxt ->
  Context.set ctxt ["a"; "c"] (MBytes.of_string "Juin") >>= fun ctxt ->
  Context.set ctxt ["version";] (MBytes.of_string "0.0") >>= fun ctxt ->
  Lwt.return ctxt

let create_block3a ctxt =
  Context.del ctxt ["a"; "b"] >>= fun ctxt ->
  Context.set ctxt ["a"; "d"] (MBytes.of_string "Mars") >>= fun ctxt ->
  Lwt.return ctxt

let create_block3b ctxt =
  Context.del ctxt ["a"; "c"] >>= fun ctxt ->
  Context.set ctxt ["a"; "d"] (MBytes.of_string "Février") >>= fun ctxt ->
  Lwt.return ctxt

type t = {
  genesis: Context.t ;
  block2: Context.t ;
  block3a: Context.t ;
  block3b: Context.t ;
}

let wrap_context_init f _ () =
  let genesis = Context.empty in
  create_block2 genesis >>= fun block2 ->
  create_block3a block2 >>= fun block3a ->
  create_block3b block2 >>= fun block3b ->
  f { genesis; block2 ; block3a; block3b } >>= fun result ->
  Lwt.return result

(** Simple test *)

let c = function
  | None -> None
  | Some s -> Some (MBytes.to_string s)

let test_simple { block2 = ctxt } =
  Context.get ctxt ["version"] >>= fun version ->
  Assert.equal_string_option ~msg:__LOC__ (c version) (Some "0.0") ;
  Context.get ctxt ["a";"b"] >>= fun novembre ->
  Assert.equal_string_option (Some "Novembre") (c novembre) ;
  Context.get ctxt ["a";"c"] >>= fun juin ->
  Assert.equal_string_option ~msg:__LOC__ (Some "Juin") (c juin) ;
  Lwt.return_unit

let test_continuation { block3a = ctxt } =
  Context.get ctxt ["version"] >>= fun version ->
  Assert.equal_string_option ~msg:__LOC__ (Some "0.0") (c version) ;
  Context.get ctxt ["a";"b"] >>= fun novembre ->
  Assert.is_none ~msg:__LOC__ (c novembre) ;
  Context.get ctxt ["a";"c"] >>= fun juin ->
  Assert.equal_string_option ~msg:__LOC__ (Some "Juin") (c juin) ;
  Context.get ctxt ["a";"d"] >>= fun mars ->
  Assert.equal_string_option ~msg:__LOC__  (Some "Mars") (c mars) ;
  Lwt.return_unit

let test_fork { block3b = ctxt } =
  Context.get ctxt ["version"] >>= fun version ->
  Assert.equal_string_option ~msg:__LOC__ (Some "0.0") (c version) ;
  Context.get ctxt ["a";"b"] >>= fun novembre ->
  Assert.equal_string_option ~msg:__LOC__ (Some "Novembre") (c novembre) ;
  Context.get ctxt ["a";"c"] >>= fun juin ->
  Assert.is_none ~msg:__LOC__ (c juin) ;
  Context.get ctxt ["a";"d"] >>= fun mars ->
  Assert.equal_string_option ~msg:__LOC__ (Some "Février") (c mars) ;
  Lwt.return_unit

let test_replay { genesis = ctxt0 }  =
  Context.set ctxt0 ["version"] (MBytes.of_string "0.0") >>= fun ctxt1 ->
  Context.set ctxt1 ["a"; "b"] (MBytes.of_string "Novembre") >>= fun ctxt2 ->
  Context.set ctxt2 ["a"; "c"] (MBytes.of_string "Juin") >>= fun ctxt3 ->
  Context.set ctxt3 ["a"; "d"] (MBytes.of_string "July") >>= fun ctxt4a ->
  Context.set ctxt3 ["a"; "d"] (MBytes.of_string "Juillet") >>= fun ctxt4b ->
  Context.set ctxt4a ["a"; "b"] (MBytes.of_string "November") >>= fun ctxt5a ->
  Context.get ctxt4a ["a";"b"] >>= fun novembre ->
  Assert.equal_string_option ~msg:__LOC__ (Some "Novembre") (c novembre) ;
  Context.get ctxt5a ["a";"b"] >>= fun november ->
  Assert.equal_string_option ~msg:__LOC__ (Some "November") (c november) ;
  Context.get ctxt5a ["a";"d"] >>= fun july ->
  Assert.equal_string_option ~msg:__LOC__ (Some "July") (c july) ;
  Context.get ctxt4b ["a";"b"] >>= fun novembre ->
  Assert.equal_string_option ~msg:__LOC__ (Some "Novembre") (c novembre) ;
  Context.get ctxt4b ["a";"d"] >>= fun juillet ->
  Assert.equal_string_option ~msg:__LOC__ (Some "Juillet") (c juillet) ;
  Lwt.return_unit

let fold_keys s k ~init ~f =
  let rec loop k acc =
    Context.fold s k ~init:acc
      ~f:(fun file acc ->
          match file with
          | `Key k -> f k acc
          | `Dir k -> loop k acc) in
  loop k init
let keys t = fold_keys t ~init:[] ~f:(fun k acc -> Lwt.return (k :: acc))

let test_fold { genesis = ctxt } =
  Context.set ctxt ["a"; "b"] (MBytes.of_string "Novembre") >>= fun ctxt ->
  Context.set ctxt ["a"; "c"] (MBytes.of_string "Juin") >>= fun ctxt ->
  Context.set ctxt ["a"; "d"; "e"] (MBytes.of_string "Septembre") >>= fun ctxt ->
  Context.set ctxt ["f";] (MBytes.of_string "Avril") >>= fun ctxt ->
  Context.set ctxt ["g"; "h"] (MBytes.of_string "Avril") >>= fun ctxt ->
  keys ctxt [] >>= fun l ->
  Assert.equal_string_list_list ~msg:__LOC__
    [["a";"b"];
     ["a";"c"];
     ["a";"d";"e"];
     ["f"];
     ["g";"h"]] (List.sort compare l) ;
  keys ctxt ["a"] >>= fun l ->
  Assert.equal_string_list_list
    ~msg:__LOC__ [["a";"b"]; ["a";"c"]; ["a";"d";"e"]]
    (List.sort compare l) ;
  keys ctxt ["f"] >>= fun l ->
  Assert.equal_string_list_list ~msg:__LOC__ [] l ;
  keys ctxt ["g"] >>= fun l ->
  Assert.equal_string_list_list ~msg:__LOC__ [["g";"h"]] l ;
  keys ctxt ["i"] >>= fun l ->
  Assert.equal_string_list_list ~msg:__LOC__ [] l ;
  Lwt.return_unit

(******************************************************************************)

let tests = [
  "simple", test_simple ;
  "continuation", test_continuation ;
  "fork", test_fork ;
  "replay", test_replay ;
  "fold", test_fold ;
]

let tests =
  List.map
    (fun (n, f) -> Alcotest_lwt.test_case n `Quick (wrap_context_init f))
    tests