(* TODO refacto common parts with secmod_eddsa *)
open Syntax
open Crypto
open Time
module Cfg = Config.Exchange_secmod_rsa

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

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

(* -- util -- *)

let time_abs_of_string s =
  int_of_string_opt s |> Option.map (fun n -> Absolute.of_s (Int64.of_int n))

let t1_t2_of_fpath fpath =
  let fname = Fpath.filename fpath in
  match String.split_on_char '-' fname with
  | [] -> Fmt.failwith "not possible"
  | [ t1; t2 ] -> (
      match (time_abs_of_string t1, time_abs_of_string t2) with
      | None, _ | _, None -> None
      | Some t1, Some t2 -> Some (t1, t2))
  | _ -> None

let time_abs_to_string abs =
  abs
  |> Timestamp.of_absolute
  |> Timestamp.to_s
  |> Option.get
  |> Int64.to_int
  |> string_of_int

let key_fpath k =
  let t1 = time_abs_to_string k.t1 in
  let t2 = time_abs_to_string k.t2 in
  let fname = Fmt.str "%s-%s" t1 t2 in
  Fpath.(v Cfg.key_dir / k.section_name / fname)

(* -- IO -- *)

let read_eddsa fpath =
  let* data = Bos.OS.File.read fpath |> unwrap_err_msg in
  EddsaPrivateKey.of_octets data

let read_rsa fpath =
  let* data = Bos.OS.File.read fpath |> unwrap_err_msg in
  RsaPrivateKey.of_octets data

let write_eddsa fpath priv =
  let data = EddsaPrivateKey.to_octets priv in
  Bos.OS.File.write fpath data |> unwrap_err_msg

let write_rsa fpath priv =
  let data = RsaPrivateKey.to_octets priv in
  Bos.OS.File.write fpath data |> unwrap_err_msg

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

let delete_key_file k =
  let+ () =
    Bos.OS.File.delete ~must_exist:true (key_fpath k) |> unwrap_err_msg
  in
  ()

let get_key_dir_contents dir_fpath =
  let* b = Bos.OS.Dir.create ~mode:0o700 dir_fpath |> unwrap_err_msg in
  if b then
    Logs.info (fun m ->
        m "secmod_rsa: created directory `%a`" Fpath.pp dir_fpath);
  let+ l =
    Bos.OS.Dir.contents ~dotfiles:false ~rel:true dir_fpath |> unwrap_err_msg
  in
  l

(* -- *)

let sort_keys l = List.sort (fun a b -> Absolute.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 = Absolute.add start Cfg.duration in
  let acc = [ (t1, t2) ] in
  let start = t2 in
  let rec go acc start end_ =
    let t1 = Absolute.sub start Cfg.overlap_duration in
    let t2 = Absolute.add start Cfg.duration 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 :: _ -> Absolute.sub hd.t2 Cfg.overlap_duration
  in
  let end_ = Absolute.add now Cfg.lookahead_sign in
  let periodes = split_in_periodes ~start ~end_ in
  let new_keys =
    List.map
      (fun (t1, t2) ->
        let priv, pub = RsaPrivateKey.generate ~bits:Cfg.rsa_keysize () in
        { section_name; priv; pub; t1; t2 })
      periodes
  in
  new_keys

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

(* we load sm_key separately
   we don't accept non-key files in key_dir *)
let load_key ~section_name fpath =
  match Fpath.equal (Fpath.normalize fpath) sm_key_fpath with
  | true -> Ok None
  | false -> (
      match t1_t2_of_fpath fpath with
      | None -> Fmt.error "invalid file `%a`" Fpath.pp fpath
      | Some (t1, t2) ->
          let* priv = read_rsa fpath in
          let pub = RsaPrivateKey.pub_of_priv priv in
          Ok (Some { section_name; priv; pub; t1; t2 }))

let load_section section_name =
  let section_fpath = Fpath.(v Cfg.key_dir / section_name) in
  let* l = get_key_dir_contents section_fpath in
  let* l = list_map (load_key ~section_name) l in
  let keys = List.filter_map Fun.id l in
  Ok keys

let load () =
  let* keys_l = list_map load_section Cfg.sections in
  let keys = List.concat keys_l in
  match keys with
  | [] -> Ok None
  | _l ->
      let* sm_key_priv = read_eddsa 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.pub k) keys in
      Ok (Some { sm_key_priv; sm_pub; ht })

let init () =
  let* opt = load () in
  let* t =
    match opt with
    | Some t -> Ok t
    | None ->
        let sm_key_priv, sm_pub = EddsaPrivateKey.generate () in
        let* () = write_eddsa sm_key_fpath sm_key_priv in
        let ht = Hashtbl.create 0xff in
        Ok { sm_key_priv; sm_pub; ht }
  in
  let now = Absolute.of_ptime (Ptime_clock.now ()) 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
  let () = List.iter (fun k -> Hashtbl.replace t.ht k.pub k) new_keys in
  let+ () = list_iter write_key new_keys in
  t

type info = string * rsa_pub * Time.Absolute.t * Time.Absolute.t

module Make () = struct
  type pub = RsaPublicKey.t
  type info = string * RsaPublicKey.t * Time.Absolute.t * Time.Absolute.t

  (* --- *)
  let t =
    match init () with
    | Error e -> Fmt.failwith "secmod_rsa initialization failure: %s." e
    | Ok t -> t

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

  let delete pub =
    let* k = find pub in
    Hashtbl.remove t.ht k.pub; delete_key_file k

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

  let add section_name t1 t2 =
    let priv, pub = RsaPrivateKey.generate ~bits:Cfg.rsa_keysize () in
    let k = { section_name; priv; pub; t1; t2 } in
    Hashtbl.replace t.ht k.pub k;
    ()

  let sm_pub = t.sm_pub

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

  let sign_secmod s = EddsaSignature.sign ~key:t.sm_key_priv s

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

  let revoke pub =
    let* k = find pub in
    let* () = delete pub in
    add k.section_name k.t1 k.t2;
    Ok ()
end