//! Core implementation of keypair functionality.

use k256::ecdsa::{SigningKey, VerifyingKey, signature::{Signer, Verifier}, Signature};
use rand::rngs::OsRng;
use serde::{Serialize, Deserialize};
use std::collections::HashMap;
use once_cell::sync::Lazy;
use std::sync::Mutex;
use sha2::{Sha256, Digest};

use super::error::CryptoError;

/// A keypair for signing and verifying messages.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeyPair {
    pub name: String,
    #[serde(with = "verifying_key_serde")]
    pub verifying_key: VerifyingKey,
    #[serde(with = "signing_key_serde")]
    pub signing_key: SigningKey,
}

// Serialization helpers for VerifyingKey
mod verifying_key_serde {
    use super::*;
    use serde::{Serializer, Deserializer};
    use serde::de::{self, Visitor};
    use std::fmt;

    pub fn serialize<S>(key: &VerifyingKey, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let bytes = key.to_sec1_bytes();
        serializer.serialize_bytes(&bytes)
    }

    struct VerifyingKeyVisitor;

    impl<'de> Visitor<'de> for VerifyingKeyVisitor {
        type Value = VerifyingKey;

        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
            formatter.write_str("a byte array representing a verifying key")
        }

        fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
        where
            E: de::Error,
        {
            VerifyingKey::from_sec1_bytes(v).map_err(|_| E::custom("invalid verifying key"))
        }
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<VerifyingKey, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_bytes(VerifyingKeyVisitor)
    }
}

// Serialization helpers for SigningKey
mod signing_key_serde {
    use super::*;
    use serde::{Serializer, Deserializer};
    use serde::de::{self, Visitor};
    use std::fmt;

    pub fn serialize<S>(key: &SigningKey, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let bytes = key.to_bytes();
        serializer.serialize_bytes(&bytes)
    }

    struct SigningKeyVisitor;

    impl<'de> Visitor<'de> for SigningKeyVisitor {
        type Value = SigningKey;

        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
            formatter.write_str("a byte array representing a signing key")
        }

        fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
        where
            E: de::Error,
        {
            SigningKey::from_bytes(v.into()).map_err(|_| E::custom("invalid signing key"))
        }
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<SigningKey, D::Error>
    where
        D: Deserializer<'de>,
    {
        deserializer.deserialize_bytes(SigningKeyVisitor)
    }
}

impl KeyPair {
    /// Creates a new keypair with the given name.
    pub fn new(name: &str) -> Self {
        let signing_key = SigningKey::random(&mut OsRng);
        let verifying_key = VerifyingKey::from(&signing_key);
        
        KeyPair {
            name: name.to_string(),
            verifying_key,
            signing_key,
        }
    }

    /// Gets the public key bytes.
    pub fn pub_key(&self) -> Vec<u8> {
        self.verifying_key.to_sec1_bytes().to_vec()
    }
    
    /// Derives a public key from a private key.
    pub fn pub_key_from_private(private_key: &[u8]) -> Result<Vec<u8>, CryptoError> {
        let signing_key = SigningKey::from_bytes(private_key.into())
            .map_err(|_| CryptoError::InvalidKeyLength)?;
        let verifying_key = VerifyingKey::from(&signing_key);
        Ok(verifying_key.to_sec1_bytes().to_vec())
    }

    /// Signs a message.
    pub fn sign(&self, message: &[u8]) -> Vec<u8> {
        let signature: Signature = self.signing_key.sign(message);
        signature.to_bytes().to_vec()
    }

    /// Verifies a message signature.
    pub fn verify(&self, message: &[u8], signature_bytes: &[u8]) -> Result<bool, CryptoError> {
        let signature = Signature::from_bytes(signature_bytes.into())
            .map_err(|_| CryptoError::SignatureFormatError)?;
            
        match self.verifying_key.verify(message, &signature) {
            Ok(_) => Ok(true),
            Err(_) => Ok(false), // Verification failed, but operation was successful
        }
    }
    
    /// Verifies a message signature using only a public key.
    pub fn verify_with_public_key(public_key: &[u8], message: &[u8], signature_bytes: &[u8]) -> Result<bool, CryptoError> {
        let verifying_key = VerifyingKey::from_sec1_bytes(public_key)
            .map_err(|_| CryptoError::InvalidKeyLength)?;
        
        let signature = Signature::from_bytes(signature_bytes.into())
            .map_err(|_| CryptoError::SignatureFormatError)?;
            
        match verifying_key.verify(message, &signature) {
            Ok(_) => Ok(true),
            Err(_) => Ok(false), // Verification failed, but operation was successful
        }
    }
    
    /// Encrypts a message using the recipient's public key.
    /// This implements ECIES (Elliptic Curve Integrated Encryption Scheme):
    /// 1. Generate an ephemeral keypair
    /// 2. Derive a shared secret using ECDH
    /// 3. Derive encryption key from the shared secret
    /// 4. Encrypt the message using symmetric encryption
    /// 5. Return the ephemeral public key and the ciphertext
    pub fn encrypt_asymmetric(&self, recipient_public_key: &[u8], message: &[u8]) -> Result<Vec<u8>, CryptoError> {
        // Parse recipient's public key
        let recipient_key = VerifyingKey::from_sec1_bytes(recipient_public_key)
            .map_err(|_| CryptoError::InvalidKeyLength)?;
        
        // Generate ephemeral keypair
        let ephemeral_signing_key = SigningKey::random(&mut OsRng);
        let ephemeral_public_key = VerifyingKey::from(&ephemeral_signing_key);
        
        // Derive shared secret (this is a simplified ECDH)
        // In a real implementation, we would use proper ECDH, but for this example:
        let shared_point = recipient_key.to_encoded_point(false);
        let shared_secret = {
            let mut hasher = Sha256::default();
            hasher.update(ephemeral_signing_key.to_bytes());
            hasher.update(shared_point.as_bytes());
            hasher.finalize().to_vec()
        };
        
        // Encrypt the message using the derived key
        let ciphertext = super::symmetric::encrypt_with_key(&shared_secret, message)
            .map_err(|_| CryptoError::EncryptionFailed)?;
        
        // Format: ephemeral_public_key || ciphertext
        let mut result = ephemeral_public_key.to_sec1_bytes().to_vec();
        result.extend_from_slice(&ciphertext);
        
        Ok(result)
    }
    
    /// Decrypts a message using the recipient's private key.
    /// This is the counterpart to encrypt_asymmetric.
    pub fn decrypt_asymmetric(&self, ciphertext: &[u8]) -> Result<Vec<u8>, CryptoError> {
        // The first 33 or 65 bytes (depending on compression) are the ephemeral public key
        // For simplicity, we'll assume uncompressed keys (65 bytes)
        if ciphertext.len() <= 65 {
            return Err(CryptoError::DecryptionFailed);
        }
        
        // Extract ephemeral public key and actual ciphertext
        let ephemeral_public_key = &ciphertext[..65];
        let actual_ciphertext = &ciphertext[65..];
        
        // Parse ephemeral public key
        let sender_key = VerifyingKey::from_sec1_bytes(ephemeral_public_key)
            .map_err(|_| CryptoError::InvalidKeyLength)?;
        
        // Derive shared secret (simplified ECDH)
        let shared_point = sender_key.to_encoded_point(false);
        let shared_secret = {
            let mut hasher = Sha256::default();
            hasher.update(self.signing_key.to_bytes());
            hasher.update(shared_point.as_bytes());
            hasher.finalize().to_vec()
        };
        
        // Decrypt the message using the derived key
        super::symmetric::decrypt_with_key(&shared_secret, actual_ciphertext)
            .map_err(|_| CryptoError::DecryptionFailed)
    }
}

/// A collection of keypairs.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct KeySpace {
    pub name: String,
    pub keypairs: HashMap<String, KeyPair>,
}

impl KeySpace {
    /// Creates a new key space with the given name.
    pub fn new(name: &str) -> Self {
        KeySpace {
            name: name.to_string(),
            keypairs: HashMap::new(),
        }
    }

    /// Adds a new keypair to the space.
    pub fn add_keypair(&mut self, name: &str) -> Result<(), CryptoError> {
        if self.keypairs.contains_key(name) {
            return Err(CryptoError::KeypairAlreadyExists);
        }
        
        let keypair = KeyPair::new(name);
        self.keypairs.insert(name.to_string(), keypair);
        Ok(())
    }

    /// Gets a keypair by name.
    pub fn get_keypair(&self, name: &str) -> Result<&KeyPair, CryptoError> {
        self.keypairs.get(name).ok_or(CryptoError::KeypairNotFound)
    }

    /// Lists all keypair names in the space.
    pub fn list_keypairs(&self) -> Vec<String> {
        self.keypairs.keys().cloned().collect()
    }
}

/// Session state for the current key space and selected keypair.
pub struct Session {
    pub current_space: Option<KeySpace>,
    pub selected_keypair: Option<String>,
}

impl Default for Session {
    fn default() -> Self {
        Session {
            current_space: None,
            selected_keypair: None,
        }
    }
}

/// Global session state.
static SESSION: Lazy<Mutex<Session>> = Lazy::new(|| {
    Mutex::new(Session::default())
});

/// Creates a new key space with the given name.
pub fn create_space(name: &str) -> Result<(), CryptoError> {
    let mut session = SESSION.lock().unwrap();
    
    // Create a new space
    let space = KeySpace::new(name);
    
    // Set as current space
    session.current_space = Some(space);
    session.selected_keypair = None;
    
    Ok(())
}

/// Sets the current key space.
pub fn set_current_space(space: KeySpace) -> Result<(), CryptoError> {
    let mut session = SESSION.lock().unwrap();
    session.current_space = Some(space);
    session.selected_keypair = None;
    Ok(())
}

/// Gets the current key space.
pub fn get_current_space() -> Result<KeySpace, CryptoError> {
    let session = SESSION.lock().unwrap();
    session.current_space.clone().ok_or(CryptoError::NoActiveSpace)
}

/// Clears the current session (logout).
pub fn clear_session() {
    let mut session = SESSION.lock().unwrap();
    session.current_space = None;
    session.selected_keypair = None;
}

/// Creates a new keypair in the current space.
pub fn create_keypair(name: &str) -> Result<(), CryptoError> {
    let mut session = SESSION.lock().unwrap();
    
    if let Some(ref mut space) = session.current_space {
        if space.keypairs.contains_key(name) {
            return Err(CryptoError::KeypairAlreadyExists);
        }
        
        let keypair = KeyPair::new(name);
        space.keypairs.insert(name.to_string(), keypair);
        
        // Automatically select the new keypair
        session.selected_keypair = Some(name.to_string());
        
        Ok(())
    } else {
        Err(CryptoError::NoActiveSpace)
    }
}

/// Selects a keypair for use.
pub fn select_keypair(name: &str) -> Result<(), CryptoError> {
    let mut session = SESSION.lock().unwrap();
    
    if let Some(ref space) = session.current_space {
        if !space.keypairs.contains_key(name) {
            return Err(CryptoError::KeypairNotFound);
        }
        
        session.selected_keypair = Some(name.to_string());
        Ok(())
    } else {
        Err(CryptoError::NoActiveSpace)
    }
}

/// Gets the currently selected keypair.
pub fn get_selected_keypair() -> Result<KeyPair, CryptoError> {
    let session = SESSION.lock().unwrap();
    
    if let Some(ref space) = session.current_space {
        if let Some(ref keypair_name) = session.selected_keypair {
            if let Some(keypair) = space.keypairs.get(keypair_name) {
                return Ok(keypair.clone());
            }
            return Err(CryptoError::KeypairNotFound);
        }
        return Err(CryptoError::NoKeypairSelected);
    }
    
    Err(CryptoError::NoActiveSpace)
}

/// Lists all keypair names in the current space.
pub fn list_keypairs() -> Result<Vec<String>, CryptoError> {
    let session = SESSION.lock().unwrap();
    
    if let Some(ref space) = session.current_space {
        Ok(space.keypairs.keys().cloned().collect())
    } else {
        Err(CryptoError::NoActiveSpace)
    }
}

/// Gets the public key of the selected keypair.
pub fn keypair_pub_key() -> Result<Vec<u8>, CryptoError> {
    let keypair = get_selected_keypair()?;
    Ok(keypair.pub_key())
}

/// Derives a public key from a private key.
pub fn derive_public_key(private_key: &[u8]) -> Result<Vec<u8>, CryptoError> {
    KeyPair::pub_key_from_private(private_key)
}

/// Signs a message with the selected keypair.
pub fn keypair_sign(message: &[u8]) -> Result<Vec<u8>, CryptoError> {
    let keypair = get_selected_keypair()?;
    Ok(keypair.sign(message))
}

/// Verifies a message signature with the selected keypair.
pub fn keypair_verify(message: &[u8], signature_bytes: &[u8]) -> Result<bool, CryptoError> {
    let keypair = get_selected_keypair()?;
    keypair.verify(message, signature_bytes)
}

/// Verifies a message signature with a public key.
pub fn verify_with_public_key(public_key: &[u8], message: &[u8], signature_bytes: &[u8]) -> Result<bool, CryptoError> {
    KeyPair::verify_with_public_key(public_key, message, signature_bytes)
}

/// Encrypts a message for a recipient using their public key.
pub fn encrypt_asymmetric(recipient_public_key: &[u8], message: &[u8]) -> Result<Vec<u8>, CryptoError> {
    let keypair = get_selected_keypair()?;
    keypair.encrypt_asymmetric(recipient_public_key, message)
}

/// Decrypts a message that was encrypted with the current keypair's public key.
pub fn decrypt_asymmetric(ciphertext: &[u8]) -> Result<Vec<u8>, CryptoError> {
    let keypair = get_selected_keypair()?;
    keypair.decrypt_asymmetric(ciphertext)
}