(* TODO
   ! use lock
   !? Scanf not thread-safe
   schedule tasks
   sign: check timestamps before signing *)
(* IMPROVE
   Bin encode/decode
     - catch failure
     - should use little-endian
   list_issue_date: save timestamp of key generation
   key validity period:
     more checks + do not exceed lookahead
   refacto common parts with secmod_eddsa *)
let src = Logs.Src.create "mte.secmod_eddsa"

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

(* - *)
open Syntax
open Crypto
open Time
module Sfn = Mfat.Sfn
module Spath = Mfat.Spath
module Cfg = Config.Exchange_secmod_eddsa

type key = {
  priv: EddsaPrivateKey.t;
  pub: EddsaPublicKey.t;
  t1: TimeAbsolute.t;
  t2: TimeAbsolute.t;
}

type t = {
  fs: Fat.t;
  sm_priv: EddsaPrivateKey.t;
  sm_pub: EddsaPublicKey.t;
  ht: (EddsaPublicKey.t, key) Hashtbl.t;
}

(* todo: this should be encoded in little-endian *)
let key_bin =
  let open Bin in
  record (fun t1 t2 priv ->
      let pub = EddsaPrivateKey.pub_of_priv priv in
      { t1; t2; priv; pub })
  |+ field TimeAbsolute.bin (fun t -> t.t1)
  |+ field TimeAbsolute.bin (fun t -> t.t2)
  |+ field EddsaPrivateKey.bin (fun t -> t.priv)
  |> sealr

(* for sm_key only *)
let read_eddsa fs spath =
  Log.debug (fun m -> m "reading key file `%a`" Spath.pp spath);
  let* data = Fat.read fs spath |> unwrap_msg in
  let* priv = EddsaPrivateKey.of_octets data in
  let pub = EddsaPrivateKey.pub_of_priv priv in
  Ok (priv, pub)

let write_eddsa fs spath priv =
  Log.debug (fun m -> m "writing key file `%a`" Spath.pp spath);
  let data = EddsaPrivateKey.to_octets priv in
  Fat.write fs spath data |> unwrap_msg

let key_spath k =
  let sfn_res = String.sub (EddsaPublicKey.to_b32 k.pub) 0 8 |> Sfn.of_string in
  match sfn_res with
  | Error _ -> failwith "not possible"
  | Ok sfn -> Spath.(Cfg.key_dir / sfn)

let read_key fs spath =
  Log.debug (fun m -> m "reading key file `%a`" Spath.pp spath);
  let* data = Fat.read fs spath |> unwrap_msg in
  (* todo: catch failure *)
  let k = Bin.decode key_bin data (ref 0) in
  Ok k

let write_key fs k =
  let spath = key_spath k in
  Log.debug (fun m -> m "writing key file `%a`" Spath.pp spath);
  let data = Bin.to_string key_bin k in
  Fat.write fs spath data |> unwrap_msg

let delete_file fs spath =
  Log.debug (fun m -> m "delete key file `%a`" Spath.pp spath);
  let+ () = Fat.remove fs spath |> unwrap_msg in
  ()

let gen_key t1 t2 =
  let priv, pub = EddsaPrivateKey.generate () in
  let k = { priv; pub; t1; t2 } in
  Log.debug (fun m -> m "generated key `%s`" (EddsaPublicKey.to_b32 pub));
  k

let sort_keys l = List.sort (fun a b -> TimeAbsolute.compare a.t2 b.t2) l

let split_in_periodes ~start ~end_ =
  assert (start < end_);
  (* no overlap on first periode *)
  let t1 = start in
  let t2 = TimeAbsolute.add start Cfg.duration in
  let acc = [ (t1, t2) ] in
  let start = t2 in
  let rec go acc start end_ =
    let t1 = TimeAbsolute.sub start Cfg.overlap_duration in
    let t2 = TimeAbsolute.add start Cfg.duration in
    if t2 > end_ then acc else go ((t1, t2) :: acc) t2 end_
  in
  go acc start end_

(* try to not generate keys with validity start in the past *)
let gen_additional_keys_until_lookahead ~now l =
  let start =
    match List.rev (sort_keys l) with
    | [] -> now
    | hd :: _ -> TimeAbsolute.sub hd.t2 Cfg.overlap_duration
  in
  let end_ = TimeAbsolute.add now Cfg.lookahead_sign in
  if TimeAbsolute.compare start end_ >= 0 then []
  else
    let periodes = split_in_periodes ~start ~end_ in
    let new_keys = List.map (fun (t1, t2) -> gen_key t1 t2) periodes in
    new_keys

let load fs =
  let* () =
    if Fat.exists fs Cfg.key_dir then Ok ()
    else Fat.mkdir fs Cfg.key_dir |> unwrap_msg
  in
  let* l = Fat.ls fs Cfg.key_dir |> unwrap_msg in
  let l = List.map (fun entry -> entry.Fat.name) l in
  let l = List.map (Spath.add Cfg.key_dir) l in
  let l =
    List.filter (fun spath -> not @@ Spath.equal spath Cfg.sm_priv_key) l
  in
  let* keys = list_map (read_key fs) l in
  match keys with
  | [] -> Ok None
  | _l ->
      let* sm_priv, sm_pub = read_eddsa fs Cfg.sm_priv_key in
      let ht = Hashtbl.create 0xff in
      let () = List.iter (fun k -> Hashtbl.replace ht k.pub k) keys in
      Ok (Some { fs; sm_priv; sm_pub; ht })

let init fs =
  let* opt = load fs in
  let* t =
    match opt with
    | Some t -> Ok t
    | None ->
        let sm_priv, sm_pub = EddsaPrivateKey.generate () in
        Log.debug (fun m ->
            m "generated secmod key: `%s`" (EddsaPublicKey.to_b32 sm_pub));
        let ht = Hashtbl.create 0xff in
        let t = { fs; sm_priv; sm_pub; ht } in
        let* () = write_eddsa fs Cfg.sm_priv_key t.sm_priv in
        Ok t
  in
  let now = TimeAbsolute.of_ptime (Mirage_ptime.now ()) in
  let keys = List.of_seq @@ Hashtbl.to_seq_values t.ht in
  let new_keys = gen_additional_keys_until_lookahead ~now keys in
  List.iter (fun k -> Hashtbl.replace t.ht k.pub k) new_keys;
  let+ () = list_iter (write_key fs) new_keys in
  t

module Make (Fs : Fat.FS) = struct
  let t =
    match init Fs.t with
    | Error e -> Fmt.failwith "secmod_eddsa initialization failure: %s." e
    | Ok t -> t

  let find pub =
    Hashtbl.find_opt t.ht pub |> Option.to_result ~none:"key not found"

  let add t1 t2 =
    let k = gen_key t1 t2 in
    Hashtbl.replace t.ht k.pub k;
    let+ () = write_key t.fs k in
    ()

  let delete pub =
    let* k = find pub in
    Hashtbl.remove t.ht k.pub;
    delete_file t.fs (key_spath k)

  let _delete_outdated ~now =
    Hashtbl.to_seq_values t.ht
    |> List.of_seq
    |> List.filter (fun k -> TimeAbsolute.compare now k.t2 >= 0)
    |> List.map (fun k -> k.pub)
    |> list_iter delete

  (* ---- *)

  let sm_pub = t.sm_pub
  let sign_secmod s = EddsaSignature.sign ~key:t.sm_priv s

  let sign pub s =
    let+ k = find pub in
    let data = EddsaSignature.sign ~key:k.priv s in
    data

  (* delete and replace *)
  let revoke pub =
    let* k = find pub in
    let* () = delete pub in
    let* () = add k.t1 k.t2 in
    Ok ()

  let conv = fun { priv= _; pub; t1; t2 } -> (pub, (t1, t2))
  let keys () = Hashtbl.to_seq_values t.ht |> List.of_seq |> List.map conv
  let find_key pub = Hashtbl.find_opt t.ht pub |> Option.map conv
end