//! Encrypted local vault — byte-compatible port of `security/vaultStore.js`.
//!
//! Format (unchanged, so the 0.1.0 master password keeps working):
//! * `vault/core-v1.meta.json`: scrypt parameters (N=32768, r=8, p=1, 32 B key) and a 16 B salt.
//! * `vault/core-v1.data.json`: AES-256-GCM, 12 B IV, 16 B tag, AAD `CHJ_CORE_VAULT_V1`,
//!   plaintext = compact JSON `{schemaVersion:1, profiles, hostKeys, secrets}`.
//! The key exists only in Core memory and is zeroized on lock.

use crate::core::error::{CoreError, CoreResult};
use crate::core::storage::atomic_write_json;
use aes_gcm::aead::{Aead, KeyInit, Payload};
use aes_gcm::{Aes256Gcm, Nonce};
use base64::engine::general_purpose::STANDARD as B64;
use base64::Engine;
use parking_lot::Mutex;
use rand::RngCore;
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::fs;
use std::path::{Path, PathBuf};
use zeroize::{Zeroize, Zeroizing};

pub const KEY_LENGTH: usize = 32;
const SALT_LENGTH: usize = 16;
const IV_LENGTH: usize = 12;
const TAG_LENGTH: usize = 16;
const AAD: &[u8] = b"CHJ_CORE_VAULT_V1";
const SCRYPT_N: u64 = 1 << 15;
const SCRYPT_R: u32 = 8;
const SCRYPT_P: u32 = 1;
pub const MIN_PASSWORD_LENGTH: usize = 4;
pub const MAX_PASSWORD_LENGTH: usize = 64;

pub type VaultKey = Zeroizing<[u8; KEY_LENGTH]>;

pub fn initial_vault() -> Value {
    json!({ "schemaVersion": 1, "profiles": [], "hostKeys": {}, "secrets": {} })
}

pub fn validate_password(password: &str) -> CoreResult<String> {
    if password.contains('\0') {
        return Err(CoreError::plain("Password contains an invalid character."));
    }
    let length = password.chars().count();
    if !(MIN_PASSWORD_LENGTH..=MAX_PASSWORD_LENGTH).contains(&length) {
        return Err(CoreError::plain(format!(
            "Vault password must contain {MIN_PASSWORD_LENGTH} to {MAX_PASSWORD_LENGTH} characters."
        )));
    }
    Ok(password.to_string())
}

pub fn validate_vault_data(mut value: Value) -> CoreResult<Value> {
    let object = value.as_object_mut().ok_or_else(|| CoreError::plain("Vault data is invalid."))?;
    if object.get("schemaVersion").and_then(Value::as_f64) != Some(1.0) {
        return Err(CoreError::plain("Unsupported vault data schema."));
    }
    if !object.get("profiles").map(Value::is_array).unwrap_or(false) {
        return Err(CoreError::plain("Vault profiles are invalid."));
    }
    if !object.get("hostKeys").map(Value::is_object).unwrap_or(false) {
        return Err(CoreError::plain("Vault host keys are invalid."));
    }
    if !object.get("secrets").map(Value::is_object).unwrap_or(false) {
        object.insert("secrets".into(), json!({}));
    }
    Ok(value)
}

fn derive_key(password: &str, salt: &[u8], n: u64, r: u32, p: u32) -> CoreResult<VaultKey> {
    if n < 2 || !n.is_power_of_two() || n > (1 << 20) || r == 0 || r > 32 || p == 0 || p > 16 {
        return Err(CoreError::plain("Unsupported vault KDF parameters."));
    }
    let log_n = n.trailing_zeros() as u8;
    let params = scrypt::Params::new(log_n, r, p, KEY_LENGTH).map_err(|_| CoreError::plain("Unsupported vault KDF parameters."))?;
    let mut key: VaultKey = Zeroizing::new([0u8; KEY_LENGTH]);
    scrypt::scrypt(password.as_bytes(), salt, &params, key.as_mut()).map_err(|_| CoreError::plain("Key derivation failed."))?;
    Ok(key)
}

#[derive(Default)]
struct Session {
    key: Option<VaultKey>,
    data: Option<Value>,
}

pub struct VaultStore {
    root: PathBuf,
    meta_path: PathBuf,
    data_path: PathBuf,
    session: Mutex<Session>,
    generation: Mutex<u64>,
    /// Serializes read-modify-write updates: two concurrent SSH connects must not lose a trusted host key.
    update_lock: Mutex<()>,
}

#[derive(serde::Serialize, Clone, Debug)]
pub struct VaultStatus {
    pub initialized: bool,
    #[serde(rename = "needsSetup")]
    pub needs_setup: bool,
    pub damaged: bool,
    pub unlocked: bool,
}

fn locked_error() -> CoreError {
    CoreError::new("VAULT_LOCKED", "Vault is locked.")
}

impl VaultStore {
    pub fn new(storage_root: &Path) -> Self {
        let root = storage_root.join("vault");
        Self {
            update_lock: Mutex::new(()),
            meta_path: root.join("core-v1.meta.json"),
            data_path: root.join("core-v1.data.json"),
            root,
            session: Mutex::new(Session::default()),
            generation: Mutex::new(0),
        }
    }

    pub fn generation(&self) -> u64 {
        *self.generation.lock()
    }

    pub fn status(&self) -> VaultStatus {
        let meta = self.meta_path.exists();
        let data = self.data_path.exists();
        let session = self.session.lock();
        VaultStatus {
            initialized: meta && data,
            needs_setup: !meta && !data,
            damaged: meta != data,
            unlocked: session.key.is_some() && session.data.is_some(),
        }
    }

    pub fn is_unlocked(&self) -> bool {
        let session = self.session.lock();
        session.key.is_some() && session.data.is_some()
    }

    fn read_meta(&self) -> CoreResult<Value> {
        Ok(serde_json::from_str(&fs::read_to_string(&self.meta_path)?)?)
    }

    /// Stable identity of this vault instance (same formula as 0.1.0, used by biometric enrollment).
    pub fn instance_id(&self) -> Option<String> {
        if !self.status().initialized {
            return None;
        }
        let meta = self.read_meta().ok()?;
        let identity = json!([meta.get("formatVersion"), meta.get("createdAt"), meta.get("salt")]);
        let text = serde_json::to_string(&identity).ok()?;
        Some(hex::encode(Sha256::digest(text.as_bytes())))
    }

    pub async fn create(&self, password: &str) -> CoreResult<VaultStatus> {
        let generation = self.generation();
        let status = self.status();
        if status.initialized {
            return Err(CoreError::plain("Vault is already initialized."));
        }
        if status.damaged {
            return Err(CoreError::plain("Vault files are incomplete. Restore or remove them before setup."));
        }
        let normalized = validate_password(password)?;
        let mut salt = [0u8; SALT_LENGTH];
        rand::rngs::OsRng.fill_bytes(&mut salt);
        let key = tokio::task::spawn_blocking(move || derive_key(&normalized, &salt, SCRYPT_N, SCRYPT_R, SCRYPT_P))
            .await
            .map_err(|_| CoreError::plain("Key derivation failed."))??;
        if generation != self.generation() || !self.status().needs_setup {
            return Err(CoreError::plain("Vault operation cancelled."));
        }
        let now = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
        let data = initial_vault();
        let result = (|| -> CoreResult<()> {
            atomic_write_json(
                &self.meta_path,
                &json!({
                    "formatVersion": 1,
                    "dataSchemaVersion": 1,
                    "createdAt": now,
                    "updatedAt": now,
                    "cipher": "aes-256-gcm",
                    "kdf": { "name": "scrypt", "N": SCRYPT_N, "r": SCRYPT_R, "p": SCRYPT_P, "keyLength": KEY_LENGTH },
                    "salt": B64.encode(salt)
                }),
            )?;
            self.write_encrypted(&data, &key)
        })();
        if let Err(error) = result {
            let _ = fs::remove_file(&self.meta_path);
            let _ = fs::remove_file(&self.data_path);
            return Err(error);
        }
        self.set_session(key, data);
        Ok(self.status())
    }

    fn kdf_from_meta(meta: &Value) -> CoreResult<(Vec<u8>, u64, u32, u32)> {
        if meta.get("formatVersion").and_then(Value::as_f64) != Some(1.0)
            || meta.get("cipher").and_then(Value::as_str) != Some("aes-256-gcm")
            || meta.pointer("/kdf/name").and_then(Value::as_str) != Some("scrypt")
        {
            return Err(CoreError::plain("Unsupported vault format."));
        }
        let salt = B64.decode(meta.get("salt").and_then(Value::as_str).unwrap_or("")).map_err(|_| CoreError::plain("Vault salt is invalid."))?;
        if salt.len() != SALT_LENGTH {
            return Err(CoreError::plain("Vault salt is invalid."));
        }
        let num = |key: &str, fallback: u64| meta.pointer(&format!("/kdf/{key}")).and_then(Value::as_u64).filter(|v| *v > 0).unwrap_or(fallback);
        Ok((salt, num("N", SCRYPT_N), num("r", SCRYPT_R as u64) as u32, num("p", SCRYPT_P as u64) as u32))
    }

    pub async fn unlock(&self, password: &str) -> CoreResult<VaultStatus> {
        let generation = self.generation();
        let instance = self.instance_id();
        if !self.status().initialized {
            return Err(CoreError::plain("Vault setup is required or its files are incomplete."));
        }
        let normalized = validate_password(password)?;
        let failure = || CoreError::new("VAULT_UNLOCK_FAILED", "Vault password is incorrect or the vault is damaged.");
        let meta = self.read_meta().map_err(|_| failure())?;
        let (salt, n, r, p) = Self::kdf_from_meta(&meta).map_err(|_| failure())?;
        let key = tokio::task::spawn_blocking(move || derive_key(&normalized, &salt, n, r, p))
            .await
            .map_err(|_| failure())?
            .map_err(|_| failure())?;
        if generation != self.generation() || instance != self.instance_id() {
            return Err(failure());
        }
        let data = self.read_encrypted(&key).map_err(|_| failure())?;
        self.set_session(key, data);
        Ok(self.status())
    }

    /// Copy of the unlock key for a Core-internal operation (biometric enrollment). Never exposed over IPC.
    pub fn unlock_key_copy(&self) -> CoreResult<VaultKey> {
        let session = self.session.lock();
        match (&session.key, &session.data) {
            (Some(key), Some(_)) => Ok(Zeroizing::new(**key)),
            _ => Err(locked_error()),
        }
    }

    pub async fn verify_password(&self, password: &str) -> CoreResult<bool> {
        self.assert_unlocked()?;
        let generation = self.generation();
        let required = || CoreError::new("BIOMETRIC_MASTER_PASSWORD_REQUIRED", "Master password required.");
        let normalized = validate_password(password)?;
        let meta = self.read_meta()?;
        let (salt, n, r, p) = Self::kdf_from_meta(&meta)?;
        let key = tokio::task::spawn_blocking(move || derive_key(&normalized, &salt, n, r, p))
            .await
            .map_err(|_| required())??;
        self.assert_unlocked()?;
        if generation != self.generation() {
            return Err(required());
        }
        let current = self.unlock_key_copy()?;
        if !bool::from(subtle::ConstantTimeEq::ct_eq(key.as_ref(), current.as_ref())) {
            return Err(required());
        }
        self.read_encrypted(&key)?;
        Ok(true)
    }

    pub fn unlock_with_key(&self, key: &[u8]) -> CoreResult<VaultStatus> {
        if !self.status().initialized || key.len() != KEY_LENGTH {
            return Err(CoreError::new("BIOMETRIC_INVALID_KEY", "Invalid protected unlocking key."));
        }
        let mut copy: VaultKey = Zeroizing::new([0u8; KEY_LENGTH]);
        copy.copy_from_slice(key);
        match self.read_encrypted(&copy) {
            Ok(data) => {
                self.set_session(copy, data);
                Ok(self.status())
            }
            Err(_) => Err(CoreError::new(
                "BIOMETRIC_ENROLLMENT_INVALIDATED",
                "Biometric enrollment is no longer valid. Use the master password.",
            )),
        }
    }

    pub fn lock(&self) -> VaultStatus {
        *self.generation.lock() += 1;
        {
            let mut session = self.session.lock();
            session.key = None; // Zeroizing wipes the key on drop.
            session.data = None;
        }
        self.status()
    }

    pub fn reset(&self, confirmation: &str) -> CoreResult<VaultStatus> {
        if confirmation != "SMAZAT" {
            return Err(CoreError::new("VAULT_RESET_CONFIRMATION_REQUIRED", "Vault reset confirmation is invalid."));
        }
        self.lock();
        if self.root.exists() {
            fs::remove_dir_all(&self.root)?;
        }
        let status = self.status();
        if !status.needs_setup {
            return Err(CoreError::new("VAULT_RESET_FAILED", "Vault data could not be removed completely."));
        }
        Ok(status)
    }

    pub fn assert_unlocked(&self) -> CoreResult<()> {
        if self.is_unlocked() {
            Ok(())
        } else {
            Err(locked_error())
        }
    }

    pub fn get_data(&self) -> CoreResult<Value> {
        let session = self.session.lock();
        match (&session.key, &session.data) {
            (Some(_), Some(data)) => Ok(data.clone()),
            _ => Err(locked_error()),
        }
    }

    pub fn replace_data(&self, next: Value) -> CoreResult<Value> {
        let normalized = validate_vault_data(next)?;
        let mut session = self.session.lock();
        let key = match (&session.key, &session.data) {
            (Some(key), Some(_)) => Zeroizing::new(**key),
            _ => return Err(locked_error()),
        };
        self.write_encrypted(&normalized, &key)?;
        session.data = Some(normalized.clone());
        Ok(normalized)
    }

    /// Mutator → encrypt → commit, like `VaultStore.update` in 0.1.0.
    pub fn update<R>(&self, mutator: impl FnOnce(&mut Value) -> CoreResult<R>) -> CoreResult<R> {
        let _update = self.update_lock.lock();
        let mut draft = self.get_data()?;
        let result = mutator(&mut draft)?;
        self.replace_data(draft)?;
        Ok(result)
    }

    fn write_encrypted(&self, data: &Value, key: &[u8; KEY_LENGTH]) -> CoreResult<()> {
        let data = validate_vault_data(data.clone())?;
        let mut plaintext = Zeroizing::new(serde_json::to_vec(&data)?);
        let mut iv = [0u8; IV_LENGTH];
        rand::rngs::OsRng.fill_bytes(&mut iv);
        let cipher = Aes256Gcm::new_from_slice(key).map_err(|_| CoreError::plain("Invalid vault key."))?;
        let mut sealed = cipher
            .encrypt(Nonce::from_slice(&iv), Payload { msg: &plaintext, aad: AAD })
            .map_err(|_| CoreError::plain("Vault encryption failed."))?;
        plaintext.zeroize();
        let tag = sealed.split_off(sealed.len() - TAG_LENGTH);
        let result = atomic_write_json(
            &self.data_path,
            &json!({
                "formatVersion": 1,
                "updatedAt": chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true),
                "iv": B64.encode(iv),
                "authTag": B64.encode(&tag),
                "ciphertext": B64.encode(&sealed)
            }),
        );
        sealed.zeroize();
        result
    }

    fn read_encrypted(&self, key: &[u8; KEY_LENGTH]) -> CoreResult<Value> {
        let blob: Value = serde_json::from_str(&fs::read_to_string(&self.data_path)?)?;
        if blob.get("formatVersion").and_then(Value::as_f64) != Some(1.0) {
            return Err(CoreError::plain("Unsupported encrypted vault data."));
        }
        let field = |name: &str| B64.decode(blob.get(name).and_then(Value::as_str).unwrap_or("")).unwrap_or_default();
        let iv = field("iv");
        let tag = field("authTag");
        let mut ciphertext = field("ciphertext");
        if iv.len() != IV_LENGTH || tag.len() != TAG_LENGTH || ciphertext.is_empty() {
            return Err(CoreError::plain("Encrypted vault data is invalid."));
        }
        ciphertext.extend_from_slice(&tag);
        let cipher = Aes256Gcm::new_from_slice(key).map_err(|_| CoreError::plain("Invalid vault key."))?;
        let plaintext = Zeroizing::new(
            cipher
                .decrypt(Nonce::from_slice(&iv), Payload { msg: &ciphertext, aad: AAD })
                .map_err(|_| CoreError::plain("Vault authentication failed."))?,
        );
        ciphertext.zeroize();
        validate_vault_data(serde_json::from_slice(&plaintext)?)
    }

    fn set_session(&self, key: VaultKey, data: Value) {
        self.lock();
        let mut session = self.session.lock();
        session.key = Some(key);
        session.data = Some(data);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn create_lock_unlock_roundtrip() {
        let dir = tempfile::tempdir().unwrap();
        let vault = VaultStore::new(dir.path());
        assert!(vault.status().needs_setup);
        vault.create("heslo123").await.unwrap();
        vault.update(|data| {
            data["profiles"].as_array_mut().unwrap().push(json!({"id": "p1"}));
            Ok(())
        })
        .unwrap();
        vault.lock();
        assert!(vault.get_data().is_err());
        assert_eq!(vault.unlock("wrong-pass").await.unwrap_err().code, "VAULT_UNLOCK_FAILED");
        vault.unlock("heslo123").await.unwrap();
        assert_eq!(vault.get_data().unwrap()["profiles"][0]["id"], "p1");
        assert!(vault.verify_password("heslo123").await.unwrap());
        assert!(vault.verify_password("jine-heslo").await.is_err());
    }

    #[tokio::test]
    async fn password_rules_count_code_points() {
        assert!(validate_password("ččč").is_err());
        assert!(validate_password("čččč").is_ok());
        assert!(validate_password(&"a".repeat(65)).is_err());
        assert!(validate_password("ab\0c").is_err());
    }

    #[tokio::test]
    async fn reset_requires_confirmation() {
        let dir = tempfile::tempdir().unwrap();
        let vault = VaultStore::new(dir.path());
        vault.create("heslo123").await.unwrap();
        assert!(vault.reset("DELETE").is_err());
        assert!(vault.reset("SMAZAT").unwrap().needs_setup);
    }
}