//! Cryptographic operations for signing and verification

use crate::encoding::{base64url_decode, base64url_encode};
use crate::error::{PlcError, Result};
use k256::ecdsa::{
    signature::Signer as K256Signer, signature::Verifier as K256Verifier,
    Signature as K256Signature, SigningKey as K256SigningKey, VerifyingKey as K256VerifyingKey,
};
use p256::ecdsa::{
    Signature as P256Signature, SigningKey as P256SigningKey, VerifyingKey as P256VerifyingKey,
};
use zeroize::Zeroizing;

/// Multicodec prefix for secp256k1 public key
const SECP256K1_MULTICODEC: &[u8] = &[0xe7, 0x01];

/// Multicodec prefix for P-256 public key
const P256_MULTICODEC: &[u8] = &[0x80, 0x24];

/// Multibase prefix for base58btc encoding
const MULTIBASE_BASE58BTC: u8 = b'z';

/// A signing key that can be either P-256 or secp256k1
#[derive(Clone)]
pub enum SigningKey {
    /// NIST P-256 signing key
    P256(P256SigningKey),
    /// secp256k1 signing key
    K256(K256SigningKey),
}

impl SigningKey {
    /// Generate a new P-256 key pair
    ///
    /// # Examples
    ///
    /// ```
    /// use atproto_plc::crypto::SigningKey;
    ///
    /// let key = SigningKey::generate_p256();
    /// ```
    pub fn generate_p256() -> Self {
        let key = P256SigningKey::random(&mut rand::thread_rng());
        SigningKey::P256(key)
    }

    /// Generate a new secp256k1 key pair
    ///
    /// # Examples
    ///
    /// ```
    /// use atproto_plc::crypto::SigningKey;
    ///
    /// let key = SigningKey::generate_k256();
    /// ```
    pub fn generate_k256() -> Self {
        let key = K256SigningKey::random(&mut rand::thread_rng());
        SigningKey::K256(key)
    }

    /// Convert this signing key to a did:key string
    ///
    /// The did:key format uses multibase (base58btc) and multicodec encoding
    ///
    /// # Examples
    ///
    /// ```
    /// use atproto_plc::crypto::SigningKey;
    ///
    /// let key = SigningKey::generate_p256();
    /// let did_key = key.to_did_key();
    /// assert!(did_key.starts_with("did:key:"));
    /// ```
    pub fn to_did_key(&self) -> String {
        let verifying_key = self.verifying_key();
        verifying_key.to_did_key()
    }

    /// Get the verifying key (public key) for this signing key
    pub fn verifying_key(&self) -> VerifyingKey {
        match self {
            SigningKey::P256(key) => VerifyingKey::P256(*key.verifying_key()),
            SigningKey::K256(key) => VerifyingKey::K256(*key.verifying_key()),
        }
    }

    /// Sign data with this key
    ///
    /// # Errors
    ///
    /// Returns `PlcError::CryptoError` if signing fails
    pub fn sign(&self, data: &[u8]) -> Result<Vec<u8>> {
        match self {
            SigningKey::P256(key) => {
                let signature: P256Signature = key.sign(data);
                Ok(signature.to_vec())
            }
            SigningKey::K256(key) => {
                let signature: K256Signature = key.sign(data);
                Ok(signature.to_vec())
            }
        }
    }

    /// Sign data and return base64url-encoded signature
    pub fn sign_base64url(&self, data: &[u8]) -> Result<String> {
        let signature = self.sign(data)?;
        Ok(base64url_encode(&signature))
    }
}

impl Drop for SigningKey {
    fn drop(&mut self) {
        // Zeroize the key material when dropped
        match self {
            SigningKey::P256(key) => {
                let bytes = Zeroizing::new(key.to_bytes());
                drop(bytes);
            }
            SigningKey::K256(key) => {
                let bytes = Zeroizing::new(key.to_bytes());
                drop(bytes);
            }
        }
    }
}

/// A verifying key (public key) that can be either P-256 or secp256k1
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VerifyingKey {
    /// NIST P-256 verifying key
    P256(P256VerifyingKey),
    /// secp256k1 verifying key
    K256(K256VerifyingKey),
}

impl VerifyingKey {
    /// Parse a verifying key from a did:key string
    ///
    /// # Errors
    ///
    /// Returns `PlcError::InvalidDidKey` if the string is not a valid did:key
    pub fn from_did_key(did_key: &str) -> Result<Self> {
        if !did_key.starts_with("did:key:") {
            return Err(PlcError::InvalidDidKey(
                "DID key must start with 'did:key:'".to_string(),
            ));
        }

        let multibase_str = &did_key[8..]; // Skip "did:key:"

        // Decode multibase (base58btc)
        if !multibase_str.starts_with(char::from(MULTIBASE_BASE58BTC)) {
            return Err(PlcError::InvalidDidKey(
                "Only base58btc multibase encoding is supported".to_string(),
            ));
        }

        let base58_str = &multibase_str[1..]; // Skip 'z' prefix
        let decoded = bs58::decode(base58_str)
            .into_vec()
            .map_err(|e| PlcError::InvalidDidKey(format!("Base58 decode error: {}", e)))?;

        // Check multicodec prefix and extract public key bytes
        if decoded.starts_with(SECP256K1_MULTICODEC) {
            let key_bytes = &decoded[SECP256K1_MULTICODEC.len()..];
            let verifying_key = K256VerifyingKey::from_sec1_bytes(key_bytes)
                .map_err(|e| PlcError::InvalidDidKey(format!("Invalid secp256k1 key: {}", e)))?;
            Ok(VerifyingKey::K256(verifying_key))
        } else if decoded.starts_with(P256_MULTICODEC) {
            let key_bytes = &decoded[P256_MULTICODEC.len()..];
            let verifying_key = P256VerifyingKey::from_sec1_bytes(key_bytes)
                .map_err(|e| PlcError::InvalidDidKey(format!("Invalid P-256 key: {}", e)))?;
            Ok(VerifyingKey::P256(verifying_key))
        } else {
            Err(PlcError::UnsupportedKeyType(format!(
                "Unknown multicodec prefix: {:?}",
                &decoded[..2.min(decoded.len())]
            )))
        }
    }

    /// Convert this verifying key to a did:key string
    pub fn to_did_key(&self) -> String {
        let (multicodec, key_bytes) = match self {
            VerifyingKey::P256(key) => (P256_MULTICODEC, key.to_sec1_bytes().to_vec()),
            VerifyingKey::K256(key) => (SECP256K1_MULTICODEC, key.to_sec1_bytes().to_vec()),
        };

        // Combine multicodec prefix and key bytes
        let mut combined = Vec::with_capacity(multicodec.len() + key_bytes.len());
        combined.extend_from_slice(multicodec);
        combined.extend_from_slice(&key_bytes);

        // Encode with multibase (base58btc)
        let encoded = format!("{}{}", MULTIBASE_BASE58BTC as char, bs58::encode(combined).into_string());

        format!("did:key:{}", encoded)
    }

    /// Verify a signature using this key
    ///
    /// # Errors
    ///
    /// Returns `PlcError::SignatureVerificationFailed` if verification fails
    pub fn verify(&self, data: &[u8], signature: &[u8]) -> Result<()> {
        match self {
            VerifyingKey::P256(key) => {
                let sig = P256Signature::from_slice(signature)
                    .map_err(|e| PlcError::CryptoError(format!("Invalid signature format: {}", e)))?;
                key.verify(data, &sig)
                    .map_err(|_| PlcError::SignatureVerificationFailed)?;
            }
            VerifyingKey::K256(key) => {
                let sig = K256Signature::from_slice(signature)
                    .map_err(|e| PlcError::CryptoError(format!("Invalid signature format: {}", e)))?;
                key.verify(data, &sig)
                    .map_err(|_| PlcError::SignatureVerificationFailed)?;
            }
        }
        Ok(())
    }

    /// Verify a base64url-encoded signature
    pub fn verify_base64url(&self, data: &[u8], signature_b64: &str) -> Result<()> {
        let signature = base64url_decode(signature_b64)?;
        self.verify(data, &signature)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_p256_keygen_and_sign() {
        let key = SigningKey::generate_p256();
        let data = b"hello world";
        let signature = key.sign(data).unwrap();
        assert!(!signature.is_empty());
    }

    #[test]
    fn test_k256_keygen_and_sign() {
        let key = SigningKey::generate_k256();
        let data = b"hello world";
        let signature = key.sign(data).unwrap();
        assert!(!signature.is_empty());
    }

    #[test]
    fn test_p256_sign_verify() {
        let key = SigningKey::generate_p256();
        let data = b"hello world";
        let signature = key.sign(data).unwrap();

        let verifying_key = key.verifying_key();
        assert!(verifying_key.verify(data, &signature).is_ok());

        // Wrong data should fail
        let wrong_data = b"goodbye world";
        assert!(verifying_key.verify(wrong_data, &signature).is_err());
    }

    #[test]
    fn test_k256_sign_verify() {
        let key = SigningKey::generate_k256();
        let data = b"hello world";
        let signature = key.sign(data).unwrap();

        let verifying_key = key.verifying_key();
        assert!(verifying_key.verify(data, &signature).is_ok());
    }

    #[test]
    fn test_did_key_roundtrip_p256() {
        let key = SigningKey::generate_p256();
        let did_key = key.to_did_key();
        assert!(did_key.starts_with("did:key:z"));

        let parsed = VerifyingKey::from_did_key(&did_key).unwrap();
        assert_eq!(parsed, key.verifying_key());
    }

    #[test]
    fn test_did_key_roundtrip_k256() {
        let key = SigningKey::generate_k256();
        let did_key = key.to_did_key();
        assert!(did_key.starts_with("did:key:z"));

        let parsed = VerifyingKey::from_did_key(&did_key).unwrap();
        assert_eq!(parsed, key.verifying_key());
    }

    #[test]
    fn test_base64url_sign_verify() {
        let key = SigningKey::generate_p256();
        let data = b"hello world";
        let signature_b64 = key.sign_base64url(data).unwrap();

        let verifying_key = key.verifying_key();
        assert!(verifying_key.verify_base64url(data, &signature_b64).is_ok());
    }

    #[test]
    fn test_invalid_did_key() {
        assert!(VerifyingKey::from_did_key("invalid").is_err());
        assert!(VerifyingKey::from_did_key("did:web:example.com").is_err());
    }
}