// Copyright (c) 2021 Paul Miller
// Copyright (c) 2022-2023 Yuki Kishimoto
// Copyright (c) 2023-2025 Rust Nostr Developers
// Distributed under the MIT software license

//! Keys

use alloc::boxed::Box;
use core::cmp::Ordering;
use core::fmt;
use core::hash::{Hash, Hasher};
#[cfg(feature = "std")]
use core::str::FromStr;

#[cfg(all(feature = "std", feature = "os-rng"))]
use rand::TryRngCore;
#[cfg(all(feature = "std", feature = "os-rng"))]
use rand::rngs::OsRng;
#[cfg(feature = "rand")]
use rand::{CryptoRng, RngCore};
use secp256k1::schnorr::Signature;
use secp256k1::{self, Keypair, Message, Secp256k1, Signing, XOnlyPublicKey};

pub mod public_key;
pub mod secret_key;

pub use self::public_key::PublicKey;
pub use self::secret_key::SecretKey;
#[cfg(all(feature = "std", feature = "os-rng"))]
use crate::event::{Event, EventId, UnsignedEvent};
#[cfg(all(feature = "std", feature = "os-rng"))]
use crate::nips::nip04::{AsyncNip04, Nip04};
#[cfg(all(feature = "std", feature = "os-rng"))]
use crate::nips::nip44::{AsyncNip44, Nip44};
use crate::signer::{AsyncGetPublicKey, GetPublicKey, SignerError};
#[cfg(all(feature = "std", feature = "os-rng"))]
use crate::signer::{AsyncNostrSigner, AsyncSignEvent, NostrSigner, SignEvent, SignerBackend};
#[cfg(feature = "rand")]
use crate::util;
use crate::util::BoxedFuture;
#[cfg(feature = "std")]
use crate::util::SECP256K1;

/// [`Keys`] error
#[derive(Debug, PartialEq)]
pub enum Error {
    /// Secp256k1 error
    Secp256k1(secp256k1::Error),
    /// Hex decode error
    Hex(faster_hex::Error),
    /// Invalid secret key
    InvalidSecretKey,
    /// Invalid public key
    InvalidPublicKey,
}

impl core::error::Error for Error {}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Secp256k1(e) => e.fmt(f),
            Self::Hex(e) => e.fmt(f),
            Self::InvalidSecretKey => f.write_str("Invalid secret key"),
            Self::InvalidPublicKey => f.write_str("Invalid public key"),
        }
    }
}

impl From<secp256k1::Error> for Error {
    fn from(e: secp256k1::Error) -> Self {
        Self::Secp256k1(e)
    }
}

impl From<faster_hex::Error> for Error {
    fn from(e: faster_hex::Error) -> Self {
        Self::Hex(e)
    }
}

/// Nostr keys
#[derive(Clone)]
pub struct Keys {
    /// Public key
    pub public_key: PublicKey,
    secret_key: SecretKey,
    keypair: Keypair,
}

impl fmt::Debug for Keys {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Keys")
            .field("public_key", &self.public_key)
            .finish()
    }
}

impl PartialEq for Keys {
    fn eq(&self, other: &Self) -> bool {
        self.public_key == other.public_key
    }
}

impl Eq for Keys {}

impl PartialOrd for Keys {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Keys {
    fn cmp(&self, other: &Self) -> Ordering {
        self.public_key.cmp(&other.public_key)
    }
}

impl Hash for Keys {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.public_key.hash(state)
    }
}

impl Keys {
    /// Construct from a secret key.
    ///
    /// This method internally constructs the [`Keypair`] and derives the [`PublicKey`].
    #[inline]
    #[cfg(feature = "std")]
    pub fn new(secret_key: SecretKey) -> Self {
        Self::new_with_ctx(&SECP256K1, secret_key)
    }

    /// Construct from a secret key.
    ///
    /// This method internally constructs the [`Keypair`] and derives the [`PublicKey`].
    pub fn new_with_ctx<C>(secp: &Secp256k1<C>, secret_key: SecretKey) -> Self
    where
        C: Signing,
    {
        let keypair: Keypair = Keypair::from_secret_key(secp, &secret_key);
        let public_key: XOnlyPublicKey = XOnlyPublicKey::from_keypair(&keypair).0;

        Self {
            public_key: PublicKey::from(public_key),
            secret_key,
            keypair,
        }
    }

    /// Parse secret key and construct keys.
    ///
    /// Check [`SecretKey::parse`] to learn more about secret key parsing.
    #[inline]
    #[cfg(feature = "std")]
    pub fn parse(secret_key: &str) -> Result<Self, Error> {
        Self::parse_with_ctx(&SECP256K1, secret_key)
    }

    /// Parse secret key and construct keys.
    ///
    /// Check [`SecretKey::parse`] to learn more about secret key parsing.
    #[inline]
    pub fn parse_with_ctx<C>(secp: &Secp256k1<C>, secret_key: &str) -> Result<Self, Error>
    where
        C: Signing,
    {
        let secret_key: SecretKey = SecretKey::parse(secret_key)?;
        Ok(Self::new_with_ctx(secp, secret_key))
    }

    /// Generate random keys
    ///
    /// This constructor uses a random number generator that retrieves randomness from the operating system (see [`OsRng`]).
    ///
    /// Use [`Keys::generate_with_rng`] to specify a custom random source.
    ///
    /// This internally construct a keypair, so for faster secret generation (i.e., for vanity pubkey mining),
    /// it's suggested to use [`SecretKey::generate`] instead.
    #[inline]
    #[cfg(all(feature = "std", feature = "os-rng"))]
    pub fn generate() -> Self {
        Self::generate_with_rng(&SECP256K1, &mut OsRng.unwrap_err())
    }

    /// Generate random keys
    ///
    /// This internally construct a keypair, so for faster secret generation (i.e., for vanity pubkey mining),
    /// it's suggested to use [`SecretKey::generate_with_rng`] instead.
    #[inline]
    #[cfg(feature = "rand")]
    pub fn generate_with_rng<C, R>(secp: &Secp256k1<C>, rng: &mut R) -> Self
    where
        C: Signing,
        R: RngCore,
    {
        let secret_key: SecretKey = SecretKey::generate_with_rng(rng);
        Self::new_with_ctx(secp, secret_key)
    }

    /// Get public key
    #[inline]
    pub fn public_key(&self) -> PublicKey {
        self.public_key
    }

    /// Get secret key
    #[inline]
    pub fn secret_key(&self) -> &SecretKey {
        &self.secret_key
    }

    /// Get keypair
    #[inline]
    pub fn keypair(&self) -> &Keypair {
        &self.keypair
    }

    /// Creates a schnorr signature of the [`Message`].
    ///
    /// This method uses a random number generator that retrieves randomness from the operating system (see [`OsRng`]).
    #[inline]
    #[cfg(all(feature = "std", feature = "os-rng"))]
    pub fn sign_schnorr(&self, message: &Message) -> Signature {
        self.sign_schnorr_with_rng(&SECP256K1, message, &mut OsRng.unwrap_err())
    }

    /// Creates a schnorr signature of the [`Message`] using a custom random number generation source.
    #[cfg(feature = "rand")]
    pub fn sign_schnorr_with_rng<C, R>(
        &self,
        secp: &Secp256k1<C>,
        message: &Message,
        rng: &mut R,
    ) -> Signature
    where
        C: Signing,
        R: RngCore + CryptoRng,
    {
        let aux: [u8; 32] = util::random_32_bytes(rng);
        self.sign_schnorr_with_aux_rand(secp, message, &aux)
    }

    /// Creates a schnorr signature using the given auxiliary random data.
    pub fn sign_schnorr_with_aux_rand<C>(
        &self,
        secp: &Secp256k1<C>,
        message: &Message,
        aux: &[u8; 32],
    ) -> Signature
    where
        C: Signing,
    {
        secp.sign_schnorr_with_aux_rand(message, &self.keypair, aux)
    }
}

#[cfg(feature = "std")]
impl FromStr for Keys {
    type Err = Error;

    /// Try to parse [Keys] from **secret key** `hex` or `bech32`
    #[inline]
    fn from_str(secret_key: &str) -> Result<Self, Self::Err> {
        Self::parse(secret_key)
    }
}

impl GetPublicKey for Keys {
    #[inline]
    fn get_public_key(&self) -> Result<PublicKey, SignerError> {
        Ok(self.public_key)
    }
}

#[cfg(all(feature = "std", feature = "os-rng"))]
impl SignEvent for Keys {
    fn sign_event(&self, unsigned: UnsignedEvent) -> Result<Event, SignerError> {
        let id: EventId = unsigned.id.unwrap_or_else(|| unsigned.compute_id());
        let message: Message = Message::from_digest(id.to_bytes());
        let sig: Signature = self.sign_schnorr(&message);
        unsigned.add_signature(sig).map_err(SignerError::backend)
    }
}

#[cfg(all(feature = "std", feature = "os-rng"))]
impl Nip04 for Keys {
    type Error = SignerError;

    fn nip04_encrypt(
        &self,
        _public_key: &PublicKey,
        _content: &str,
    ) -> Result<String, SignerError> {
        #[cfg(feature = "nip04")]
        {
            let secret_key: &SecretKey = self.secret_key();
            crate::nips::nip04::encrypt(secret_key, _public_key, _content)
                .map_err(SignerError::backend)
        }

        #[cfg(not(feature = "nip04"))]
        Err(SignerError::from("NIP04 feature is not enabled"))
    }

    fn nip04_decrypt(
        &self,
        _public_key: &PublicKey,
        _encrypted_content: &str,
    ) -> Result<String, SignerError> {
        #[cfg(feature = "nip04")]
        {
            let secret_key: &SecretKey = self.secret_key();
            crate::nips::nip04::decrypt(secret_key, _public_key, _encrypted_content)
                .map_err(SignerError::backend)
        }

        #[cfg(not(feature = "nip04"))]
        Err(SignerError::from("NIP04 feature is not enabled"))
    }
}

#[cfg(all(feature = "std", feature = "os-rng"))]
impl Nip44 for Keys {
    type Error = SignerError;

    fn nip44_encrypt(
        &self,
        _public_key: &PublicKey,
        _content: &str,
    ) -> Result<String, SignerError> {
        #[cfg(feature = "nip44")]
        {
            use crate::nips::nip44::{self, Version};
            let secret_key: &SecretKey = self.secret_key();
            nip44::encrypt(secret_key, _public_key, _content, Version::default())
                .map_err(SignerError::backend)
        }

        #[cfg(not(feature = "nip44"))]
        Err(SignerError::from("NIP44 feature is not enabled"))
    }

    fn nip44_decrypt(
        &self,
        _public_key: &PublicKey,
        _payload: &str,
    ) -> Result<String, SignerError> {
        #[cfg(feature = "nip44")]
        {
            let secret_key: &SecretKey = self.secret_key();
            crate::nips::nip44::decrypt(secret_key, _public_key, _payload)
                .map_err(SignerError::backend)
        }

        #[cfg(not(feature = "nip44"))]
        Err(SignerError::from("NIP44 feature is not enabled"))
    }
}

#[cfg(all(feature = "std", feature = "os-rng"))]
impl NostrSigner for Keys {
    fn backend(&self) -> SignerBackend<'_> {
        SignerBackend::Keys
    }
}

impl AsyncGetPublicKey for Keys {
    fn get_public_key_async(&self) -> BoxedFuture<'_, Result<PublicKey, SignerError>> {
        Box::pin(async move { GetPublicKey::get_public_key(self) })
    }
}

#[cfg(all(feature = "std", feature = "os-rng"))]
impl AsyncSignEvent for Keys {
    fn sign_event_async(
        &self,
        unsigned: UnsignedEvent,
    ) -> BoxedFuture<'_, Result<Event, SignerError>> {
        Box::pin(async move { SignEvent::sign_event(self, unsigned) })
    }
}

#[cfg(all(feature = "std", feature = "os-rng"))]
impl AsyncNip04 for Keys {
    type Error = SignerError;

    fn nip04_encrypt_async<'a>(
        &'a self,
        public_key: &'a PublicKey,
        content: &'a str,
    ) -> BoxedFuture<'a, Result<String, SignerError>> {
        Box::pin(async move { Nip04::nip04_encrypt(self, public_key, content) })
    }

    fn nip04_decrypt_async<'a>(
        &'a self,
        public_key: &'a PublicKey,
        encrypted_content: &'a str,
    ) -> BoxedFuture<'a, Result<String, SignerError>> {
        Box::pin(async move { Nip04::nip04_decrypt(self, public_key, encrypted_content) })
    }
}

#[cfg(all(feature = "std", feature = "os-rng"))]
impl AsyncNip44 for Keys {
    type Error = SignerError;

    fn nip44_encrypt_async<'a>(
        &'a self,
        public_key: &'a PublicKey,
        content: &'a str,
    ) -> BoxedFuture<'a, Result<String, SignerError>> {
        Box::pin(async move { Nip44::nip44_encrypt(self, public_key, content) })
    }

    fn nip44_decrypt_async<'a>(
        &'a self,
        public_key: &'a PublicKey,
        payload: &'a str,
    ) -> BoxedFuture<'a, Result<String, SignerError>> {
        Box::pin(async move { Nip44::nip44_decrypt(self, public_key, payload) })
    }
}

#[cfg(all(feature = "std", feature = "os-rng"))]
impl AsyncNostrSigner for Keys {
    fn backend(&self) -> SignerBackend<'_> {
        SignerBackend::Keys
    }
}

#[cfg(test)]
#[cfg(feature = "std")]
mod tests {
    use super::*;

    const SECRET_KEY_BECH32: &str =
        "nsec1j4c6269y9w0q2er2xjw8sv2ehyrtfxq3jwgdlxj6qfn8z4gjsq5qfvfk99";
    const SECRET_KEY_HEX: &str = "6b911fd37cdf5c81d4c0adb1ab7fa822ed253ab0ad9aa18d77257c88b29b718e";

    #[test]
    fn parse_keys() -> Result<(), Error> {
        Keys::parse(SECRET_KEY_BECH32)?;
        Keys::parse(SECRET_KEY_HEX)?;
        Ok(())
    }

    #[test]
    fn parse_invalid_keys() {
        assert_eq!(Keys::parse("nsec...").unwrap_err(), Error::InvalidSecretKey);
        assert_eq!(
            Keys::parse("npub14f8usejl26twx0dhuxjh9cas7keav9vr0v8nvtwtrjqx3vycc76qqh9nsy")
                .unwrap_err(),
            Error::InvalidSecretKey
        );
        assert_eq!(
            Keys::parse("6b911fd37cdf5c8").unwrap_err(),
            Error::InvalidSecretKey
        );
    }
}

#[cfg(bench)]
#[cfg(all(feature = "std", feature = "os-rng"))]
mod benches {
    use test::{Bencher, black_box};

    use super::*;

    #[bench]
    pub fn generate_keys(bh: &mut Bencher) {
        bh.iter(|| {
            black_box(Keys::generate());
        });
    }
}
