//! Session manager for the vault crate (optional)
//! Provides ergonomic, stateful access to unlocked keyspaces and keypairs for interactive applications.
//! All state is local to the SessionManager instance. No global state.

use crate::{KVStore, KeyEntry, KeyspaceData, Vault, VaultError};
use std::collections::HashMap;
use zeroize::Zeroize;

/// SessionManager: Ergonomic, stateful wrapper over the Vault stateless API.
#[cfg(not(target_arch = "wasm32"))]
pub struct SessionManager<S: KVStore + Send + Sync> {
    // ... existing fields

    vault: Vault<S>,
    unlocked_keyspace: Option<(String, Vec<u8>, KeyspaceData)>, // (name, password, data)
    current_keypair: Option<String>,
}

#[cfg(target_arch = "wasm32")]
pub struct SessionManager<S: KVStore> {
    vault: Vault<S>,
    unlocked_keyspace: Option<(String, Vec<u8>, KeyspaceData)>, // (name, password, data)
    current_keypair: Option<String>,
}

#[cfg(target_arch = "wasm32")]
impl<S: KVStore> SessionManager<S> {
    pub fn get_vault_mut(&mut self) -> &mut Vault<S> {
        &mut self.vault
    }
}

#[cfg(not(target_arch = "wasm32"))]
impl<S: KVStore + Send + Sync> SessionManager<S> {
    pub fn new(vault: Vault<S>) -> Self {
        Self {
            vault,
            unlocked_keyspace: None,
            current_keypair: None,
        }
    }

    pub async fn create_keyspace(
        &mut self,
        name: &str,
        password: &[u8],
        tags: Option<Vec<String>>,
    ) -> Result<(), VaultError> {
        self.vault.create_keyspace(name, password, tags).await?;
        self.unlock_keyspace(name, password).await
    }

    pub async fn unlock_keyspace(&mut self, name: &str, password: &[u8]) -> Result<(), VaultError> {
        let data = self.vault.unlock_keyspace(name, password).await?;
        self.unlocked_keyspace = Some((name.to_string(), password.to_vec(), data));
        self.current_keypair = None;
        Ok(())
    }

    pub fn select_keypair(&mut self, key_id: &str) -> Result<(), VaultError> {
        let data = self
            .unlocked_keyspace
            .as_ref()
            .map(|(_, _, d)| d)
            .ok_or_else(|| VaultError::Crypto("No keyspace unlocked".to_string()))?;
        if data.keypairs.iter().any(|k| k.id == key_id) {
            self.current_keypair = Some(key_id.to_string());
            Ok(())
        } else {
            Err(VaultError::Crypto("Keypair not found".to_string()))
        }
    }

    pub async fn add_keypair(
        &mut self,
        key_type: Option<crate::KeyType>,
        metadata: Option<crate::KeyMetadata>,
    ) -> Result<String, VaultError> {
        let (name, password, _) = self
            .unlocked_keyspace
            .as_ref()
            .ok_or_else(|| VaultError::Crypto("No keyspace unlocked".to_string()))?;
        let id = self
            .vault
            .add_keypair(name, password, key_type, metadata.clone())
            .await?;
        let data = self.vault.unlock_keyspace(name, password).await?;
        self.unlocked_keyspace = Some((name.clone(), password.clone(), data));
        Ok(id)
    }

    pub fn list_keypairs(&self) -> Option<&[KeyEntry]> {
        self.current_keyspace().map(|ks| ks.keypairs.as_slice())
    }

    pub fn current_keyspace(&self) -> Option<&KeyspaceData> {
        self.unlocked_keyspace.as_ref().map(|(_, _, data)| data)
    }

    /// Returns the name of the currently unlocked keyspace, if any.
    pub fn current_keyspace_name(&self) -> Option<&str> {
        self.unlocked_keyspace
            .as_ref()
            .map(|(name, _, _)| name.as_str())
    }

    pub fn current_keypair(&self) -> Option<&KeyEntry> {
        let keyspace = self.current_keyspace()?;
        let key_id = self.current_keypair.as_ref()?;
        keyspace.keypairs.iter().find(|k| &k.id == key_id)
    }

    /// Returns the metadata of the current selected keypair, if any.
    pub fn current_keypair_metadata(&self) -> Option<crate::KeyMetadata> {
        self.current_keypair().and_then(|k| k.metadata.clone())
    }

    /// Returns the public key of the current selected keypair, if any.
    pub fn current_keypair_public_key(&self) -> Option<Vec<u8>> {
        self.current_keypair().map(|k| k.public_key.clone())
    }

    /// Returns true if a keyspace is currently unlocked.
    pub fn is_unlocked(&self) -> bool {
        self.unlocked_keyspace.is_some()
    }

    /// Returns the default keypair (first keypair) for client identity, if any.
    pub fn default_keypair(&self) -> Option<&KeyEntry> {
        self.current_keyspace()
            .and_then(|ks| ks.keypairs.first())
    }
    
    /// Selects the default keypair (first keypair) as the current keypair.
    pub fn select_default_keypair(&mut self) -> Result<(), VaultError> {
        let default_id = self
            .default_keypair()
            .map(|k| k.id.clone())
            .ok_or_else(|| VaultError::Crypto("No default keypair found".to_string()))?;
        
        self.select_keypair(&default_id)
    }
    
    /// Returns true if the current keypair is the default keypair (first keypair).
    pub fn is_default_keypair_selected(&self) -> bool {
        match (self.current_keypair(), self.default_keypair()) {
            (Some(current), Some(default)) => current.id == default.id,
            _ => false,
        }
    }

    pub async fn sign(&self, message: &[u8]) -> Result<Vec<u8>, VaultError> {
        let (name, password, _) = self
            .unlocked_keyspace
            .as_ref()
            .ok_or(VaultError::Crypto("No keyspace unlocked".to_string()))?;
        let keypair = self
            .current_keypair()
            .ok_or(VaultError::Crypto("No keypair selected".to_string()))?;
        self.vault.sign(name, password, &keypair.id, message).await
    }

    /// Verify a signature using the currently selected keypair
    pub async fn verify(&self, message: &[u8], signature: &[u8]) -> Result<bool, VaultError> {
        let (name, password, _) = self
            .unlocked_keyspace
            .as_ref()
            .ok_or(VaultError::Crypto("No keyspace unlocked".to_string()))?;
        let keypair = self
            .current_keypair()
            .ok_or(VaultError::Crypto("No keypair selected".to_string()))?;
        self.vault.verify(name, password, &keypair.id, message, signature).await
    }

    /// Encrypt data using the keyspace symmetric cipher
    /// Returns the encrypted data with the nonce prepended
    pub async fn encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>, VaultError> {
        let (name, password, _) = self
            .unlocked_keyspace
            .as_ref()
            .ok_or(VaultError::Crypto("No keyspace unlocked".to_string()))?;
        self.vault.encrypt(name, password, plaintext).await
    }

    /// Decrypt data using the keyspace symmetric cipher
    /// Expects the nonce to be prepended to the ciphertext (as returned by encrypt)
    pub async fn decrypt(&self, ciphertext: &[u8]) -> Result<Vec<u8>, VaultError> {
        let (name, password, _) = self
            .unlocked_keyspace
            .as_ref()
            .ok_or(VaultError::Crypto("No keyspace unlocked".to_string()))?;
        self.vault.decrypt(name, password, ciphertext).await
    }

    pub fn get_vault(&self) -> &Vault<S> {
        &self.vault
    }

    pub fn logout(&mut self) {
        if let Some((_, mut password, mut data)) = self.unlocked_keyspace.take() {
            password.zeroize();
            data.zeroize();
        }
        self.current_keypair = None;
    }
}

#[cfg(not(target_arch = "wasm32"))]
impl<S: KVStore + Send + Sync> Drop for SessionManager<S> {
    fn drop(&mut self) {
        self.logout();
    }
}

#[cfg(target_arch = "wasm32")]
impl<S: KVStore> SessionManager<S> {
    pub fn new(vault: Vault<S>) -> Self {
        Self {
            vault,
            unlocked_keyspace: None,
            current_keypair: None,
        }
    }

    pub async fn create_keyspace(
        &mut self,
        name: &str,
        password: &[u8],
        tags: Option<Vec<String>>,
    ) -> Result<(), VaultError> {
        self.vault.create_keyspace(name, password, tags).await?;
        self.unlock_keyspace(name, password).await
    }

    pub async fn unlock_keyspace(&mut self, name: &str, password: &[u8]) -> Result<(), VaultError> {
        let data = self.vault.unlock_keyspace(name, password).await?;
        self.unlocked_keyspace = Some((name.to_string(), password.to_vec(), data));
        self.current_keypair = None;
        Ok(())
    }

    pub fn select_keypair(&mut self, key_id: &str) -> Result<(), VaultError> {
        let data = self
            .unlocked_keyspace
            .as_ref()
            .map(|(_, _, d)| d)
            .ok_or_else(|| VaultError::Crypto("No keyspace unlocked".to_string()))?;
        if data.keypairs.iter().any(|k| k.id == key_id) {
            self.current_keypair = Some(key_id.to_string());
            Ok(())
        } else {
            Err(VaultError::Crypto("Keypair not found".to_string()))
        }
    }

    pub async fn add_keypair(
        &mut self,
        key_type: Option<crate::KeyType>,
        metadata: Option<crate::KeyMetadata>,
    ) -> Result<String, VaultError> {
        let (name, password, _) = self
            .unlocked_keyspace
            .as_ref()
            .ok_or_else(|| VaultError::Crypto("No keyspace unlocked".to_string()))?;
        let id = self
            .vault
            .add_keypair(name, password, key_type, metadata.clone())
            .await?;
        let data = self.vault.unlock_keyspace(name, password).await?;
        self.unlocked_keyspace = Some((name.clone(), password.clone(), data));
        Ok(id)
    }

    pub fn list_keypairs(&self) -> Option<&[KeyEntry]> {
        self.current_keyspace().map(|ks| ks.keypairs.as_slice())
    }

    pub fn current_keyspace(&self) -> Option<&KeyspaceData> {
        self.unlocked_keyspace.as_ref().map(|(_, _, data)| data)
    }

    /// Returns the name of the currently unlocked keyspace, if any.
    pub fn current_keyspace_name(&self) -> Option<&str> {
        self.unlocked_keyspace
            .as_ref()
            .map(|(name, _, _)| name.as_str())
    }

    pub fn current_keypair(&self) -> Option<&KeyEntry> {
        let keyspace = self.current_keyspace()?;
        let key_id = self.current_keypair.as_ref()?;
        keyspace.keypairs.iter().find(|k| &k.id == key_id)
    }

    /// Returns the metadata of the current selected keypair, if any.
    pub fn current_keypair_metadata(&self) -> Option<crate::KeyMetadata> {
        self.current_keypair().and_then(|k| k.metadata.clone())
    }

    /// Returns the public key of the current selected keypair, if any.
    pub fn current_keypair_public_key(&self) -> Option<Vec<u8>> {
        self.current_keypair().map(|k| k.public_key.clone())
    }

    /// Returns true if a keyspace is currently unlocked.
    pub fn is_unlocked(&self) -> bool {
        self.unlocked_keyspace.is_some()
    }

    /// Returns the default keypair (first keypair) for client identity, if any.
    pub fn default_keypair(&self) -> Option<&KeyEntry> {
        self.current_keyspace()
            .and_then(|ks| ks.keypairs.first())
    }
    
    /// Selects the default keypair (first keypair) as the current keypair.
    pub fn select_default_keypair(&mut self) -> Result<(), VaultError> {
        let default_id = self
            .default_keypair()
            .map(|k| k.id.clone())
            .ok_or_else(|| VaultError::Crypto("No default keypair found".to_string()))?;
        
        self.select_keypair(&default_id)
    }
    
    /// Returns true if the current keypair is the default keypair (first keypair).
    pub fn is_default_keypair_selected(&self) -> bool {
        match (self.current_keypair(), self.default_keypair()) {
            (Some(current), Some(default)) => current.id == default.id,
            _ => false,
        }
    }

    pub async fn sign(&self, message: &[u8]) -> Result<Vec<u8>, VaultError> {
        let (name, password, _) = self
            .unlocked_keyspace
            .as_ref()
            .ok_or(VaultError::Crypto("No keyspace unlocked".to_string()))?;
        let keypair = self
            .current_keypair()
            .ok_or(VaultError::Crypto("No keypair selected".to_string()))?;
        self.vault.sign(name, password, &keypair.id, message).await
    }

    /// Verify a signature using the currently selected keypair
    pub async fn verify(&self, message: &[u8], signature: &[u8]) -> Result<bool, VaultError> {
        let (name, password, _) = self
            .unlocked_keyspace
            .as_ref()
            .ok_or(VaultError::Crypto("No keyspace unlocked".to_string()))?;
        let keypair = self
            .current_keypair()
            .ok_or(VaultError::Crypto("No keypair selected".to_string()))?;
        self.vault.verify(name, password, &keypair.id, message, signature).await
    }

    /// Encrypt data using the keyspace symmetric cipher
    /// Returns the encrypted data with the nonce prepended
    pub async fn encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>, VaultError> {
        let (name, password, _) = self
            .unlocked_keyspace
            .as_ref()
            .ok_or(VaultError::Crypto("No keyspace unlocked".to_string()))?;
        self.vault.encrypt(name, password, plaintext).await
    }

    /// Decrypt data using the keyspace symmetric cipher
    /// Expects the nonce to be prepended to the ciphertext (as returned by encrypt)
    pub async fn decrypt(&self, ciphertext: &[u8]) -> Result<Vec<u8>, VaultError> {
        let (name, password, _) = self
            .unlocked_keyspace
            .as_ref()
            .ok_or(VaultError::Crypto("No keyspace unlocked".to_string()))?;
        self.vault.decrypt(name, password, ciphertext).await
    }

    pub fn get_vault(&self) -> &Vault<S> {
        &self.vault
    }

    pub fn logout(&mut self) {
        if let Some((_, mut password, mut data)) = self.unlocked_keyspace.take() {
            password.zeroize();
            data.zeroize();
        }
        self.current_keypair = None;
    }
}

#[cfg(target_arch = "wasm32")]
impl<S: KVStore> Drop for SessionManager<S> {
    fn drop(&mut self) {
        self.logout();
    }
}