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_path = "/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;
  mutex: Miou.Mutex.t;
}

open struct
  let find_coin_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_coin_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 read_sm_key fs spath =
    let* s = Fat.read fs spath in
    Bbin.decode Eddsa.priv_bin s

  let write_sm_key fs spath t =
    let* s = Bbin.encode Eddsa.priv_bin t 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 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_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 epochs t0 tmax overlap duration =
    let rec go t acc =
      match t >= tmax with
      | true -> acc
      | false ->
          let t1 = TimeAbsolute.sub t overlap in
          let t2 = TimeAbsolute.add t duration in
          let acc = (t1, t2) :: acc in
          go t2 acc
    in
    go t0 []

  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 load fs =
    let open Fat in
    let open Cfg in
    let* () = if exists fs key_dir then Ok () else mkdir fs key_dir in
    let* l = ls fs key_dir in
    let l = List.map (fun entry -> Path.add key_dir entry.name) l in
    let* keys = list_map (read_key fs) l in
    match keys with
    | [] -> Ok None
    | _l ->
        let* sm_priv = read_sm_key fs sm_key_path in
        let sm_pub = Eddsa.pub_of_priv sm_priv in
        let ht = Hashtbl.create 0xff in
        List.iter (fun k -> Hashtbl.replace ht k.h_pub k) keys;
        let mutex = Miou.Mutex.create () in
        Ok (Some { fs; sm_priv; sm_pub; ht; mutex })

  let create 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_sm_key fs Cfg.sm_key_path sm_priv in
          let ht = Hashtbl.create 0xff in
          let mutex = Miou.Mutex.create () in
          Ok { fs; sm_priv; sm_pub; ht; mutex }
    in
    let keys = Hashtbl.to_seq_values t.ht |> List.of_seq in
    let new_keys =
      Coin.all_coins
      |> Iarray.to_list
      |> List.concat_map (fun coin ->
          let t0 =
            keys
            |> List.filter (fun k ->
                String.equal coin.Coin.section_name k.coin.section_name)
            |> List.fold_left (fun acc k -> TimeAbsolute.max acc k.t2) now
          in
          let tmax = TimeAbsolute.add t0 Cfg.lookahead in
          let l = epochs t0 tmax Cfg.overlap coin.duration_withdraw in
          List.map (fun (t1, t2) -> gen_key coin t1 t2) 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

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

  let conv k = (k.h_pub, (k.coin, k.pub, k.t1))
end

(* ### *)

let create fs =
  match create fs with
  | Error e ->
      Fmt.failwith "secmod_rsa initialization failure: %a." Result.pp_err e
  | Ok t -> t

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

let keys t =
  Miou.Mutex.protect t.mutex @@ fun () ->
  Hashtbl.to_seq_values t.ht |> List.of_seq |> List.map conv

let find_key t h_pub =
  Miou.Mutex.protect t.mutex @@ fun () ->
  Hashtbl.find_opt t.ht h_pub |> Option.map conv

let sign t h_pub s =
  let k = Miou.Mutex.protect t.mutex @@ fun () -> find_exn t h_pub in
  Rsa.sign ~key:k.priv s

let revoke t h_pub =
  Miou.Mutex.protect t.mutex @@ fun () ->
  let k = find_exn t h_pub in
  Hashtbl.remove t.ht h_pub;
  let* () = delete_file t.fs (key_spath k) in
  let k = gen_key k.coin k.t1 k.t2 in
  Hashtbl.replace t.ht k.h_pub k;
  let+ () = write_key t.fs k in
  ()