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

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

(* - *)
open Syntax
open Crypto
open Time

module Cfg = struct
  open Config
  include Exchange_secmod_rsa

  let sections = Coin.all_coins |> List.map (fun coin -> coin.Coin.section_name)

  (* helper functions to get config value from section_name: *)

  let duration_withdraw =
    let duration_withdraw_assoc =
      Coin.all_coins
      |> List.map (fun coin -> (coin.Coin.section_name, coin.duration_withdraw))
    in
    fun ~section_name ->
      match List.assoc_opt section_name duration_withdraw_assoc with
      | None ->
          Fmt.failwith "section_name `%s` not found in config" section_name
      | Some v -> v

  let rsa_keysize =
    let rsa_keysize_assoc =
      Coin.all_coins
      |> List.map (fun coin -> (coin.Coin.section_name, coin.rsa_keysize))
    in
    fun ~section_name ->
      match List.assoc_opt section_name rsa_keysize_assoc with
      | None ->
          Fmt.failwith "section_name `%s` not found in config" section_name
      | Some v -> v
end

type key = {
  section_name: string;
  priv: RsaPrivateKey.t;
  pub: RsaPublicKey.t;
  h_pub: DenominationHash.t;
  t1: TimeAbsolute.t;
  t2: TimeAbsolute.t;
}

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

let parse_filename =
  let scan_filename s =
    Scanf.sscanf_opt s "%Lu-%Lu" (fun t1 t2 ->
        (TimeAbsolute.of_s t1, TimeAbsolute.of_s t2))
  in
  fun fpath -> scan_filename (Fpath.filename fpath)

let pp_filename =
  let to_int64 abs =
    abs |> Timestamp.of_absolute |> Timestamp.to_s |> function
    | None ->
        (* (= `never`) this should not happen given resonable config value *)
        Fmt.failwith "encountered timestamp with value `never`"
    | Some i -> i
  in
  fun ppf (t1, t2) -> Fmt.pf ppf "%Lu-%Lu" (to_int64 t1) (to_int64 t2)

let key_fpath k =
  let fname = Fmt.str "%a" pp_filename (k.t1, k.t2) in
  Fpath.(v Cfg.key_dir / k.section_name / fname)

(* -- IO -- *)

let read_eddsa fs fpath =
  Log.debug (fun m -> m "reading key file `%a`" Fpath.pp fpath);
  let fpath = Fpath.to_string fpath in
  let* data = Fat.read fs fpath |> unwrap_msg in
  EddsaPrivateKey.of_octets data

let read_rsa fs fpath =
  Log.debug (fun m -> m "reading key file `%a`" Fpath.pp fpath);
  let fpath = Fpath.to_string fpath in
  let* data = Fat.read fs fpath |> unwrap_msg in
  RsaPrivateKey.of_octets data

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

let write_rsa fs fpath priv =
  let data = RsaPrivateKey.to_octets priv in
  let fpath = Fpath.to_string fpath in
  Fat.write fs fpath data |> unwrap_msg

let write_key fs k = write_rsa fs (key_fpath k) k.priv

let delete_file fs fpath =
  (* check fpath just to be safe *)
  let () =
    let root = Fpath.v Cfg.key_dir in
    if not @@ Fpath.is_rooted ~root fpath then
      Fmt.failwith
        "delete_file failure: file `%a` is not contained in secmod directory"
        Fpath.pp fpath
  in
  Log.debug (fun m -> m "delete key file `%a`" Fpath.pp fpath);
  let fpath = Fpath.to_string fpath in
  let+ () = Fat.remove fs fpath |> unwrap_msg in
  ()

let get_key_dir_contents fs dir_fpath =
  let dir_fpath = Fpath.to_string dir_fpath in
  let* () = Fat.mkdir fs dir_fpath |> unwrap_msg in
  Log.info (fun m -> m "created directory `%s`" dir_fpath);
  let+ l = Fat.ls fs dir_fpath |> unwrap_msg in
  (*   List.map Fpath.normalize l *)
  let l = List.map (fun entry -> entry.Fat.name) l in
  let l = List.map Fpath.v l in
  l

(* -- *)

let gen_key ~section_name t1 t2 =
  let bits = Cfg.rsa_keysize ~section_name in
  let priv, pub = RsaPrivateKey.generate ~bits () in
  let h_pub = DenominationHash.hash_of_rsa pub in
  Log.debug (fun m ->
      m "generated key (%a):@,`%s`" pp_filename (t1, t2)
        (DenominationHash.to_octets h_pub |> B32.encode));
  { section_name; priv; pub; h_pub; t1; t2 }

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

let split_in_periodes ~start ~end_ ~duration_withdraw =
  assert (start < end_);
  (* 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 ~now ~section_name l =
  (* try to _not_ generate keys with validity start in the past
     (probably not important) *)
  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 duration_withdraw = Cfg.duration_withdraw ~section_name in
    let periodes = split_in_periodes ~start ~end_ ~duration_withdraw in
    let new_keys =
      List.map (fun (t1, t2) -> gen_key ~section_name t1 t2) periodes
    in
    new_keys

let sm_key_fpath =
  Result.get_ok
  @@
  let+ fpath = Fpath.of_string Cfg.sm_priv_key |> unwrap_msg in
  Fpath.normalize fpath

(* we load sm_key separately
   we don't accept non-key files in key_dir *)
let load_key fs ~section_name fpath =
  match parse_filename fpath with
  | None -> Fmt.error "invalid file `%a`" Fpath.pp fpath
  | Some (t1, t2) ->
      let+ priv = read_rsa fs fpath in
      let pub = RsaPrivateKey.pub_of_priv priv in
      let h_pub = DenominationHash.hash_of_rsa pub in
      { section_name; priv; pub; h_pub; t1; t2 }

let load_section fs section_name =
  let section_fpath = Fpath.(v Cfg.key_dir / section_name) in
  let* l = get_key_dir_contents fs section_fpath in
  let l = List.filter (fun fpath -> not @@ Fpath.equal fpath sm_key_fpath) l in
  let* keys = list_map (load_key fs ~section_name) l in
  Ok keys

let load fs =
  let* keys_l = list_map (load_section fs) Cfg.sections in
  let keys = List.concat keys_l in
  match keys with
  | [] -> Ok None
  | _l ->
      let* sm_key_priv = read_eddsa fs sm_key_fpath in
      let sm_pub = EddsaPrivateKey.pub_of_priv sm_key_priv 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_key_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_key_priv, sm_pub = EddsaPrivateKey.generate () in
        Log.debug (fun m ->
            m "generated secmod key: `%s`" (EddsaPublicKey.to_b32 sm_pub));
        let* () = write_eddsa fs sm_key_fpath sm_key_priv in
        let ht = Hashtbl.create 0xff in
        Ok { fs; sm_key_priv; sm_pub; ht }
  in
  let all_keys = List.of_seq @@ Hashtbl.to_seq_values t.ht in
  let new_keys_l =
    List.map
      (fun section_name ->
        let keys =
          List.filter (fun k -> k.section_name = section_name) all_keys
        in
        gen_additional_keys_until_lookahead ~now ~section_name keys)
      Cfg.sections
  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: %s." e
    | Ok t -> t

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

  let delete h_pub =
    let* k = find h_pub in
    Hashtbl.remove t.ht h_pub;
    delete_file t.fs (key_fpath 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 section_name t1 t2 =
    let k = gen_key ~section_name 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 = EddsaSignature.sign ~key:t.sm_key_priv s

  let sign h_pub msg =
    let+ k = find h_pub in
    let data = RsaSignature.sign ~key:k.priv msg in
    data

  let revoke h_pub =
    Log.debug (fun m ->
        m "revoke `%s`" (DenominationHash.to_octets h_pub |> B32.encode));
    let* k = find h_pub in
    let* () = delete h_pub in
    let* () = add k.section_name k.t1 k.t2 in
    Ok ()

  let conv =
   fun { section_name; priv= _; pub; h_pub; t1; t2= _ } ->
    (h_pub, (section_name, pub, 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