use once_cell::sync::Lazy;
use std::sync::Mutex;

use crate::vault::error::CryptoError;
use crate::vault::keyspace::keypair_types::{KeyPair, KeySpace}; // Assuming KeyPair and KeySpace will be in keypair_types.rs

/// 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.
pub static SESSION: Lazy<Mutex<Session>> = Lazy::new(|| Mutex::new(Session::default()));

// Session management and selected keypair operation functions will be added here
/// 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(name.to_string()));
        }

        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(name.to_string()));
        }

        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(keypair_name.clone()));
        }
        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)
}