(* TODO
   - key rotation
   - how many signkey to use?
     we just use 1 for now
   - does the secmod's own key as metadata/expiration date?
   - something to refer to valid sk/dn

   - eddsa.ml with phantom type for key-kind + signed-data-kind *)
open Syntax
open Crypto

type signkey = {
  priv: eddsa_priv;
  sk_data: Signkey_data.t;
}

type denom = {
  priv: rsa_priv;
  dn_data: Denom_data.t;
}

type t = {
  lock: Miou.Mutex.t;
  sm_key_priv: eddsa_priv;
  sm_key_pub: eddsa_pub;
  sk_ht: (eddsa_pub, signkey) Hashtbl.t;
  sk_revoked_ht:
    (eddsa_pub, signkey * Bin_sig.MasterSigningKeyRevocation.t) Hashtbl.t;
  dn_ht: (denomination_hash, denom) Hashtbl.t;
  dn_revoked_ht:
    ( denomination_hash,
      denom * Bin_sig.MasterDenominationKeyRevocation.t )
    Hashtbl.t;
}

(* note: don't expose a signing function if we want a real "security module" one day *)
let sign_with_sm_key t s = EddsaSignature.sign ~key:t.sm_key_priv s
let verify_with_sm_key t s ~msg = EddsaSignature.verify ~key:t.sm_key_pub s ~msg

let verify_with_master_key s ~msg =
  EddsaSignature.verify ~key:Config.master_public_key s ~msg

(* TODO
   - do something to force `pub` to be one of the valid signkey
     how to handle revocation?
   raise exn for now *)
let sign_with_signkey t ~pub s =
  Miou.Mutex.protect t.lock @@ fun () ->
  match Hashtbl.find_opt t.sk_ht pub with
  | None -> Fmt.failwith "secmod failure: public key not found."
  | Some signkey ->
      let v = EddsaSignature.sign ~key:signkey.priv s in
      v

let verify_with_signkey t ~pub s ~msg =
  Miou.Mutex.protect t.lock @@ fun () ->
  match Hashtbl.find_opt t.sk_ht pub with
  | None -> Error "secmod failure: public key not found."
  | Some signkey -> EddsaSignature.verify ~key:signkey.sk_data.pub s ~msg

let get_sm_key_pub t = t.sm_key_pub

let get_signkeys t =
  Miou.Mutex.protect t.lock @@ fun () ->
  Hashtbl.to_seq_values t.sk_ht |> List.of_seq

let get_denoms t =
  Miou.Mutex.protect t.lock @@ fun () ->
  Hashtbl.to_seq_values t.dn_ht |> List.of_seq

let find_signkey t pub =
  Miou.Mutex.protect t.lock @@ fun () -> Hashtbl.find_opt t.sk_ht pub

let find_denom t h_denom =
  Miou.Mutex.protect t.lock @@ fun () -> Hashtbl.find_opt t.dn_ht h_denom

let get_signkeys_data t = get_signkeys t |> List.map (fun v -> v.sk_data)
let get_denoms_data t = get_denoms t |> List.map (fun v -> v.dn_data)

let find_signkey_data t pub =
  find_signkey t pub |> Option.map (fun v -> v.sk_data)

let find_denom_data t h_denom =
  find_denom t h_denom |> Option.map (fun v -> v.dn_data)

let add_signkey_master_signatures conn t l =
  Miou.Mutex.protect t.lock @@ fun () ->
  list_iter
    (fun (pub, master_sig) ->
      match Hashtbl.find_opt t.sk_ht pub with
      | None -> Error "secmod failure: public key not found."
      | Some signkey ->
          let sk_data = { signkey.sk_data with master_sig= Some master_sig } in
          let signkey = { signkey with sk_data } in
          let* () =
            Pg.activate_signing_key conn ~master_sig sk_data |> unwrap_err_caqti
          in
          Hashtbl.replace t.sk_ht pub signkey;
          Ok ())
    l

let add_denom_master_signatures conn t l =
  Miou.Mutex.protect t.lock @@ fun () ->
  list_iter
    (fun (h_denom_pub, master_sig) ->
      match Hashtbl.find_opt t.dn_ht h_denom_pub with
      | None -> Error "secmod failure: denomination hash not found."
      | Some denom ->
          let dn_data = { denom.dn_data with master_sig= Some master_sig } in
          let denom = { denom with dn_data } in
          let* () =
            Pg.add_denomination_key conn ~master_sig dn_data |> unwrap_err_caqti
          in
          Hashtbl.replace t.dn_ht h_denom_pub denom;
          Ok ())
    l

let revoke_signkey t exchange_pub master_sig =
  Miou.Mutex.protect t.lock @@ fun () ->
  match Hashtbl.find_opt t.sk_ht exchange_pub with
  | None -> Error "secmod failure: denomination hash not found."
  | Some signkey ->
      Hashtbl.remove t.sk_ht exchange_pub;
      Hashtbl.replace t.sk_revoked_ht exchange_pub (signkey, master_sig);
      Ok ()

let revoke_denomination t h_denom_pub master_sig =
  Miou.Mutex.protect t.lock @@ fun () ->
  match Hashtbl.find_opt t.dn_ht h_denom_pub with
  | None -> Error "secmod failure: denomination hash not found."
  | Some denom ->
      Hashtbl.remove t.dn_ht h_denom_pub;
      Hashtbl.replace t.dn_revoked_ht h_denom_pub (denom, master_sig);
      Ok ()

let dir = Fpath.(v "data" / "secmod ")

let db_lookup_signkey_data conn fname priv =
  let pub = EddsaPrivateKey.pub_of_priv priv in
  let* opt = Pg.lookup_signing_key conn pub in
  match opt with
  | None ->
      Fmt.error_msg
        "load_signkey error no associated metadata found in database for \
         signkey `%s`."
        (Fpath.to_string fname)
  | Some (stamp_start, stamp_expire, stamp_end) ->
      (* TODO master_sig *)
      let master_sig = None in
      let v =
        Signkey_data.{ pub; stamp_start; stamp_expire; stamp_end; master_sig }
      in
      Ok v

let db_lookup_denom_data conn ~section_name priv =
  let pub = RsaPrivateKey.pub_of_priv priv in
  let h_pub = Bin_type.DenominationHash.hash (RsaPublicKey.to_octets pub) in
  let* opt = Pg.lookup_denomination_key conn h_pub in
  match opt with
  | None ->
      Fmt.error_msg
        "load_denom error no associated metadata found in database for \
         denomination `%s`."
        section_name
  | Some
      ( stamp_start,
        stamp_expire_withdraw,
        stamp_expire_deposit,
        stamp_expire_legal,
        value,
        fee_withdraw,
        fee_deposit,
        fee_refresh,
        fee_refund,
        age_mask ) ->
      let master_sig = None in
      let v =
        Denom_data.
          {
            pub;
            section_name;
            value;
            stamp_start;
            stamp_expire_withdraw;
            stamp_expire_deposit;
            stamp_expire_legal;
            fee_withdraw;
            fee_deposit;
            fee_refresh;
            fee_refund;
            age_mask;
            h_pub;
            master_sig;
          }
      in
      Ok v

let load_signkey conn fname =
  let* opt = Data_file.read_eddsa fname in
  match opt with
  | None -> Ok None
  | Some priv ->
      let* sk_data = db_lookup_signkey_data conn fname priv in
      let signkey = { priv; sk_data } in
      Ok (Some signkey)

let _store t =
  let* () = Data_file.write_eddsa Fpath.(dir / "sm_key") t.sm_key_priv in
  let* () =
    get_signkeys t
    |> List.mapi (fun i (key : signkey) ->
        let fname = Fpath.(dir / string_of_int i) in
        (fname, key.priv))
    |> list_iter (fun (fname, key) -> Data_file.write_eddsa fname key)
  in
  let* () =
    get_denoms t
    |> List.mapi (fun i (key : denom) ->
        let fname = Fpath.(dir / string_of_int i) in
        (fname, key.priv))
    |> list_iter (fun (fname, key) -> Data_file.write_rsa fname key)
  in
  Ok ()

let load conn =
  let error_invalid_state =
    Fmt.error_msg "secmod load error: invalid store state."
  in
  let* sm_key_priv = Data_file.read_eddsa Fpath.(dir / "sm_key") in
  let* signkeys =
    let l = List.init 1 (fun i -> Fpath.(dir / string_of_int i)) in
    let* l = list_map (fun fname -> load_signkey conn fname) l in
    match opt_list l with Error () -> error_invalid_state | Ok opt -> Ok opt
  in
  let* denoms =
    let* l =
      let open Config.Coin in
      list_map
        (fun coin ->
          let section_name = coin.section_name in
          let fname = Fpath.(dir / section_name) in
          let* opt = Data_file.read_rsa fname in
          match opt with
          | None -> Ok None
          | Some priv ->
              (* todo: could check that coin config match db values *)
              let* dn_data = db_lookup_denom_data conn ~section_name priv in
              let denom = { priv; dn_data } in
              Ok (Some denom))
        all_coins
    in
    match Syntax.opt_list l with
    | Error () -> error_invalid_state
    | Ok opt -> Ok opt
  in
  match (sm_key_priv, signkeys, denoms) with
  | None, None, None -> Ok None
  | Some sm_key_priv, Some signkeys, Some denoms ->
      let sm_key_pub = EddsaPrivateKey.pub_of_priv sm_key_priv in
      let lock = Miou.Mutex.create () in
      let sk_ht =
        signkeys
        |> List.map (fun v -> (v.sk_data.pub, v))
        |> List.to_seq
        |> Hashtbl.of_seq
      in
      let dn_ht =
        denoms
        |> List.map (fun v -> (v.dn_data.h_pub, v))
        |> List.to_seq
        |> Hashtbl.of_seq
      in
      let sk_revoked_ht = Hashtbl.create 0xff in
      let dn_revoked_ht = Hashtbl.create 0xff in
      Ok
        (Some
           {
             lock;
             sm_key_priv;
             sm_key_pub;
             sk_ht;
             sk_revoked_ht;
             dn_ht;
             dn_revoked_ht;
           })
  | _, _, _ -> error_invalid_state

let make_new_signkey () =
  let stamp_start = Ptime_clock.now () |> Option.some in
  let stamp_expire =
    Timestamp.add_span_exn stamp_start
      (Some Config.Exchange.signkey_legal_duration)
  in
  let stamp_end = stamp_expire in
  let priv, pub = Mirage_crypto_ec.Ed25519.generate () in
  let master_sig = None in
  let sk_data =
    Signkey_data.{ pub; stamp_start; stamp_expire; stamp_end; master_sig }
  in
  { priv; sk_data }

let make_new_denom
    Config.Coin.
      {
        section_name;
        value;
        duration_withdraw;
        duration_spend;
        duration_legal;
        fee_withdraw;
        fee_deposit;
        fee_refresh;
        fee_refund;
        cipher;
        rsa_keysize;
        age_restricted= _;
      } =
  assert (cipher = `RSA);
  let stamp_start = Ptime_clock.now () |> Option.some in
  let stamp_expire_withdraw =
    Timestamp.add_span_exn stamp_start (Some duration_withdraw)
  in
  let stamp_expire_deposit =
    Timestamp.add_span_exn stamp_start (Some duration_spend)
  in
  let stamp_expire_legal =
    Timestamp.add_span_exn stamp_start (Some duration_legal)
  in
  let priv, pub = RsaPrivateKey.generate ~bits:rsa_keysize () in
  let h_pub = Bin_type.DenominationHash.hash (RsaPublicKey.to_octets pub) in
  let master_sig = None in
  let dn_data =
    Denom_data.
      {
        pub;
        section_name;
        value;
        stamp_start;
        stamp_expire_withdraw;
        stamp_expire_deposit;
        stamp_expire_legal;
        fee_withdraw;
        fee_deposit;
        fee_refresh;
        fee_refund;
        age_mask= 0;
        h_pub;
        master_sig;
      }
  in
  { priv; dn_data }

let make_new () =
  let lock = Miou.Mutex.create () in
  let sm_key_priv, sm_key_pub = Mirage_crypto_ec.Ed25519.generate () in
  let sk_ht =
    [ make_new_signkey () ]
    |> List.map (fun v -> (v.sk_data.pub, v))
    |> List.to_seq
    |> Hashtbl.of_seq
  in
  let dn_ht =
    Config.Coin.all_coins
    |> List.map make_new_denom
    |> List.map (fun v -> (v.dn_data.h_pub, v))
    |> List.to_seq
    |> Hashtbl.of_seq
  in
  let sk_revoked_ht = Hashtbl.create 0xff in
  let dn_revoked_ht = Hashtbl.create 0xff in
  { lock; sm_key_priv; sm_key_pub; sk_ht; sk_revoked_ht; dn_ht; dn_revoked_ht }

let init conn =
  match load conn with
  | Ok None ->
      let t = make_new () in
      Logs.info (fun m -> m "secmod initialized with fresh keys");
      t
  | Ok (Some v) ->
      Logs.info (fun m -> m "secmod initialized from storage");
      v
  | Error _ ->
      (* TODO error: pretty print *)
      Fmt.failwith "secmod init failure."