//! vault: Cryptographic keyspace and operations


//! vault: Cryptographic keyspace and operations

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

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

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

/// 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
fn encrypt_with_nonce_prepended(key: &[u8], plaintext: &[u8], cipher: &str) -> Result<Vec<u8>, VaultError> {
    use crate::crypto::random_salt;
    use crate::crypto;
    let nonce = random_salt(12);
    debug!("nonce: {}", hex::encode(&nonce));
    let (ct, _key_hex) = match cipher {
        "chacha20poly1305" => {
            let ct = encrypt_chacha20(key, plaintext, &nonce)
                .map_err(|e| VaultError::Crypto(e))?;
            debug!("ct: {}", hex::encode(&ct));
            debug!("key: {}", hex::encode(key));
            (ct, hex::encode(key))
        },
        "aes-gcm" => {
            let ct = encrypt_aes_gcm(key, plaintext, &nonce)
                .map_err(|e| VaultError::Crypto(e))?;
            debug!("ct: {}", hex::encode(&ct));
            debug!("key: {}", hex::encode(key));
            (ct, hex::encode(key))
        },
        _ => {
            debug!("unsupported cipher: {}", cipher);
            return Err(VaultError::Other(format!("Unsupported cipher: {cipher}")));
        }
    };
    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.
    pub async fn create_keyspace(&mut self, name: &str, password: &[u8], kdf: &str, cipher: &str, 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::{random_salt, kdf};
        use crate::data::{KeyspaceMetadata, KeyspaceData};
        use serde_json;
        
        // 1. Generate salt
        let salt = random_salt(16);
        debug!("salt: {:?}", salt);
        // 2. Derive key
        let key = match kdf {
            "scrypt" => match kdf::derive_key_scrypt(password, &salt, 32) {
                Ok(val) => val,
                Err(e) => {
                    debug!("kdf scrypt error: {}", e);
                    return Err(VaultError::Crypto(e));
                }
            },
            "pbkdf2" => kdf::derive_key_pbkdf2(password, &salt, 32, 10_000),
            _ => {
                debug!("unsupported KDF: {}", kdf);
                return Err(VaultError::Other(format!("Unsupported KDF: {kdf}")));
            }
        };
        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
        let encrypted_blob = encrypt_with_nonce_prepended(&key, &plaintext, cipher)?;
        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.try_into().unwrap_or([0u8; 16]),
            kdf: kdf.to_string(),
            cipher: cipher.to_string(),
            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");
        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
    pub async fn unlock_keyspace(&self, name: &str, password: &[u8]) -> Result<KeyspaceData, VaultError> {
        debug!("unlock_keyspace entry: name={}", name);
        use crate::crypto::{kdf};
        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 = match metadata.kdf.as_str() {
            "scrypt" => match kdf::derive_key_scrypt(password, &metadata.salt, 32) {
                Ok(val) => val,
                Err(e) => {
                    debug!("kdf scrypt error: {}", e);
                    return Err(VaultError::Crypto(e));
                }
            },
            "pbkdf2" => kdf::derive_key_pbkdf2(password, &metadata.salt, 32, 10_000),
            _ => {
                debug!("unsupported KDF: {}", metadata.kdf);
                return Err(VaultError::Other(format!("Unsupported KDF: {}", metadata.kdf)));
            }
        };
        debug!("derived key: {} bytes", key.len());
        debug!("derived key (hex): {}", hex::encode(&key));
        // 3. Split nonce and ciphertext
        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));
        // 4. Decrypt
        let plaintext = match metadata.cipher.as_str() {
            "chacha20poly1305" => match decrypt_chacha20(&key, ct, nonce) {
                Ok(val) => val,
                Err(e) => {
                    debug!("chacha20poly1305 error: {}", e);
                    return Err(VaultError::Crypto(e));
                }
            },
            "aes-gcm" => match decrypt_aes_gcm(&key, ct, nonce) {
                Ok(val) => val,
                Err(e) => {
                    debug!("aes-gcm error: {}", e);
                    return Err(VaultError::Crypto(e));
                }
            },
            _ => {
                debug!("unsupported cipher: {}", metadata.cipher);
                return Err(VaultError::Other(format!("Unsupported cipher: {}", metadata.cipher)));
            }
        };
        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 ---

    /// Add a new keypair to a keyspace (generates and stores a new keypair)
    pub async fn add_keypair(&mut self, keyspace: &str, password: &[u8], key_type: 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 (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;
        // 1. Fetch metadata
        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: kdf={} cipher={} salt={:?}", metadata.kdf, metadata.cipher, 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 = match metadata.kdf.as_str() {
            "scrypt" => match kdf::derive_key_scrypt(password, &metadata.salt, 32) {
                Ok(val) => val,
                Err(e) => {
                    debug!("kdf scrypt error: {}", e);
                    return Err(VaultError::Crypto(e));
                }
            },
            "pbkdf2" => kdf::derive_key_pbkdf2(password, &metadata.salt, 32, 10_000),
            _ => {
                debug!("unsupported KDF: {}", metadata.kdf);
                return Err(VaultError::Other(format!("Unsupported KDF: {}", metadata.kdf)));
            }
        };
        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, &metadata.cipher)?;
        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::{SigningKey, Signer};
                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::{SigningKey, signature::Signer};
                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::{VerifyingKey, Signature, Verifier};
                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::{VerifyingKey, Signature, signature::Verifier};
                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");
    debug!("keyspace={}", keyspace);
    use crate::crypto::{kdf};
    // 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={:?} cipher={} (hex salt: {})", meta.salt, meta.cipher, hex::encode(&meta.salt));
    // 2. Derive key
    let key = match meta.kdf.as_str() {
        "scrypt" => match kdf::derive_key_scrypt(password, &meta.salt, 32) {
            Ok(val) => val,
            Err(e) => {
                debug!("kdf scrypt error: {}", e);
                return Err(VaultError::Crypto(e));
            }
        },
        "pbkdf2" => kdf::derive_key_pbkdf2(password, &meta.salt, 32, 10_000),
        _ => {
            debug!("unsupported KDF: {}", meta.kdf);
            return Err(VaultError::Other(format!("Unsupported KDF: {}", meta.kdf)));
        }
    };
    // 3. Generate nonce
    let nonce = random_salt(12);
    debug!("nonce={:?} (hex nonce: {})", nonce, hex::encode(&nonce));
    // 4. Encrypt
    let ciphertext = match meta.cipher.as_str() {
        "chacha20poly1305" => match encrypt_chacha20(&key, plaintext, &nonce) {
            Ok(val) => val,
            Err(e) => {
                debug!("chacha20poly1305 error: {}", e);
                return Err(VaultError::Crypto(e));
            }
        },
        "aes-gcm" => match encrypt_aes_gcm(&key, plaintext, &nonce) {
            Ok(val) => val,
            Err(e) => {
                debug!("aes-gcm error: {}", e);
                return Err(VaultError::Crypto(e));
            }
        },
        _ => {
            debug!("unsupported cipher: {}", meta.cipher);
            return Err(VaultError::Other(format!("Unsupported cipher: {}", meta.cipher)));
        }
    };
    // 5. Prepend nonce to ciphertext
    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");
    debug!("keyspace={}", keyspace);
    use crate::crypto::{kdf};
    // 1. Fetch metadata
    let meta_bytes = self.storage.get(keyspace).await.map_err(|e| VaultError::Storage(format!("{e:?}")))?;
    let meta_bytes = meta_bytes.ok_or(VaultError::KeyspaceNotFound(keyspace.to_string()))?;
    let metadata: KeyspaceMetadata = serde_json::from_slice(&meta_bytes).map_err(|e| VaultError::Serialization(e.to_string()))?;
    debug!("salt={:?} cipher={} (hex salt: {})", metadata.salt, metadata.cipher, hex::encode(&metadata.salt));
    // 2. Derive key
    let key = match metadata.kdf.as_str() {
        "scrypt" => match kdf::derive_key_scrypt(password, &metadata.salt, 32) {
            Ok(val) => val,
            Err(e) => {
                debug!("storage error: {:?}", e);
                return Err(VaultError::Crypto(e));
            }
        },
        "pbkdf2" => kdf::derive_key_pbkdf2(password, &metadata.salt, 32, 10_000),
        _ => {
            debug!("unsupported KDF: {}", metadata.kdf);
            return Err(VaultError::Other(format!("Unsupported KDF: {}", metadata.kdf)));
        }
    };
    // 3. Split nonce and ciphertext
    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 nonce: {})", nonce, hex::encode(nonce));
    // 4. Decrypt
    let plaintext = match metadata.cipher.as_str() {
        "chacha20poly1305" => match decrypt_chacha20(&key, ct, nonce) {
            Ok(val) => val,
            Err(e) => {
                debug!("chacha20poly1305 error: {}", e);
                return Err(VaultError::Crypto(e));
            }
        },
        "aes-gcm" => match decrypt_aes_gcm(&key, ct, nonce) {
            Ok(val) => val,
            Err(e) => {
                debug!("aes-gcm error: {}", e);
                return Err(VaultError::Crypto(e));
            }
        },
        _ => {
            debug!("unsupported cipher: {}", metadata.cipher);
            return Err(VaultError::Other(format!("Unsupported cipher: {}", metadata.cipher)));
        }
    };
    Ok(plaintext)
}
} // <-- Close the impl block