//! vault: Cryptographic keyspace and operations

//! vault: Cryptographic keyspace and operations

pub mod data;
pub use crate::data::{KeyEntry, KeyMetadata, KeyType};
pub use crate::session::SessionManager;
mod crypto;
mod error;
pub mod rhai_bindings;
mod rhai_sync_helpers;
pub mod session;
mod utils;

#[cfg(target_arch = "wasm32")]
pub mod session_singleton;
#[cfg(target_arch = "wasm32")]
pub mod wasm_helpers;

use crate::crypto::kdf;
use crate::crypto::random_salt;
use data::*;
use error::VaultError;
pub use kvstore::traits::KVStore;

use crate::crypto::cipher::{decrypt_chacha20, encrypt_chacha20};
use signature::SignatureEncoding;
// TEMP: File-based debug logger for crypto troubleshooting
use log::debug;

/// Vault: Cryptographic keyspace and operations
pub struct Vault<S: KVStore> {
    storage: S,
    // Optionally: cache of unlocked keyspaces, etc.
}

/// Helper to encrypt and prepend nonce to ciphertext for keyspace storage
/// Helper to encrypt and prepend nonce to ciphertext for keyspace storage
/// Always uses ChaCha20Poly1305.
fn encrypt_with_nonce_prepended(key: &[u8], plaintext: &[u8]) -> Result<Vec<u8>, VaultError> {
    let nonce = random_salt(12);
    debug!("nonce: {}", hex::encode(&nonce));
    // Always use ChaCha20Poly1305 for encryption
    let ct = encrypt_chacha20(key, plaintext, &nonce).map_err(|e| VaultError::Crypto(e))?;
    debug!("ct: {}", hex::encode(&ct));
    debug!("key: {}", hex::encode(key));
    let mut blob = nonce.clone();
    blob.extend_from_slice(&ct);
    debug!("ENCRYPTED (nonce|ct): {}", hex::encode(&blob));
    Ok(blob)
}

impl<S: KVStore> Vault<S> {
    pub fn new(storage: S) -> Self {
        Self { storage }
    }

    /// Create a new keyspace with the given name, password, and options.
    /// Create a new keyspace with the given name and password. Always uses PBKDF2 and ChaCha20Poly1305.
    pub async fn create_keyspace(
        &mut self,
        name: &str,
        password: &[u8],
        tags: Option<Vec<String>>,
    ) -> Result<(), VaultError> {
        // Check if keyspace already exists
        if self
            .storage
            .get(name)
            .await
            .map_err(|e| VaultError::Storage(format!("{e:?}")))?            
            .is_some()
        {
            debug!("keyspace '{}' already exists", name);
            return Err(VaultError::Crypto("Keyspace already exists".to_string()));
        }
        debug!("entry: name={}", name);
        use crate::crypto::{kdf, random_salt};
        use crate::data::{KeyspaceData, KeyspaceMetadata};
        use serde_json;

        // 1. Generate salt
        let salt = random_salt(16);
        debug!("salt: {:?}", salt);
        // 2. Derive key
        // Always use PBKDF2 for key derivation
        let key = kdf::keyspace_key(password, &salt);
        debug!("derived key: {} bytes", key.len());
        // 3. Prepare initial keyspace data
        let keyspace_data = KeyspaceData { keypairs: vec![] };
        let plaintext = match serde_json::to_vec(&keyspace_data) {
            Ok(val) => val,
            Err(e) => {
                debug!("serde_json error: {}", e);
                return Err(VaultError::Serialization(e.to_string()));
            }
        };
        debug!("plaintext serialized: {} bytes", plaintext.len());
        // 4. Generate nonce (12 bytes for both ciphers)
        let nonce = random_salt(12);
        debug!("nonce: {}", hex::encode(&nonce));
        // 5. Encrypt
        // Always use ChaCha20Poly1305 for encryption
        let encrypted_blob = encrypt_with_nonce_prepended(&key, &plaintext)?;
        debug!("encrypted_blob: {} bytes", encrypted_blob.len());
        debug!("encrypted_blob (hex): {}", hex::encode(&encrypted_blob));
        // 6. Compose metadata
        let metadata = KeyspaceMetadata {
            name: name.to_string(),
            salt: salt.clone().try_into().unwrap_or([0u8; 16]),
            encrypted_blob,
            created_at: Some(crate::utils::now()),
            tags,
        };
        // 7. Store in kvstore (keyed by keyspace name)
        let meta_bytes = match serde_json::to_vec(&metadata) {
            Ok(val) => val,
            Err(e) => {
                debug!("serde_json metadata error: {}", e);
                return Err(VaultError::Serialization(e.to_string()));
            }
        };
        self.storage
            .set(name, &meta_bytes)
            .await
            .map_err(|e| VaultError::Storage(format!("{e:?}")))?;
        debug!("success");
        
        // 8. Create default keypair, passing the salt we already have
        self.create_default_keypair(name, password, &salt).await?;
        
        Ok(())
    }

    /// List all keyspaces (metadata only, not decrypted)
    pub async fn list_keyspaces(&self) -> Result<Vec<KeyspaceMetadata>, VaultError> {
        use serde_json;
        // 1. List all keys in kvstore
        let keys = self
            .storage
            .keys()
            .await
            .map_err(|e| VaultError::Storage(format!("{e:?}")))?;
        let mut keyspaces = Vec::new();
        for key in keys {
            if let Some(bytes) = self
                .storage
                .get(&key)
                .await
                .map_err(|e| VaultError::Storage(format!("{e:?}")))?
            {
                if let Ok(meta) = serde_json::from_slice::<KeyspaceMetadata>(&bytes) {
                    keyspaces.push(meta);
                }
            }
        }
        Ok(keyspaces)
    }

    /// Unlock a keyspace by name and password, returning the decrypted data
    /// Unlock a keyspace by name and password, returning the decrypted data
    /// Always uses PBKDF2 and ChaCha20Poly1305.
    pub async fn unlock_keyspace(
        &self,
        name: &str,
        password: &[u8],
    ) -> Result<KeyspaceData, VaultError> {
        debug!("unlock_keyspace entry: name={}", name);
        // use crate::crypto::kdf; // removed if not needed
        use serde_json;
        // 1. Fetch keyspace metadata
        let meta_bytes = self
            .storage
            .get(name)
            .await
            .map_err(|e| VaultError::Storage(format!("{e:?}")))?;
        let meta_bytes = meta_bytes.ok_or(VaultError::KeyspaceNotFound(name.to_string()))?;
        let metadata: KeyspaceMetadata = serde_json::from_slice(&meta_bytes)
            .map_err(|e| VaultError::Serialization(e.to_string()))?;
        if metadata.salt.len() != 16 {
            debug!("salt length {} != 16", metadata.salt.len());
            return Err(VaultError::Crypto(
                "Salt length must be 16 bytes".to_string(),
            ));
        }
        // 2. Derive key
        let key = kdf::keyspace_key(password, &metadata.salt);
        debug!("derived key: {} bytes", key.len());

        let ciphertext = &metadata.encrypted_blob;
        if ciphertext.len() < 12 {
            debug!("ciphertext too short: {}", ciphertext.len());
            return Err(VaultError::Crypto("Ciphertext too short".to_string()));
        }

        let (nonce, ct) = ciphertext.split_at(12);
        debug!("nonce: {}", hex::encode(nonce));
        let plaintext = decrypt_chacha20(&key, ct, nonce).map_err(VaultError::Crypto)?;
        debug!("plaintext decrypted: {} bytes", plaintext.len());
        // 4. Deserialize keyspace data
        let keyspace_data: KeyspaceData = match serde_json::from_slice(&plaintext) {
            Ok(val) => val,
            Err(e) => {
                debug!("serde_json data error: {}", e);
                return Err(VaultError::Serialization(e.to_string()));
            }
        };
        debug!("success");
        Ok(keyspace_data)
    }

    /// Lock a keyspace (remove from cache, if any)
    /// Lock a keyspace (remove from cache, if any)
    pub fn lock_keyspace(&mut self, _name: &str) {
        // Optional: clear from in-memory cache
    }

    // --- Keypair Management APIs ---

    /// Create a default Ed25519 keypair for client identity
    /// This keypair is deterministically generated from the password and salt
    /// and will always be the first keypair in the keyspace
    async fn create_default_keypair(
        &mut self,
        keyspace: &str,
        password: &[u8],
        salt: &[u8],
    ) -> Result<String, VaultError> {
        // 1. Derive a deterministic seed using standard PBKDF2
        let seed = kdf::keyspace_key(password, salt);
        
        // 2. Generate Ed25519 keypair from the seed
        use ed25519_dalek::{SigningKey, VerifyingKey};
        
        // Use the seed to create a deterministic keypair
        let signing = SigningKey::from_bytes(seed.as_slice().try_into().unwrap());
        let verifying: VerifyingKey = (&signing).into();
        
        let priv_bytes = signing.to_bytes().to_vec();
        let pub_bytes = verifying.to_bytes().to_vec();
        
        // Create an ID for the default keypair
        let id = hex::encode(&pub_bytes);
        
        // 3. Unlock the keyspace to get its data
        let mut data = self.unlock_keyspace(keyspace, password).await?;
        
        // 4. Add to keypairs (as the first entry)
        let entry = KeyEntry {
            id: id.clone(),
            key_type: KeyType::Ed25519,
            private_key: priv_bytes,
            public_key: pub_bytes,
            metadata: Some(KeyMetadata {
                name: Some("Default Identity".to_string()),
                created_at: Some(crate::utils::now()),
                tags: Some(vec!["default".to_string(), "identity".to_string()]),
            }),
        };
        
        // Ensure it's the first keypair by inserting at index 0
        data.keypairs.insert(0, entry);
        
        // 5. Re-encrypt and store
        self.save_keyspace(keyspace, password, &data).await?;
        
        Ok(id)
    }

    /// Add a new keypair to a keyspace (generates and stores a new keypair)
    /// Add a new keypair to a keyspace (generates and stores a new keypair)
    /// If key_type is None, defaults to Secp256k1.
    pub async fn add_keypair(
        &mut self,
        keyspace: &str,
        password: &[u8],
        key_type: Option<KeyType>,
        metadata: Option<KeyMetadata>,
    ) -> Result<String, VaultError> {
        use crate::data::KeyEntry;
        use rand_core::OsRng;
        use rand_core::RngCore;

        // 1. Unlock keyspace
        let mut data = self.unlock_keyspace(keyspace, password).await?;
        // 2. Generate keypair
        let key_type = key_type.unwrap_or(KeyType::Secp256k1);
        let (private_key, public_key, id) = match key_type {
            KeyType::Ed25519 => {
                use ed25519_dalek::{SigningKey, VerifyingKey};
                let mut bytes = [0u8; 32];
                OsRng.fill_bytes(&mut bytes);
                let signing = SigningKey::from_bytes(&bytes);
                let verifying: VerifyingKey = (&signing).into();
                let priv_bytes = signing.to_bytes().to_vec();
                let pub_bytes = verifying.to_bytes().to_vec();
                let id = hex::encode(&pub_bytes);
                (priv_bytes, pub_bytes, id)
            }
            KeyType::Secp256k1 => {
                use k256::ecdsa::SigningKey;

                let sk = SigningKey::random(&mut OsRng);
                let pk = sk.verifying_key();
                let priv_bytes = sk.to_bytes().to_vec();
                let pub_bytes = pk.to_encoded_point(false).as_bytes().to_vec();
                let id = hex::encode(&pub_bytes);
                (priv_bytes, pub_bytes, id)
            }
        };
        // 3. Add to keypairs
        let entry = KeyEntry {
            id: id.clone(),
            key_type,
            private_key,
            public_key,
            metadata,
        };
        data.keypairs.push(entry);
        // 4. Re-encrypt and store
        self.save_keyspace(keyspace, password, &data).await?;
        Ok(id)
    }

    /// Remove a keypair by id from a keyspace
    pub async fn remove_keypair(
        &mut self,
        keyspace: &str,
        password: &[u8],
        key_id: &str,
    ) -> Result<(), VaultError> {
        let mut data = self.unlock_keyspace(keyspace, password).await?;
        data.keypairs.retain(|k| k.id != key_id);
        self.save_keyspace(keyspace, password, &data).await
    }

    /// List all keypairs in a keyspace (public info only)
    pub async fn list_keypairs(
        &self,
        keyspace: &str,
        password: &[u8],
    ) -> Result<Vec<(String, KeyType)>, VaultError> {
        let data = self.unlock_keyspace(keyspace, password).await?;
        Ok(data
            .keypairs
            .iter()
            .map(|k| (k.id.clone(), k.key_type.clone()))
            .collect())
    }

    /// Export a keypair's private and public key by id
    pub async fn export_keypair(
        &self,
        keyspace: &str,
        password: &[u8],
        key_id: &str,
    ) -> Result<(Vec<u8>, Vec<u8>), VaultError> {
        let data = self.unlock_keyspace(keyspace, password).await?;
        let key = data
            .keypairs
            .iter()
            .find(|k| k.id == key_id)
            .ok_or(VaultError::KeyNotFound(key_id.to_string()))?;
        Ok((key.private_key.clone(), key.public_key.clone()))
    }

    /// Save the updated keyspace data (helper)
    async fn save_keyspace(
        &mut self,
        keyspace: &str,
        password: &[u8],
        data: &KeyspaceData,
    ) -> Result<(), VaultError> {
        debug!("save_keyspace entry: keyspace={}", keyspace);
        use crate::crypto::kdf;
        use serde_json;

        let meta_bytes = self
            .storage
            .get(keyspace)
            .await
            .map_err(|e| VaultError::Storage(format!("{e:?}")))?;
        debug!(
            "got meta_bytes: {}",
            meta_bytes.as_ref().map(|v| v.len()).unwrap_or(0)
        );
        let meta_bytes = meta_bytes.ok_or(VaultError::KeyspaceNotFound(keyspace.to_string()))?;
        let mut metadata: KeyspaceMetadata = serde_json::from_slice(&meta_bytes)
            .map_err(|e| VaultError::Serialization(e.to_string()))?;
        debug!("metadata: salt={:?}", metadata.salt);
        if metadata.salt.len() != 16 {
            debug!("salt length {} != 16", metadata.salt.len());
            return Err(VaultError::Crypto(
                "Salt length must be 16 bytes".to_string(),
            ));
        }
        // 2. Derive key
        let key = kdf::keyspace_key(password, &metadata.salt);
        debug!("derived key: {} bytes", key.len());
        // 3. Serialize plaintext
        let plaintext = match serde_json::to_vec(data) {
            Ok(val) => val,
            Err(e) => {
                debug!("serde_json data error: {}", e);
                return Err(VaultError::Serialization(e.to_string()));
            }
        };
        debug!("plaintext serialized: {} bytes", plaintext.len());
        // 4. Generate nonce
        let nonce = random_salt(12);
        debug!("nonce: {}", hex::encode(&nonce));
        // 5. Encrypt
        let encrypted_blob = encrypt_with_nonce_prepended(&key, &plaintext)?;
        debug!("encrypted_blob: {} bytes", encrypted_blob.len());
        // 6. Store new encrypted blob
        metadata.encrypted_blob = encrypted_blob;
        let meta_bytes = match serde_json::to_vec(&metadata) {
            Ok(val) => val,
            Err(e) => {
                debug!("serde_json metadata error: {}", e);
                return Err(VaultError::Serialization(e.to_string()));
            }
        };
        self.storage
            .set(keyspace, &meta_bytes)
            .await
            .map_err(|e| VaultError::Storage(format!("{e:?}")))?;
        debug!("success");
        Ok(())
    }

    /// Sign a message with a stored keypair in a keyspace
    ///
    /// # Arguments
    /// * `keyspace` - Keyspace name
    /// * `password` - Keyspace password
    /// * `key_id` - Keypair ID
    /// * `message` - Message to sign
    pub async fn sign(
        &self,
        keyspace: &str,
        password: &[u8],
        key_id: &str,
        message: &[u8],
    ) -> Result<Vec<u8>, VaultError> {
        let data = self.unlock_keyspace(keyspace, password).await?;
        let key = data
            .keypairs
            .iter()
            .find(|k| k.id == key_id)
            .ok_or(VaultError::KeyNotFound(key_id.to_string()))?;
        match key.key_type {
            KeyType::Ed25519 => {
                use ed25519_dalek::{Signer, SigningKey};
                let signing =
                    SigningKey::from_bytes(&key.private_key.clone().try_into().map_err(|_| {
                        VaultError::Crypto("Invalid Ed25519 private key length".to_string())
                    })?);
                let sig = signing.sign(message);
                Ok(sig.to_bytes().to_vec())
            }
            KeyType::Secp256k1 => {
                use k256::ecdsa::{signature::Signer, SigningKey};
                let arr: &[u8; 32] = key.private_key.as_slice().try_into().map_err(|_| {
                    VaultError::Crypto("Invalid secp256k1 private key length".to_string())
                })?;
                let sk = SigningKey::from_bytes(arr.into())
                    .map_err(|e| VaultError::Crypto(e.to_string()))?;
                let sig: k256::ecdsa::DerSignature = sk.sign(message);
                Ok(sig.to_vec())
            }
        }
    }

    /// Verify a signature with a stored keypair in a keyspace
    ///
    /// # Arguments
    /// * `keyspace` - Keyspace name
    /// * `password` - Keyspace password
    /// * `key_id` - Keypair ID
    /// * `message` - Message that was signed
    /// * `signature` - Signature to verify
    pub async fn verify(
        &self,
        keyspace: &str,
        password: &[u8],
        key_id: &str,
        message: &[u8],
        signature: &[u8],
    ) -> Result<bool, VaultError> {
        let data = self.unlock_keyspace(keyspace, password).await?;
        let key = data
            .keypairs
            .iter()
            .find(|k| k.id == key_id)
            .ok_or(VaultError::KeyNotFound(key_id.to_string()))?;
        match key.key_type {
            KeyType::Ed25519 => {
                use ed25519_dalek::{Signature, Verifier, VerifyingKey};
                let verifying =
                    VerifyingKey::from_bytes(&key.public_key.clone().try_into().map_err(|_| {
                        VaultError::Crypto("Invalid Ed25519 public key length".to_string())
                    })?)
                    .map_err(|e| VaultError::Crypto(e.to_string()))?;
                let sig = Signature::from_bytes(&signature.try_into().map_err(|_| {
                    VaultError::Crypto("Invalid Ed25519 signature length".to_string())
                })?);
                Ok(verifying.verify(message, &sig).is_ok())
            }
            KeyType::Secp256k1 => {
                use k256::ecdsa::{signature::Verifier, Signature, VerifyingKey};
                let pk = VerifyingKey::from_sec1_bytes(&key.public_key)
                    .map_err(|e| VaultError::Crypto(e.to_string()))?;
                let sig = Signature::from_der(signature)
                    .map_err(|e| VaultError::Crypto(e.to_string()))?;
                Ok(pk.verify(message, &sig).is_ok())
            }
        }
    }

    /// Encrypt a message using the keyspace symmetric cipher
    /// (for simplicity, uses keyspace password-derived key)
    pub async fn encrypt(
        &self,
        keyspace: &str,
        password: &[u8],
        plaintext: &[u8],
    ) -> Result<Vec<u8>, VaultError> {
        debug!("encrypt");

        // 1. Load keyspace metadata
        let meta_bytes = self
            .storage
            .get(keyspace)
            .await
            .map_err(|e| VaultError::Storage(format!("{e:?}")))?;
        let meta_bytes = match meta_bytes {
            Some(val) => val,
            None => {
                debug!("keyspace not found");
                return Err(VaultError::Other("Keyspace not found".to_string()));
            }
        };
        let meta: KeyspaceMetadata = match serde_json::from_slice(&meta_bytes) {
            Ok(val) => val,
            Err(e) => {
                debug!("serialization error: {}", e);
                return Err(VaultError::Serialization(e.to_string()));
            }
        };
        debug!(
            "salt={:?} (hex salt: {})",
            meta.salt,
            hex::encode(&meta.salt)
        );
        // 2. Derive key
        let key = kdf::keyspace_key(password, &meta.salt);
        // 3. Generate nonce
        let nonce = random_salt(12);
        debug!("nonce={:?} (hex nonce: {})", nonce, hex::encode(&nonce));
        // 4. Encrypt
        let ciphertext = encrypt_chacha20(&key, plaintext, &nonce).map_err(VaultError::Crypto)?;
        let mut out = nonce;
        out.extend_from_slice(&ciphertext);
        Ok(out)
    }

    /// Decrypt a message using the keyspace symmetric cipher
    /// (for simplicity, uses keyspace password-derived key)
    pub async fn decrypt(
        &self,
        keyspace: &str,
        password: &[u8],
        ciphertext: &[u8],
    ) -> Result<Vec<u8>, VaultError> {
        debug!("decrypt");

        // 1. Load keyspace metadata
        let meta_bytes = self
            .storage
            .get(keyspace)
            .await
            .map_err(|e| VaultError::Storage(format!("{e:?}")))?;
        let meta_bytes = match meta_bytes {
            Some(val) => val,
            None => {
                debug!("keyspace not found");
                return Err(VaultError::Other("Keyspace not found".to_string()));
            }
        };
        let meta: KeyspaceMetadata = match serde_json::from_slice(&meta_bytes) {
            Ok(val) => val,
            Err(e) => {
                debug!("serialization error: {}", e);
                return Err(VaultError::Serialization(e.to_string()));
            }
        };
        debug!(
            "salt={:?} (hex salt: {})",
            meta.salt,
            hex::encode(&meta.salt)
        );
        // 2. Derive key
        let key = kdf::keyspace_key(password, &meta.salt);
        // 3. Extract nonce
        let nonce = &ciphertext[..12];
        debug!("nonce={:?} (hex nonce: {})", nonce, hex::encode(nonce));
        // 4. Decrypt
        let plaintext =
            decrypt_chacha20(&key, &ciphertext[12..], nonce).map_err(VaultError::Crypto)?;
        Ok(plaintext)
    }
}