/// AES-256-GCM encrypted identity vault with Argon2id key derivation.

use aes_gcm::{aead::{Aead, KeyInit}, Aes256Gcm, Nonce};
use argon2::{password_hash::SaltString, Argon2, Params, Version};
use rand::rngs::OsRng;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use zeroize::Zeroize;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Identity {
    pub name: String,
    pub protocol: String,
    pub credentials: String,
    pub created_at: chrono::DateTime<chrono::Utc>,
}

#[derive(Debug, Serialize, Deserialize)]
struct VaultBlob { salt: String, nonce: String, ciphertext: String, version: u32 }

#[derive(Debug)]
pub struct Vault {
    identities: Vec<Identity>,
    key: [u8; 32],
    /// Base64 salt used to derive `key`. Must be reused verbatim on every
    /// flush -- generating a fresh salt per-flush would desync it from the
    /// key already in memory, making the vault undecryptable even with the
    /// correct password (the bug this field exists to prevent).
    salt: String,
    vault_path: PathBuf,
}

impl Drop for Vault {
    fn drop(&mut self) {
        self.key.zeroize();
    }
}

fn vault_dir() -> PathBuf { dirs::home_dir().unwrap_or_else(|| PathBuf::from(".")).join(".nirc") }
fn vault_path() -> PathBuf { vault_dir().join("vault.json") }

impl Vault {
    pub fn create(password: &str) -> anyhow::Result<Self> { Self::create_at(&vault_path(), password) }
    pub fn unlock(password: &str) -> anyhow::Result<Self> { Self::unlock_at(&vault_path(), password) }

    /// Create a new vault at an explicit path. `create()` is a thin wrapper
    /// over this using the default `~/.nirc/vault.json` location; tests use
    /// this directly with an isolated temp path so parallel test runs don't
    /// race on the same on-disk file.
    pub fn create_at(path: &std::path::Path, password: &str) -> anyhow::Result<Self> {
        if path.exists() { std::fs::remove_file(path)?; }
        if let Some(parent) = path.parent() { std::fs::create_dir_all(parent)?; }
        let salt = SaltString::generate(&mut OsRng);
        let key = derive_key(password, &salt)?;
        let vault = Self { identities: Vec::new(), key, salt: salt.to_string(), vault_path: path.to_path_buf() };
        vault.flush()?;
        Ok(vault)
    }
    /// Unlock a vault at an explicit path. See `create_at`.
    pub fn unlock_at(path: &std::path::Path, password: &str) -> anyhow::Result<Self> {
        let raw = std::fs::read_to_string(path)?;
        let blob: VaultBlob = serde_json::from_str(&raw)?;
        let salt = SaltString::from_b64(&blob.salt).map_err(|e| anyhow::anyhow!("invalid salt: {e}"))?;
        let key = derive_key(password, &salt)?;
        let cipher = Aes256Gcm::new_from_slice(&key)
            .map_err(|e| anyhow::anyhow!("cipher init: {e}"))?;
        let nonce_bytes = base64_url_decode(&blob.nonce)?;
        let nonce = Nonce::from_slice(&nonce_bytes);
        let ct_bytes = base64_url_decode(&blob.ciphertext)?;
        let pt = cipher.decrypt(nonce, ct_bytes.as_ref())
            .map_err(|_| anyhow::anyhow!("wrong password or corrupted vault"))?;
        let plaintext = String::from_utf8(pt)?;
        let identities: Vec<Identity> = if plaintext.is_empty() { Vec::new() } else { serde_json::from_str(&plaintext)? };
        Ok(Self { identities, key, salt: blob.salt, vault_path: path.to_path_buf() })
    }
    fn flush(&self) -> anyhow::Result<()> {
        let plaintext = serde_json::to_string(&self.identities)?;
        let cipher = Aes256Gcm::new_from_slice(&self.key)
            .map_err(|e| anyhow::anyhow!("cipher init: {e}"))?;
        let nonce_bytes = rand::random::<[u8; 12]>();
        let nonce = Nonce::from_slice(&nonce_bytes);
        let ct = cipher.encrypt(nonce, plaintext.as_bytes())
            .map_err(|e| anyhow::anyhow!("encrypt: {e}"))?;
        let blob = VaultBlob { salt: self.salt.clone(), nonce: base64_url_encode(&nonce_bytes), ciphertext: base64_url_encode(&ct), version: 1 };
        let json = serde_json::to_string_pretty(&blob)?;
        let tmp = self.vault_path.with_extension("json.tmp");
        std::fs::write(&tmp, &json)?;
        std::fs::rename(&tmp, &self.vault_path)?;
        Ok(())
    }
    pub fn add_identity(&mut self, id: Identity) -> anyhow::Result<()> { self.identities.push(id); self.flush() }
    pub fn list_id(&self) -> &[Identity] { &self.identities }
    pub fn remove_identity(&mut self, name: &str) -> anyhow::Result<bool> {
        let before = self.identities.len();
        self.identities.retain(|i| i.name != name);
        if self.identities.len() < before { self.flush()?; Ok(true) } else { Ok(false) }
    }
    pub fn lock(self) { drop(self); }
}

fn derive_key(password: &str, salt: &SaltString) -> anyhow::Result<[u8; 32]> {
    let params = Params::new(65536, 3, 2, Some(32))
        .map_err(|e| anyhow::anyhow!("argon2 params: {e}"))?;
    let argon2 = Argon2::new(argon2::Algorithm::Argon2id, Version::V0x13, params);
    let mut key = [0u8; 32];
    argon2.hash_password_into(password.as_bytes(), salt.as_ref().as_bytes(), &mut key)
        .map_err(|e| anyhow::anyhow!("argon2 hash: {e}"))?;
    let out = key;
    key.zeroize(); // Wipe the stack copy before returning.
    Ok(out)
}
fn base64_url_encode(data: &[u8]) -> String { use base64::Engine; base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data) }
fn base64_url_decode(s: &str) -> anyhow::Result<Vec<u8>> { use base64::Engine; Ok(base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(s)?) }

#[cfg(test)]
mod tests {
    use super::*;
    #[test] fn create_and_unlock_roundtrip() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("vault.json");
        let mut vault = Vault::create_at(&path, "testpass").unwrap();
        vault.add_identity(Identity { name: "libera".into(), protocol: "irc".into(), credentials: "nick=testuser".into(), created_at: chrono::Utc::now() }).unwrap();
        drop(vault);
        let vault2 = Vault::unlock_at(&path, "testpass").unwrap();
        assert_eq!(vault2.list_id().len(), 1);
        assert_eq!(vault2.list_id()[0].name, "libera");
    }
    #[test] fn wrong_password_fails() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("vault.json");
        Vault::create_at(&path, "correct").unwrap();
        assert!(Vault::unlock_at(&path, "wrong").is_err());
    }
    #[test] fn remove_identity() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("vault.json");
        let mut vault = Vault::create_at(&path, "pass").unwrap();
        vault.add_identity(Identity { name: "a".into(), protocol: "irc".into(), credentials: "x".into(), created_at: chrono::Utc::now() }).unwrap();
        vault.add_identity(Identity { name: "b".into(), protocol: "matrix".into(), credentials: "y".into(), created_at: chrono::Utc::now() }).unwrap();
        assert_eq!(vault.list_id().len(), 2);
        assert!(vault.remove_identity("a").unwrap());
        assert_eq!(vault.list_id().len(), 1);
        assert!(!vault.remove_identity("nonexistent").unwrap());
    }
}