module type S = Keys_intf.S

module Make (Conn : Pg.CONN) = struct
  open Syntax
  open Crypto
  module DenominationHash = Hash.DenominationHash

  let read fname = Bos.OS.File.read fname |> unwrap_err_msg
  let write fname s = Bos.OS.File.write fname s |> unwrap_err_msg
  let write_eddsa fname priv = write fname (EddsaPrivateKey.to_octets priv)
  let write_rsa fname priv = write fname (RsaPrivateKey.to_octets priv)

  let read_eddsa fname =
    let* data = read fname in
    EddsaPrivateKey.of_octets data

  let read_rsa fname =
    let* data = read fname in
    RsaPrivateKey.of_octets data

  type sk = Signkey.t
  type future_sk = Api.FutureSignKey.t
  type dn = Denomination.t
  type future_dn = Api.FutureDenom.t

  (* TODO ! use lock *)
  (* not sure what to do with coin section_name, rm if possible *)
  type t = {
    sm_key: eddsa_priv;
    sm_pubkey: eddsa_pub;
    sk_ht: (eddsa_pub, sk) Hashtbl.t;
    dn_ht: (denom_hash, dn) Hashtbl.t;
    sk_key_ht: (eddsa_pub, eddsa_priv) Hashtbl.t;
    dn_key_ht: (denom_hash, rsa_priv) Hashtbl.t;
    future_sk_ht: (eddsa_pub, future_sk) Hashtbl.t;
    future_dn_ht: (denom_hash, future_dn) Hashtbl.t;
    future_sk_key_ht: (eddsa_pub, eddsa_priv) Hashtbl.t;
    future_dn_key_ht: (denom_hash, rsa_priv) Hashtbl.t;
    dn_section_name_ht: (denom_hash, string) Hashtbl.t;
  }

  let conn = (module Conn : Pg.CONN)
  let sm_key_fname = Fpath.(v Config.Exchange_secmod_eddsa.sm_priv_key)

  let sk_fname i =
    Fpath.(v Config.Exchange_secmod_eddsa.key_dir / Fmt.str "sk_%d" i)

  let dn_fname section_name =
    Fpath.(
      v Config.Exchange_secmod_eddsa.key_dir / Fmt.str "dn_%s" section_name)

  let sign_with_sm_key t s = EddsaSignature.sign ~key:t.sm_key s

  let make_future_sk t =
    let start = Time.Absolute.of_ptime (Ptime_clock.now ()) in
    let expire =
      Time.Absolute.add start Config.Exchange.signkey_legal_duration
    in
    let stamp_start = Timestamp.of_absolute start in
    let stamp_expire = Timestamp.of_absolute expire in
    let stamp_end = stamp_expire in
    let priv, pub = Mirage_crypto_ec.Ed25519.generate () in
    let signkey_secmod_sig =
      let open Signatures.SigningKeyAnnouncement in
      let exchange_pub = pub in
      let anchor_time = stamp_start in
      let duration = Timestamp.diff stamp_start stamp_expire in
      sign_f ~f:(sign_with_sm_key t) { exchange_pub; anchor_time; duration }
    in
    let future_sk =
      Api.FutureSignKey.
        { key= pub; stamp_start; stamp_expire; stamp_end; signkey_secmod_sig }
    in
    Hashtbl.replace t.future_sk_ht pub future_sk;
    Hashtbl.replace t.future_sk_key_ht pub priv;
    ()

  let make_future_dn t
      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 start = Time.Absolute.of_ptime (Ptime_clock.now ()) in
    let stamp_start = Timestamp.of_absolute start in
    let stamp_expire_withdraw =
      Timestamp.of_absolute @@ Time.Absolute.add start duration_withdraw
    in
    let stamp_expire_deposit =
      Timestamp.of_absolute @@ Time.Absolute.add start duration_spend
    in
    let stamp_expire_legal =
      Timestamp.of_absolute @@ Time.Absolute.add start duration_legal
    in
    let priv, pub = RsaPrivateKey.generate ~bits:rsa_keysize () in
    let open Api in
    let rsa_denomination_key =
      RsaDenominationKey.{ age_mask= 0; rsa_pub= pub }
    in
    let denom_pub = DenominationKey.Rsa rsa_denomination_key in
    let h_pub = DenominationHash.hash (RsaPublicKey.to_octets pub) in
    let denom_secmod_sig =
      let open Signatures.DenominationKeyAnnouncement in
      let h_denom_pub = h_pub in
      let h_section_name = Hash.Cstring.H64.hash section_name in
      let anchor_time = stamp_start in
      let duration_withdraw =
        Timestamp.diff stamp_start stamp_expire_withdraw
      in
      sign_f ~f:(sign_with_sm_key t)
        { h_denom_pub; h_section_name; anchor_time; duration_withdraw }
    in
    let future_dn =
      FutureDenom.
        {
          section_name;
          value;
          stamp_start;
          stamp_expire_withdraw;
          stamp_expire_deposit;
          stamp_expire_legal;
          denom_pub;
          fee_withdraw;
          fee_deposit;
          fee_refresh;
          fee_refund;
          denom_secmod_sig;
        }
    in
    Hashtbl.replace t.future_dn_ht h_pub future_dn;
    Hashtbl.replace t.future_dn_key_ht h_pub priv;
    ()

  let make_new () =
    let sm_key, sm_pubkey = Mirage_crypto_ec.Ed25519.generate () in
    let t =
      {
        sm_key;
        sm_pubkey;
        sk_ht= Hashtbl.create 0xff;
        dn_ht= Hashtbl.create 0xff;
        sk_key_ht= Hashtbl.create 0xff;
        dn_key_ht= Hashtbl.create 0xff;
        future_sk_ht= Hashtbl.create 0xff;
        future_dn_ht= Hashtbl.create 0xff;
        future_sk_key_ht= Hashtbl.create 0xff;
        future_dn_key_ht= Hashtbl.create 0xff;
        dn_section_name_ht= Hashtbl.create 0xff;
      }
    in
    make_future_sk t;
    List.iter (make_future_dn t) Config.Coin.all_coins;
    t

  let database_find_sk conn pub =
    let* opt = Pg.find_signkey conn pub |> unwrap_err_caqti in
    match opt with
    | None -> Fmt.error "Keys: signkey data not found in database"
    | Some sk_data -> Ok sk_data

  let database_find_dn conn h_pub =
    let* opt = Pg.find_denom conn h_pub |> unwrap_err_caqti in
    match opt with
    | None -> Fmt.error "Keys: denomination data not found in database"
    | Some dn_data -> Ok dn_data

  let find_signkey pub = Pg.find_signkey conn pub |> unwrap_err_caqti
  let find_denom h_pub = Pg.find_denom conn h_pub |> unwrap_err_caqti

  (* TODO check stamps definitions *)
  let make_future_sk ~pub ~start ~expire =
    let stamp_start = Timestamp.of_absolute start in
    let stamp_expire = Timestamp.of_absolute expire in
    let stamp_end = stamp_expire in
    let signkey_secmod_sig =
      let open Signatures.SigningKeyAnnouncement in
      let exchange_pub = pub in
      let anchor_time = stamp_start in
      let duration = Timestamp.diff stamp_start stamp_expire in
      sign_f ~f:Secmod_eddsa.sign_with_sm_key
        { exchange_pub; anchor_time; duration }
    in
    Api.FutureSignKey.
      { key= pub; stamp_start; stamp_expire; stamp_end; signkey_secmod_sig }

  let _get_future_signkeys () =
    Secmod_eddsa.keys ()
    |> List.map (fun (pub, t1, t2) ->
        let* opt = find_signkey pub in
        match opt with
        | None ->
            let future_sk = make_future_sk ~pub ~start:t1 ~expire:t2 in
            Ok (Some future_sk)
        | Some sk -> (
            match Timestamp.of_absolute t1 = sk.stamp_start with
            | false ->
                Fmt.error
                  "secmod/database stamp_start mismatch for signkey `%s`"
                  (EddsaPublicKey.to_b32 pub)
            | true -> Ok None))

  let make_future_dn ~coin ~pub ~start =
    let Config.Coin.
          {
            section_name;
            value;
            duration_withdraw;
            duration_spend;
            duration_legal;
            fee_withdraw;
            fee_deposit;
            fee_refresh;
            fee_refund;
            cipher= _;
            rsa_keysize= _;
            age_restricted= _;
          } =
      coin
    in
    let stamp_start = Timestamp.of_absolute start in
    let stamp_expire_withdraw =
      Timestamp.of_absolute @@ Time.Absolute.add start duration_withdraw
    in
    let stamp_expire_deposit =
      Timestamp.of_absolute @@ Time.Absolute.add start duration_spend
    in
    let stamp_expire_legal =
      Timestamp.of_absolute @@ Time.Absolute.add start duration_legal
    in
    let open Api in
    let rsa_denomination_key =
      RsaDenominationKey.{ age_mask= 0; rsa_pub= pub }
    in
    let denom_pub = DenominationKey.Rsa rsa_denomination_key in
    let h_pub = DenominationHash.hash (RsaPublicKey.to_octets pub) in
    let denom_secmod_sig =
      let open Signatures.DenominationKeyAnnouncement in
      let h_denom_pub = h_pub in
      let h_section_name = Hash.Cstring.H64.hash section_name in
      let anchor_time = stamp_start in
      let duration_withdraw =
        Timestamp.diff stamp_start stamp_expire_withdraw
      in
      sign_f ~f:Secmod_rsa.sign_with_sm_key
        { h_denom_pub; h_section_name; anchor_time; duration_withdraw }
    in
    FutureDenom.
      {
        section_name;
        value;
        stamp_start;
        stamp_expire_withdraw;
        stamp_expire_deposit;
        stamp_expire_legal;
        denom_pub;
        fee_withdraw;
        fee_deposit;
        fee_refresh;
        fee_refund;
        denom_secmod_sig;
      }

  let _get_future_denominations () =
    Secmod_rsa.keys ()
    |> List.map (fun (section_name, pub, t1, _t2) ->
        let h_pub = DenominationHash.hash (RsaPublicKey.to_octets pub) in
        (*?? let duration_withdraw = Time.Absolute.diff t1 t2 in*)
        let* opt = find_denom h_pub in
        match opt with
        | None ->
            let* coin =
              Config.Coin.all_coins
              |> List.find_opt (fun coin ->
                  coin.Config.Coin.section_name = section_name)
              |> function
              | Some v -> Ok v
              | None ->
                  Fmt.error "coin `%s` not found in configuration" section_name
            in
            let future_dn = make_future_dn ~coin ~pub ~start:t1 in
            Ok (Some future_dn)
        | Some sk -> (
            match Timestamp.of_absolute t1 = sk.stamp_start with
            | false ->
                Fmt.error
                  "secmod/database stamp_start mismatch for denomination `%s`"
                  section_name
            | true -> Ok None))

  let list_to_ht l = Hashtbl.of_seq (List.to_seq l)

  let load () =
    let* sm_key = read_eddsa sm_key_fname in
    let sm_pubkey = EddsaPrivateKey.pub_of_priv sm_key in

    let* sk_keys = list_map read_eddsa (List.init 1 sk_fname) in
    let* sk_l =
      list_map
        (fun priv ->
          let pub = EddsaPrivateKey.pub_of_priv priv in
          let+ sk = database_find_sk conn pub in
          ((pub, sk), (pub, priv)))
        sk_keys
    in
    let sk_ht, sk_key_ht =
      match List.split sk_l with l1, l2 -> (list_to_ht l1, list_to_ht l2)
    in

    let dn_section_name_ht = Hashtbl.create 0xff in
    let* dn_keys =
      list_map
        (fun coin ->
          let section_name = coin.Config.Coin.section_name in
          let+ priv = read_rsa (dn_fname section_name) in
          (section_name, priv))
        Config.Coin.all_coins
    in
    let* dn_l =
      list_map
        (fun (section_name, priv) ->
          let h_pub =
            priv
            |> RsaPrivateKey.pub_of_priv
            |> RsaPublicKey.to_octets
            |> DenominationHash.hash
          in
          (* fill dn_section_name_ht *)
          Hashtbl.replace dn_section_name_ht h_pub section_name;
          let+ dn = database_find_dn conn h_pub in
          ((h_pub, dn), (h_pub, priv)))
        dn_keys
    in
    let dn_ht, dn_key_ht =
      match List.split dn_l with l1, l2 -> (list_to_ht l1, list_to_ht l2)
    in
    (* future keys are not stored anywhere until they are certified with a master_sig
       so we don't have any future key to load *)
    let t =
      {
        sm_key;
        sm_pubkey;
        sk_ht;
        dn_ht;
        sk_key_ht;
        dn_key_ht;
        future_sk_ht= Hashtbl.create 0xff;
        future_dn_ht= Hashtbl.create 0xff;
        future_sk_key_ht= Hashtbl.create 0xff;
        future_dn_key_ht= Hashtbl.create 0xff;
        dn_section_name_ht;
      }
    in
    Ok t

  let init () =
    let dir = Fpath.v Config.Exchange_secmod_eddsa.key_dir in
    let* b = Bos.OS.Dir.create ~mode:0o700 dir |> unwrap_err_msg in
    if b then Logs.info (fun m -> m "Keys: created directory `%a`" Fpath.pp dir);
    let* l =
      Bos.OS.Dir.contents ~dotfiles:false ~rel:false dir |> unwrap_err_msg
    in
    match List.is_empty l with
    | true ->
        Logs.info (fun m -> m "Keys: empty storage, generating fresh keys");
        let t = make_new () in
        Ok t
    | false ->
        Logs.info (fun m -> m "Keys: loading keys from storage");
        load ()

  let t =
    match init () with
    | Error e -> Fmt.failwith "Keys: initialization failure: `%s`." e
    | Ok t ->
        Logs.info (fun m -> m "Keys: initialized");
        t

  let sm_pubkey = t.sm_pubkey
  let sign_with_sm_key s = sign_with_sm_key t s

  let sign_with_signkey ~pub s =
    match Hashtbl.find_opt t.sk_key_ht pub with
    | None -> Fmt.failwith "Keys sign_with_signkey failure: not found."
    | Some priv -> EddsaSignature.sign ~key:priv s

  let verify_with_sm_key s ~msg = EddsaSignature.verify ~key:t.sm_pubkey s ~msg

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

  let verify_with_signkey ~pub s ~msg =
    match Hashtbl.find_opt t.sk_ht pub with
    | None -> Fmt.failwith "Keys verify_with_signkey failure: not found."
    | Some _sk -> EddsaSignature.verify ~key:pub s ~msg

  let get_signkeys () = t.sk_ht |> Hashtbl.to_seq_values |> List.of_seq
  let get_denominations () = t.dn_ht |> Hashtbl.to_seq_values |> List.of_seq

  let get_future_signkeys () =
    t.future_sk_ht |> Hashtbl.to_seq_values |> List.of_seq

  let get_future_denominations () =
    t.future_dn_ht |> Hashtbl.to_seq_values |> List.of_seq

  let find_signkey pub = Hashtbl.find_opt t.sk_ht pub
  let find_denomination h_pub = Hashtbl.find_opt t.dn_ht h_pub
  let find_future_signkey pub = Hashtbl.find_opt t.future_sk_ht pub
  let find_future_denomination h_pub = Hashtbl.find_opt t.future_dn_ht h_pub

  let certify_future_signkey pub ~master_sig =
    match
      ( Hashtbl.find_opt t.future_sk_ht pub,
        Hashtbl.find_opt t.future_sk_key_ht pub )
    with
    | None, _ | _, None ->
        Error "Keys certify_future_signkey: future signkey not found."
    | Some future_sk, Some priv -> (
        match Hashtbl.find_opt t.sk_ht pub with
        | Some _sk -> Error "Keys certify_future_signkey: already certified"
        | None ->
            let Api.FutureSignKey.
                  {
                    key;
                    stamp_start;
                    stamp_expire;
                    stamp_end;
                    signkey_secmod_sig= _;
                  } =
              future_sk
            in
            let sk =
              Signkey.
                {
                  pub= key;
                  stamp_start;
                  stamp_expire;
                  stamp_end;
                  master_sig;
                  revoked_sig= None;
                }
            in
            Hashtbl.replace t.sk_ht pub sk;
            Hashtbl.replace t.sk_key_ht pub priv;
            Hashtbl.remove t.future_sk_ht pub;
            Hashtbl.remove t.future_sk_key_ht pub;

            let* () = Pg.insert_signkey (module Conn) sk |> unwrap_err_caqti in
            Ok ())

  let certify_future_denomination h_pub ~master_sig =
    match
      ( Hashtbl.find_opt t.future_dn_ht h_pub,
        Hashtbl.find_opt t.future_dn_key_ht h_pub )
    with
    | None, _ | _, None ->
        Error "Keys certify_future_denomination: future denomination not found."
    | Some future_dn, Some priv -> (
        match Hashtbl.find_opt t.dn_ht h_pub with
        | Some _dn ->
            Error "Keys certify_future_denomination: already certified"
        | None ->
            let Api.FutureDenom.
                  {
                    section_name;
                    value;
                    stamp_start;
                    stamp_expire_withdraw;
                    stamp_expire_deposit;
                    stamp_expire_legal;
                    denom_pub;
                    fee_withdraw;
                    fee_deposit;
                    fee_refresh;
                    fee_refund;
                    denom_secmod_sig= _;
                  } =
              future_dn
            in
            let rsa_pub =
              match denom_pub with
              | Rsa Api.RsaDenominationKey.{ age_mask= _; rsa_pub } -> rsa_pub
            in
            let dn =
              Denomination.
                {
                  pub= rsa_pub;
                  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;
                  revoked_sig= None;
                }
            in
            Hashtbl.replace t.dn_ht h_pub dn;
            Hashtbl.replace t.dn_key_ht h_pub priv;
            Hashtbl.replace t.dn_section_name_ht h_pub section_name;
            Hashtbl.remove t.future_dn_ht h_pub;
            Hashtbl.remove t.future_dn_key_ht h_pub;

            let* () = Pg.insert_denom (module Conn) dn |> unwrap_err_caqti in
            Ok ())

  let revoke_signkey pub revoked_sig =
    match Hashtbl.find_opt t.sk_ht pub with
    | None -> Error "Keys revoke_signkey: signkey not found."
    | Some sk ->
        let sk = { sk with revoked_sig= Some revoked_sig } in
        Hashtbl.replace t.sk_ht pub sk;

        let+ () =
          Pg.insert_signkey_revocation conn pub revoked_sig |> unwrap_err_caqti
        in
        ()

  let revoke_denomination h_pub revoked_sig =
    match Hashtbl.find_opt t.dn_ht h_pub with
    | None -> Error "Keys revoke_denomination: denomination not found."
    | Some dn ->
        let dn = { dn with revoked_sig= Some revoked_sig } in
        Hashtbl.replace t.dn_ht h_pub dn;

        let+ () =
          Pg.insert_denomination_revocation conn h_pub revoked_sig
          |> unwrap_err_caqti
        in
        ()

  let save () =
    let* () = write_eddsa sm_key_fname t.sm_key in
    let* () =
      Hashtbl.to_seq_values t.sk_key_ht
      |> List.of_seq
      |> List.mapi (fun i priv -> write_eddsa (sk_fname i) priv)
      |> list_iter Fun.id
    in
    let* () =
      Hashtbl.to_seq t.dn_key_ht
      |> List.of_seq
      |> list_iter (fun (h_pub, priv) ->
          match Hashtbl.find_opt t.dn_section_name_ht h_pub with
          | None -> Error "Keys save: invalid state, section_name not found"
          | Some section_name -> write_rsa (dn_fname section_name) priv)
    in
    Logs.info (fun m -> m "saved private keys data");
    Ok ()
end