let src = Logs.Src.create "mte.secmod_rsa"

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

(* - *)
open Syntax
open Time
module DenominationHash = Hash.DenominationHash
module Coin = Config.Coin

module Cfg = struct
  include Config.Exchange_secmod_rsa

  let key_dir = "/RSA"
  let sm_key = "/SM_RSA"
end

type key = {
  coin: Coin.t;
  priv: Rsa.priv;
  pub: Rsa.pub;
  h_pub: DenominationHash.t;
  t1: TimeAbsolute.t;
  t2: TimeAbsolute.t;
}

type t = {
  fs: Fat.t;
  sm_priv: Eddsa.priv;
  sm_pub: Eddsa.pub;
  ht: (DenominationHash.t, key) Hashtbl.t;
}

let find_exn h_pub section_name =
  Coin.all_coins
  |> Iarray.find_opt (fun coin ->
      String.equal section_name coin.Coin.section_name)
  |> function
  | Some coin -> coin
  | None ->
      Fmt.failwith
        "Secmod_rsa denomination key loading failure on key `%a`: unknown \
         section_name [%s]"
        DenominationHash.pp h_pub section_name

(* ? enforce all to be of the same size instead *)
(* we use Bin.cstring + b32 binary encoding because
     rsa keysize is not known and then we have to escape '\x00' *)
let rsa_private_key_bin =
  let decode_exn o =
    Rsa.priv_of_crockford o |> function Error e -> invalid_arg e | Ok v -> v
  in
  let encode o = Rsa.priv_to_crockford o in
  Bin.map Bin.cstring decode_exn encode

let key_bin =
  let open Bin in
  record (fun t1 t2 section_name priv ->
      let pub = Rsa.pub_of_priv priv in
      let h_pub = DenominationHash.hash pub in
      let coin = find_exn h_pub section_name in
      { coin; t1; t2; priv; pub; h_pub })
  |+ field TimeAbsolute.bin (fun t -> t.t1)
  |+ field TimeAbsolute.bin (fun t -> t.t2)
  |+ field cstring (fun t -> t.coin.section_name)
  |+ field rsa_private_key_bin (fun t -> t.priv)
  |> sealr

let key_spath k =
  let sfn = String.sub (DenominationHash.to_crockford k.h_pub) 0 8 in
  Fat.Path.add Cfg.key_dir sfn

let read_eddsa fs spath =
  Log.debug (fun m -> m "reading key file `%s`" spath);
  let* data = Fat.read fs spath in
  let* priv = Eddsa.priv_of_octets data |> Result.map_error (fun e -> `Msg e) in
  let pub = Eddsa.pub_of_priv priv in
  Ok (priv, pub)

let write_eddsa fs spath priv =
  Log.debug (fun m -> m "writing key file `%s`" spath);
  let data = Eddsa.priv_to_octets priv in
  Fat.write fs spath data

let read_key fs spath =
  Log.debug (fun m -> m "reading key file `%s`" spath);
  let* s = Fat.read fs spath in
  let+ k = Bbin.decode key_bin s in
  k

let write_key fs k =
  let spath = key_spath k in
  Log.debug (fun m -> m "writing key file `%s`" spath);
  let* s = Bbin.encode key_bin k in
  Fat.write fs spath s

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

(* -- *)

let gen_key coin t1 t2 =
  let bits = coin.Coin.rsa_keysize in
  let priv, pub = Rsa.generate ~bits () in
  let h_pub = DenominationHash.hash pub in
  let k = { coin; priv; pub; h_pub; t1; t2 } in
  Log.debug (fun m ->
      m "generated key for coin [%s]: `%a`" coin.section_name
        DenominationHash.pp k.h_pub);
  k

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

let split_in_periodes coin ~start ~end_ =
  assert (start < end_);
  let duration_withdraw = coin.Coin.duration_withdraw in
  (* no overlap on first periode *)
  let t1 = start in
  let t2 = TimeAbsolute.add start duration_withdraw 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 duration_withdraw in
    if t2 > end_ then acc else go ((t1, t2) :: acc) t2 end_
  in
  go acc start end_

let gen_additional_keys_until_lookahead coin ~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 coin ~start ~end_ in
    let new_keys = List.map (fun (t1, t2) -> gen_key coin 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
  in
  let* l = Fat.ls fs Cfg.key_dir in
  let l = List.map (fun entry -> Fat.Path.add Cfg.key_dir entry.Fat.name) l in
  let l = List.filter (fun spath -> not @@ String.equal Cfg.sm_key spath) 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_key in
      let ht = Hashtbl.create 0xff in
      let () = List.iter (fun k -> Hashtbl.replace ht k.h_pub k) keys in
      Ok (Some { fs; sm_priv; sm_pub; ht })

let init fs =
  let* opt = load fs in
  let now = TimeAbsolute.of_ptime (Mirage_ptime.now ()) in
  let* t =
    match opt with
    | Some t -> Ok t
    | None ->
        let sm_priv, sm_pub = Eddsa.generate () in
        Log.debug (fun m -> m "generated secmod key: `%a`" Eddsa.pp_pub sm_pub);
        let* () = write_eddsa fs Cfg.sm_key sm_priv in
        let ht = Hashtbl.create 0xff in
        Ok { fs; sm_priv; sm_pub; ht }
  in
  let all_keys = List.of_seq @@ Hashtbl.to_seq_values t.ht in
  let new_keys_l =
    Coin.all_coins
    |> Iarray.to_list
    |> List.map (fun coin ->
        let keys =
          List.filter
            (fun k -> String.equal coin.Coin.section_name k.coin.section_name)
            all_keys
        in
        gen_additional_keys_until_lookahead coin ~now keys)
  in
  let new_keys = List.concat new_keys_l in
  List.iter (fun k -> Hashtbl.replace t.ht k.h_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_rsa initialization failure: %a." Result.pp_err e
    | Ok t -> t

  let find_exn h_pub =
    Log.debug (fun m -> m "find_exn: `%a`" DenominationHash.pp h_pub);
    match Hashtbl.find_opt t.ht h_pub with
    | Some v -> v
    | None -> Fmt.failwith "secmod_rsa operation on unknown key"

  let delete h_pub =
    let k = find_exn h_pub in
    Hashtbl.remove t.ht h_pub;
    delete_file t.fs (key_spath k)

  let _delete_outdated ~now =
    Hashtbl.to_seq t.ht
    |> List.of_seq
    |> List.filter (fun (_h_pub, k) -> TimeAbsolute.compare now k.t2 >= 0)
    |> List.map (fun (h_pub, _k) -> h_pub)
    |> list_iter delete

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

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

  let sign h_pub msg =
    let k = find_exn h_pub in
    Rsa.sign ~key:k.priv msg

  let revoke h_pub =
    Log.debug (fun m -> m "revoke `%a`" DenominationHash.pp h_pub);
    let k = find_exn h_pub in
    let* () = delete h_pub in
    let* () = add k.coin k.t1 k.t2 in
    Ok ()

  let conv k = (k.h_pub, (k.coin, k.pub, k.t1))
  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