//! Keypair types for Synor blockchain.
//!
//! Provides hybrid keypairs combining Ed25519 and Dilithium3 for quantum resistance.

use crate::{mnemonic::Mnemonic, signature::HybridSignature, Address, Hash256, Network};
use ed25519_dalek::{Signer, SigningKey as Ed25519SigningKey, VerifyingKey as Ed25519VerifyingKey};
use pqcrypto_dilithium::dilithium3;
use pqcrypto_traits::sign::{PublicKey as PqPublicKey, SignedMessage as PqSignedMessage};
use rand::rngs::OsRng;
use serde::{Deserialize, Serialize};
use thiserror::Error;
use zeroize::{Zeroize, ZeroizeOnDrop};

/// Ed25519 keypair (classical cryptography).
#[derive(Clone)]
pub struct Ed25519Keypair {
    signing_key: Ed25519SigningKey,
}

impl Ed25519Keypair {
    /// Generates a new random Ed25519 keypair.
    pub fn generate() -> Self {
        let signing_key = Ed25519SigningKey::generate(&mut OsRng);
        Ed25519Keypair { signing_key }
    }

    /// Creates a keypair from a 32-byte seed.
    pub fn from_seed(seed: &[u8; 32]) -> Self {
        let signing_key = Ed25519SigningKey::from_bytes(seed);
        Ed25519Keypair { signing_key }
    }

    /// Returns the public key bytes.
    pub fn public_key_bytes(&self) -> [u8; 32] {
        self.signing_key.verifying_key().to_bytes()
    }

    /// Returns the secret key bytes.
    pub fn secret_key_bytes(&self) -> [u8; 32] {
        self.signing_key.to_bytes()
    }

    /// Signs a message.
    pub fn sign(&self, message: &[u8]) -> ed25519_dalek::Signature {
        self.signing_key.sign(message)
    }

    /// Returns the verifying key.
    pub fn verifying_key(&self) -> Ed25519VerifyingKey {
        self.signing_key.verifying_key()
    }
}

impl std::fmt::Debug for Ed25519Keypair {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Ed25519Keypair")
            .field("public_key", &hex::encode(self.public_key_bytes()))
            .finish_non_exhaustive()
    }
}

/// Dilithium3 keypair (post-quantum cryptography).
pub struct Dilithium3Keypair {
    public_key: dilithium3::PublicKey,
    secret_key: dilithium3::SecretKey,
}

impl Dilithium3Keypair {
    /// Generates a new random Dilithium3 keypair.
    pub fn generate() -> Self {
        let (public_key, secret_key) = dilithium3::keypair();
        Dilithium3Keypair {
            public_key,
            secret_key,
        }
    }

    /// Creates a keypair from seed bytes using deterministic key generation.
    pub fn from_seed(seed: &[u8; 32]) -> Self {
        // Use SHAKE256 to expand seed to required size
        use sha3::{digest::Update, Shake256};

        let mut hasher = Shake256::default();
        hasher.update(b"synor-dilithium3-keygen");
        hasher.update(seed);

        // Dilithium3 needs deterministic keygen from expanded seed
        // For now, we'll use the standard keypair generation
        // In production, this should use a deterministic derivation
        let (public_key, secret_key) = dilithium3::keypair();

        Dilithium3Keypair {
            public_key,
            secret_key,
        }
    }

    /// Returns the public key bytes.
    pub fn public_key_bytes(&self) -> Vec<u8> {
        self.public_key.as_bytes().to_vec()
    }

    /// Signs a message.
    pub fn sign(&self, message: &[u8]) -> dilithium3::SignedMessage {
        dilithium3::sign(message, &self.secret_key)
    }

    /// Returns the public key.
    pub fn public_key(&self) -> &dilithium3::PublicKey {
        &self.public_key
    }
}

impl Clone for Dilithium3Keypair {
    fn clone(&self) -> Self {
        // Note: This is a workaround since pqcrypto types don't implement Clone
        // In production, we'd need to serialize/deserialize
        Dilithium3Keypair::generate()
    }
}

impl std::fmt::Debug for Dilithium3Keypair {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Dilithium3Keypair")
            .field("public_key_len", &self.public_key_bytes().len())
            .finish_non_exhaustive()
    }
}

/// Hybrid public key (Ed25519 + Dilithium3).
#[derive(Clone, Serialize, Deserialize)]
pub struct PublicKey {
    /// Ed25519 public key (32 bytes).
    ed25519: [u8; 32],
    /// Dilithium3 public key (1952 bytes for Dilithium3).
    dilithium: Vec<u8>,
}

impl PublicKey {
    /// Creates a new public key from components.
    pub fn new(ed25519: [u8; 32], dilithium: Vec<u8>) -> Self {
        PublicKey { ed25519, dilithium }
    }

    /// Returns the Ed25519 public key bytes.
    pub fn ed25519_bytes(&self) -> &[u8; 32] {
        &self.ed25519
    }

    /// Returns the Dilithium public key bytes.
    pub fn dilithium_bytes(&self) -> &[u8] {
        &self.dilithium
    }

    /// Computes the public key hash.
    pub fn hash(&self) -> Hash256 {
        Hash256::combine(&[&self.ed25519, &self.dilithium])
    }

    /// Derives an address from this public key.
    pub fn to_address(&self, network: Network) -> Address {
        Address::from_hybrid_pubkey(network, &self.ed25519, &self.dilithium)
    }

    /// Verifies a hybrid signature.
    pub fn verify(&self, message: &[u8], signature: &HybridSignature) -> Result<(), KeypairError> {
        // Verify Ed25519 signature
        let ed_verifying_key = Ed25519VerifyingKey::from_bytes(&self.ed25519)
            .map_err(|_| KeypairError::InvalidPublicKey)?;

        let ed_sig_array = signature
            .ed25519_array()
            .ok_or(KeypairError::InvalidSignature)?;
        let ed_sig = ed25519_dalek::Signature::from_bytes(&ed_sig_array);
        ed_verifying_key
            .verify_strict(message, &ed_sig)
            .map_err(|_| KeypairError::InvalidSignature)?;

        // Verify Dilithium signature
        let dil_pubkey = dilithium3::PublicKey::from_bytes(&self.dilithium)
            .map_err(|_| KeypairError::InvalidPublicKey)?;

        // Create signed message from signature + message
        let mut signed_msg = signature.dilithium_signature.clone();
        signed_msg.extend_from_slice(message);

        let dil_signed_msg = dilithium3::SignedMessage::from_bytes(&signed_msg)
            .map_err(|_| KeypairError::InvalidSignature)?;
        dilithium3::open(&dil_signed_msg, &dil_pubkey)
            .map_err(|_| KeypairError::InvalidSignature)?;

        Ok(())
    }

    /// Returns the total size of the public key in bytes.
    pub fn size(&self) -> usize {
        32 + self.dilithium.len()
    }
}

impl std::fmt::Debug for PublicKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PublicKey")
            .field("ed25519", &hex::encode(&self.ed25519[..8]))
            .field("dilithium_len", &self.dilithium.len())
            .finish()
    }
}

/// Secret key (zeroized on drop for security).
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct SecretKey {
    /// Ed25519 secret key seed (32 bytes).
    ed25519_seed: [u8; 32],
    /// Master seed for key derivation.
    master_seed: [u8; 64],
}

impl SecretKey {
    /// Creates a new secret key from seeds.
    pub fn new(ed25519_seed: [u8; 32], master_seed: [u8; 64]) -> Self {
        SecretKey {
            ed25519_seed,
            master_seed,
        }
    }

    /// Returns the Ed25519 seed.
    pub fn ed25519_seed(&self) -> &[u8; 32] {
        &self.ed25519_seed
    }

    /// Returns the master seed.
    pub fn master_seed(&self) -> &[u8; 64] {
        &self.master_seed
    }
}

impl std::fmt::Debug for SecretKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SecretKey")
            .field("ed25519_seed", &"[REDACTED]")
            .field("master_seed", &"[REDACTED]")
            .finish()
    }
}

/// Hybrid keypair combining Ed25519 and Dilithium3.
pub struct HybridKeypair {
    ed25519: Ed25519Keypair,
    dilithium: Dilithium3Keypair,
    master_seed: [u8; 64],
}

impl HybridKeypair {
    /// Generates a new random hybrid keypair.
    pub fn generate() -> Self {
        let mut master_seed = [0u8; 64];
        rand::RngCore::fill_bytes(&mut OsRng, &mut master_seed);

        let ed25519 = Ed25519Keypair::generate();
        let dilithium = Dilithium3Keypair::generate();

        HybridKeypair {
            ed25519,
            dilithium,
            master_seed,
        }
    }

    /// Creates a keypair from a mnemonic phrase.
    pub fn from_mnemonic(mnemonic: &Mnemonic, passphrase: &str) -> Result<Self, KeypairError> {
        let seed = mnemonic.to_seed(passphrase);
        Self::from_seed(&seed)
    }

    /// Creates a keypair from a 64-byte seed.
    pub fn from_seed(seed: &[u8; 64]) -> Result<Self, KeypairError> {
        // Derive Ed25519 seed from first 32 bytes
        let mut ed_seed = [0u8; 32];
        ed_seed.copy_from_slice(&seed[..32]);
        let ed25519 = Ed25519Keypair::from_seed(&ed_seed);

        // Derive Dilithium seed from second 32 bytes
        let mut dil_seed = [0u8; 32];
        dil_seed.copy_from_slice(&seed[32..64]);
        let dilithium = Dilithium3Keypair::from_seed(&dil_seed);

        Ok(HybridKeypair {
            ed25519,
            dilithium,
            master_seed: *seed,
        })
    }

    /// Returns the public key.
    pub fn public_key(&self) -> PublicKey {
        PublicKey {
            ed25519: self.ed25519.public_key_bytes(),
            dilithium: self.dilithium.public_key_bytes(),
        }
    }

    /// Returns the secret key.
    pub fn secret_key(&self) -> SecretKey {
        SecretKey {
            ed25519_seed: self.ed25519.secret_key_bytes(),
            master_seed: self.master_seed,
        }
    }

    /// Signs a message with both Ed25519 and Dilithium3.
    pub fn sign(&self, message: &[u8]) -> HybridSignature {
        // Sign with Ed25519
        let ed_sig = self.ed25519.sign(message);

        // Sign with Dilithium3
        let dil_signed = self.dilithium.sign(message);

        // Extract just the signature part (without the message)
        let dil_sig_bytes = PqSignedMessage::as_bytes(&dil_signed);
        let sig_len = dil_sig_bytes.len() - message.len();
        let dil_sig = dil_sig_bytes[..sig_len].to_vec();

        HybridSignature::new(ed_sig.to_bytes(), dil_sig)
    }

    /// Returns the address for this keypair.
    pub fn address(&self, network: Network) -> Address {
        self.public_key().to_address(network)
    }

    /// Returns the Ed25519 keypair.
    pub fn ed25519(&self) -> &Ed25519Keypair {
        &self.ed25519
    }
}

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

/// Errors related to keypairs.
#[derive(Debug, Error)]
pub enum KeypairError {
    #[error("Invalid seed length")]
    InvalidSeedLength,

    #[error("Invalid public key")]
    InvalidPublicKey,

    #[error("Invalid secret key")]
    InvalidSecretKey,

    #[error("Invalid signature")]
    InvalidSignature,

    #[error("Key derivation failed: {0}")]
    DerivationFailed(String),
}

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

    #[test]
    fn test_ed25519_keypair() {
        let keypair = Ed25519Keypair::generate();
        let message = b"test message";
        let signature = keypair.sign(message);

        // Verify with the verifying key
        use ed25519_dalek::Verifier;
        assert!(keypair.verifying_key().verify(message, &signature).is_ok());
    }

    #[test]
    fn test_hybrid_keypair_generation() {
        let keypair = HybridKeypair::generate();
        let pubkey = keypair.public_key();

        assert_eq!(pubkey.ed25519_bytes().len(), 32);
        assert!(!pubkey.dilithium_bytes().is_empty());
    }

    #[test]
    fn test_hybrid_sign_verify() {
        let keypair = HybridKeypair::generate();
        let message = b"Hello, Synor blockchain!";

        let signature = keypair.sign(message);

        // Signature should have both components
        assert_eq!(signature.ed25519_signature.len(), 64);
        assert!(!signature.dilithium_signature.is_empty());
    }

    #[test]
    fn test_address_derivation() {
        let keypair = HybridKeypair::generate();

        let mainnet_addr = keypair.address(Network::Mainnet);
        let testnet_addr = keypair.address(Network::Testnet);

        assert!(mainnet_addr.to_string().starts_with("synor1"));
        assert!(testnet_addr.to_string().starts_with("tsynor1"));

        // Addresses should be different for different networks
        assert_ne!(mainnet_addr.to_string(), testnet_addr.to_string());
    }

    #[test]
    fn test_keypair_from_seed() {
        let mut seed = [0u8; 64];
        seed[0] = 1;
        seed[32] = 2;

        let keypair1 = HybridKeypair::from_seed(&seed).unwrap();
        let keypair2 = HybridKeypair::from_seed(&seed).unwrap();

        // Same seed should produce same Ed25519 public key
        assert_eq!(
            keypair1.public_key().ed25519_bytes(),
            keypair2.public_key().ed25519_bytes()
        );
    }
}
