//! ECDSA (Elliptic Curve Digital Signature Algorithm) signature implementation over secp256k1
//! curve using Keccak to hash the messages when signing.

use alloc::{string::ToString, vec::Vec};

use k256::{
    ecdh::diffie_hellman,
    ecdsa::{RecoveryId, SigningKey, VerifyingKey, signature::hazmat::PrehashVerifier},
};
use miden_crypto_derive::{SilentDebug, SilentDisplay};
use rand::{CryptoRng, RngCore};
use thiserror::Error;

use crate::{
    Felt, SequentialCommit, Word,
    ecdh::k256::{EphemeralPublicKey, SharedSecret},
    utils::{
        ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
        bytes_to_packed_u32_elements,
        zeroize::{Zeroize, ZeroizeOnDrop},
    },
};

#[cfg(test)]
mod tests;

// CONSTANTS
// ================================================================================================

/// Length of secret key in bytes
const SECRET_KEY_BYTES: usize = 32;
/// Length of public key in bytes when using compressed format encoding
pub(crate) const PUBLIC_KEY_BYTES: usize = 33;
/// Length of signature in bytes using our custom serialization
const SIGNATURE_BYTES: usize = 65;
/// Length of signature in bytes using standard serialization i.e., `SEC1`
const SIGNATURE_STANDARD_BYTES: usize = 64;
/// Length of scalars for the `secp256k1` curve
const SCALARS_SIZE_BYTES: usize = 32;

// SECRET KEY
// ================================================================================================

/// Secret key for ECDSA signature verification over secp256k1 curve.
#[derive(Clone, SilentDebug, SilentDisplay)]
pub struct SecretKey {
    inner: SigningKey,
}

impl SecretKey {
    /// Generates a new random secret key using the OS random number generator.
    #[cfg(feature = "std")]
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        let mut rng = rand::rng();

        Self::with_rng(&mut rng)
    }

    /// Generates a new secret key using the provided random number generator.
    pub fn with_rng<R: CryptoRng + RngCore>(rng: &mut R) -> Self {
        // we use a seedable CSPRNG and seed it with `rng`
        // this is a work around the fact that the version of the `rand` dependency in our crate
        // is different than the one used in the `k256` one. This solution will no longer be needed
        // once `k256` gets a new release with a version of the `rand` dependency matching ours
        use k256::elliptic_curve::rand_core::SeedableRng;
        let mut seed = [0_u8; 32];
        rand::RngCore::fill_bytes(rng, &mut seed);
        let mut rng = rand_hc::Hc128Rng::from_seed(seed);

        let signing_key = SigningKey::random(&mut rng);

        // Zeroize the seed to prevent leaking secret material
        seed.zeroize();

        Self { inner: signing_key }
    }

    /// Gets the corresponding public key for this secret key.
    pub fn public_key(&self) -> PublicKey {
        let verifying_key = self.inner.verifying_key();
        PublicKey { inner: *verifying_key }
    }

    /// Signs a message (represented as a Word) with this secret key.
    pub fn sign(&self, message: Word) -> Signature {
        let message_digest = hash_message(message);
        self.sign_prehash(message_digest)
    }

    /// Signs a pre-hashed message with this secret key.
    pub fn sign_prehash(&self, message_digest: [u8; 32]) -> Signature {
        let (signature_inner, recovery_id) = self
            .inner
            .sign_prehash_recoverable(&message_digest)
            .expect("failed to generate signature");

        let (r, s) = signature_inner.split_scalars();

        Signature {
            r: r.to_bytes().into(),
            s: s.to_bytes().into(),
            v: recovery_id.into(),
        }
    }

    /// Computes a Diffie-Hellman shared secret from this secret key and the ephemeral public key
    /// generated by the other party.
    pub fn get_shared_secret(&self, pk_e: EphemeralPublicKey) -> SharedSecret {
        let shared_secret_inner = diffie_hellman(self.inner.as_nonzero_scalar(), pk_e.as_affine());

        SharedSecret::new(shared_secret_inner)
    }
}

// SAFETY: The inner `k256::ecdsa::SigningKey` already implements `ZeroizeOnDrop`,
// which ensures that the secret key material is securely zeroized when dropped.
impl ZeroizeOnDrop for SecretKey {}

impl PartialEq for SecretKey {
    fn eq(&self, other: &Self) -> bool {
        use subtle::ConstantTimeEq;
        self.to_bytes().ct_eq(&other.to_bytes()).into()
    }
}

impl Eq for SecretKey {}

// PUBLIC KEY
// ================================================================================================

/// Public key for ECDSA signature verification over secp256k1 curve.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PublicKey {
    pub(crate) inner: VerifyingKey,
}

impl PublicKey {
    /// Returns a commitment to the public key using the RPO256 hash function.
    ///
    /// The commitment is computed by first converting the public key to field elements (4 bytes
    /// per element), and then computing a sequential hash of the elements.
    pub fn to_commitment(&self) -> Word {
        <Self as SequentialCommit>::to_commitment(self)
    }

    /// Verifies a signature against this public key and message.
    pub fn verify(&self, message: Word, signature: &Signature) -> bool {
        let message_digest = hash_message(message);
        self.verify_prehash(message_digest, signature)
    }

    /// Verifies a signature against this public key and pre-hashed message.
    pub fn verify_prehash(&self, message_digest: [u8; 32], signature: &Signature) -> bool {
        let signature_inner = k256::ecdsa::Signature::from_scalars(*signature.r(), *signature.s());

        match signature_inner {
            Ok(signature) => self.inner.verify_prehash(&message_digest, &signature).is_ok(),
            Err(_) => false,
        }
    }

    /// Recovers from the signature the public key associated to the secret key used to sign the
    /// message.
    pub fn recover_from(message: Word, signature: &Signature) -> Result<Self, PublicKeyError> {
        let message_digest = hash_message(message);
        let signature_data = k256::ecdsa::Signature::from_scalars(*signature.r(), *signature.s())
            .map_err(|_| PublicKeyError::RecoveryFailed)?;

        let verifying_key = k256::ecdsa::VerifyingKey::recover_from_prehash(
            &message_digest,
            &signature_data,
            RecoveryId::from_byte(signature.v()).ok_or(PublicKeyError::RecoveryFailed)?,
        )
        .map_err(|_| PublicKeyError::RecoveryFailed)?;

        Ok(Self { inner: verifying_key })
    }
}

impl SequentialCommit for PublicKey {
    type Commitment = Word;

    fn to_elements(&self) -> Vec<Felt> {
        bytes_to_packed_u32_elements(&self.to_bytes())
    }
}

#[derive(Debug, Error)]
pub enum PublicKeyError {
    #[error("Could not recover the public key from the message and signature")]
    RecoveryFailed,
}

// SIGNATURE
// ================================================================================================

/// ECDSA signature over secp256k1 curve using Keccak to hash the messages when signing.
///
/// ## Serialization Formats
///
/// This implementation supports 2 serialization formats:
///
/// ### Custom Format (66 bytes):
/// - Bytes 0-31: r component (32 bytes, big-endian)
/// - Bytes 32-63: s component (32 bytes, big-endian)
/// - Byte 64: recovery ID (v) - values 0-3
///
/// ### SEC1 Format (64 bytes):
/// - Bytes 0-31: r component (32 bytes, big-endian)
/// - Bytes 32-63: s component (32 bytes, big-endian)
/// - Note: Recovery ID
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Signature {
    r: [u8; SCALARS_SIZE_BYTES],
    s: [u8; SCALARS_SIZE_BYTES],
    v: u8,
}

impl Signature {
    /// Returns the `r` scalar of this signature.
    pub fn r(&self) -> &[u8; SCALARS_SIZE_BYTES] {
        &self.r
    }

    /// Returns the `s` scalar of this signature.
    pub fn s(&self) -> &[u8; SCALARS_SIZE_BYTES] {
        &self.s
    }

    /// Returns the `v` component of this signature, which is a `u8` representing the recovery id.
    pub fn v(&self) -> u8 {
        self.v
    }

    /// Verifies this signature against a message and public key.
    pub fn verify(&self, message: Word, pub_key: &PublicKey) -> bool {
        pub_key.verify(message, self)
    }

    /// Converts signature to SEC1 format (standard 64-byte r||s format).
    ///
    /// This format is the standard one used by most ECDSA libraries but loses the recovery ID.
    pub fn to_sec1_bytes(&self) -> [u8; SIGNATURE_STANDARD_BYTES] {
        let mut bytes = [0u8; 2 * SCALARS_SIZE_BYTES];
        bytes[0..SCALARS_SIZE_BYTES].copy_from_slice(self.r());
        bytes[SCALARS_SIZE_BYTES..2 * SCALARS_SIZE_BYTES].copy_from_slice(self.s());
        bytes
    }

    /// Creates a signature from SEC1 format bytes with a given recovery id.
    ///
    /// # Arguments
    /// * `bytes` - 64-byte array containing r and s components
    /// * `recovery_id` - recovery ID (0-3)
    pub fn from_sec1_bytes_and_recovery_id(
        bytes: [u8; SIGNATURE_STANDARD_BYTES],
        v: u8,
    ) -> Result<Self, DeserializationError> {
        let mut r = [0u8; SCALARS_SIZE_BYTES];
        let mut s = [0u8; SCALARS_SIZE_BYTES];
        r.copy_from_slice(&bytes[0..SCALARS_SIZE_BYTES]);
        s.copy_from_slice(&bytes[SCALARS_SIZE_BYTES..2 * SCALARS_SIZE_BYTES]);

        if v > 3 {
            return Err(DeserializationError::InvalidValue(r#"Invalid recovery ID"#.to_string()));
        }

        Ok(Signature { r, s, v })
    }
}

// SERIALIZATION / DESERIALIZATION
// ================================================================================================

impl Serializable for SecretKey {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        let mut buffer = Vec::with_capacity(SECRET_KEY_BYTES);
        let sk_bytes: [u8; SECRET_KEY_BYTES] = self.inner.to_bytes().into();
        buffer.extend_from_slice(&sk_bytes);

        target.write_bytes(&buffer);
    }
}

impl Deserializable for SecretKey {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let mut bytes: [u8; SECRET_KEY_BYTES] = source.read_array()?;

        let signing_key = SigningKey::from_slice(&bytes)
            .map_err(|_| DeserializationError::InvalidValue("Invalid secret key".to_string()))?;
        bytes.zeroize();

        Ok(Self { inner: signing_key })
    }
}

impl Serializable for PublicKey {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        // Compressed format
        let encoded = self.inner.to_encoded_point(true);

        target.write_bytes(encoded.as_bytes());
    }
}

impl Deserializable for PublicKey {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let bytes: [u8; PUBLIC_KEY_BYTES] = source.read_array()?;

        let verifying_key = VerifyingKey::from_sec1_bytes(&bytes)
            .map_err(|_| DeserializationError::InvalidValue("Invalid public key".to_string()))?;

        Ok(Self { inner: verifying_key })
    }
}

impl Serializable for Signature {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        let mut bytes = [0u8; SIGNATURE_BYTES];
        bytes[0..SCALARS_SIZE_BYTES].copy_from_slice(self.r());
        bytes[SCALARS_SIZE_BYTES..2 * SCALARS_SIZE_BYTES].copy_from_slice(self.s());
        bytes[2 * SCALARS_SIZE_BYTES] = self.v();
        target.write_bytes(&bytes);
    }
}

impl Deserializable for Signature {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let r: [u8; SCALARS_SIZE_BYTES] = source.read_array()?;
        let s: [u8; SCALARS_SIZE_BYTES] = source.read_array()?;
        let v: u8 = source.read_u8()?;

        Ok(Signature { r, s, v })
    }
}

// HELPER
// ================================================================================================

/// Hashes a word message using Keccak.
fn hash_message(message: Word) -> [u8; 32] {
    use sha3::{Digest, Keccak256};
    let mut hasher = Keccak256::new();
    let message_bytes: [u8; 32] = message.into();
    hasher.update(message_bytes);
    hasher.finalize().into()
}
