use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;

use aes::cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use rand::RngCore;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::error::{AppError, AppResult};

const SERVICE_NAME: &str = "com.eritors.scribe.desktop";
const IV_LENGTH: usize = 16;

type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;

// ===== Secure vault: encrypted JSON file (like Electron's safeStorage) =====

#[derive(Serialize, Deserialize, Default)]
struct SecureVault {
    token: Option<String>,
    encryption_keys: HashMap<String, String>,
    last_user_id: Option<String>,
    pin_hashes: HashMap<String, String>,
    #[serde(default)]
    pin_failed_attempts: i32,
    #[serde(default)]
    pin_locked_until: Option<i64>,
}

fn vault_path() -> PathBuf {
    dirs_next::data_dir()
        .unwrap_or_else(|| PathBuf::from("."))
        .join(SERVICE_NAME)
        .join("secure-config.json")
}

const KEYRING_USER: &str = "vault-key";

/// Retrieves (or generates and stores) the vault encryption key via the OS keyring
/// (macOS Keychain, Windows DPAPI, Linux Secret Service).
/// Falls back to the old derivation method if the keyring is unavailable,
/// and attempts to migrate the key into the keyring for next time.
fn get_vault_key() -> [u8; 32] {
    if cfg!(debug_assertions) {
        return derive_machine_key_legacy();
    }
    let entry = keyring::Entry::new(SERVICE_NAME, KEYRING_USER);
    if let Ok(entry) = &entry {
        if let Ok(stored) = entry.get_password() {
            if let Ok(decoded) = BASE64.decode(stored.trim()) {
                if decoded.len() == 32 {
                    let mut key = [0u8; 32];
                    key.copy_from_slice(&decoded);
                    return key;
                }
            }
        }
    }
    let mut key = [0u8; 32];
    rand::rng().fill_bytes(&mut key);
    let encoded = BASE64.encode(key);
    if let Ok(entry) = &entry {
        let _ = entry.set_password(&encoded);
    }
    key
}

/// Legacy derivation for migrating vaults created before keyring support.
fn derive_machine_key_legacy() -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(SERVICE_NAME.as_bytes());
    if let Ok(name) = hostname::get() {
        hasher.update(name.as_encoded_bytes());
    }
    if let Some(dir) = dirs_next::home_dir() {
        hasher.update(dir.to_string_lossy().as_bytes());
    }
    hasher.finalize().into()
}

fn encrypt_vault(data: &[u8], key: &[u8; 32]) -> AppResult<Vec<u8>> {
    let mut iv = [0u8; IV_LENGTH];
    rand::rng().fill_bytes(&mut iv);
    let block_size = 16;
    let padding_len = block_size - (data.len() % block_size);
    let mut buf = Vec::with_capacity(data.len() + padding_len);
    buf.extend_from_slice(data);
    buf.extend(std::iter::repeat(padding_len as u8).take(padding_len));
    let padded_len = buf.len();
    let cipher = Aes256CbcEnc::new(key.into(), &iv.into());
    cipher.encrypt_padded_mut::<aes::cipher::block_padding::NoPadding>(&mut buf, padded_len)
        .map_err(|e| AppError::Encryption(format!("Vault encrypt failed: {}", e)))?;
    let mut result = Vec::with_capacity(IV_LENGTH + buf.len());
    result.extend_from_slice(&iv);
    result.extend_from_slice(&buf);
    Ok(result)
}

fn decrypt_vault(data: &[u8], key: &[u8; 32]) -> AppResult<Vec<u8>> {
    if data.len() < IV_LENGTH + 1 {
        return Err(AppError::Encryption("Vault data too short".into()));
    }
    let iv: [u8; 16] = data[..IV_LENGTH].try_into()
        .map_err(|_| AppError::Encryption("Invalid IV".into()))?;
    let mut buf = data[IV_LENGTH..].to_vec();
    let cipher = Aes256CbcDec::new(key.into(), &iv.into());
    let decrypted = cipher.decrypt_padded_mut::<aes::cipher::block_padding::Pkcs7>(&mut buf)
        .map_err(|e| AppError::Encryption(format!("Vault decrypt failed: {}", e)))?;
    Ok(decrypted.to_vec())
}

fn read_vault() -> AppResult<SecureVault> {
    let path = vault_path();
    if !path.exists() {
        return Ok(SecureVault::default());
    }
    let content = fs::read_to_string(&path)
        .map_err(|e| AppError::Keyring(format!("Failed to read vault: {}", e)))?;
    let raw = BASE64.decode(content.trim())
        .map_err(|e| AppError::Keyring(format!("Vault corrupted (base64): {}", e)))?;

    let key = get_vault_key();
    if let Ok(decrypted) = decrypt_vault(&raw, &key) {
        if let Ok(vault) = serde_json::from_slice::<SecureVault>(&decrypted) {
            return Ok(vault);
        }
    }

    let legacy_key = derive_machine_key_legacy();
    if let Ok(decrypted) = decrypt_vault(&raw, &legacy_key) {
        if let Ok(vault) = serde_json::from_slice::<SecureVault>(&decrypted) {
            let _ = write_vault_with_key(&vault, &key);
            return Ok(vault);
        }
    }

    Ok(SecureVault::default())
}

fn write_vault_with_key(vault: &SecureVault, key: &[u8; 32]) -> AppResult<()> {
    let path = vault_path();
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent).map_err(|e| AppError::Internal(format!("Failed to create vault dir: {}", e)))?;
    }
    let json = serde_json::to_string(vault)
        .map_err(|e| AppError::Internal(format!("Failed to serialize vault: {}", e)))?;
    let encrypted = encrypt_vault(json.as_bytes(), key)?;
    let encoded = BASE64.encode(&encrypted);
    fs::write(&path, encoded).map_err(|e| AppError::Internal(format!("Failed to write vault: {}", e)))
}

fn write_vault(vault: &SecureVault) -> AppResult<()> {
    let key = get_vault_key();
    write_vault_with_key(vault, &key)
}

// ===== Public API (same signatures as before) =====

pub fn get_user_encryption_key(user_id: &str) -> AppResult<String> {
    let vault = read_vault()?;
    vault.encryption_keys.get(user_id).cloned()
        .ok_or_else(|| AppError::Keyring(format!("No encryption key for user {}", user_id)))
}

pub fn set_user_encryption_key(user_id: &str, encryption_key: &str) -> AppResult<()> {
    let mut vault = read_vault()?;
    vault.encryption_keys.insert(user_id.to_string(), encryption_key.to_string());
    write_vault(&vault)
}

pub fn has_user_encryption_key(user_id: &str) -> AppResult<bool> {
    let vault = read_vault()?;
    Ok(vault.encryption_keys.contains_key(user_id))
}


pub fn get_token() -> AppResult<Option<String>> {
    let vault = read_vault()?;
    Ok(vault.token)
}

pub fn set_token(token: &str) -> AppResult<()> {
    let mut vault = read_vault()?;
    vault.token = Some(token.to_string());
    write_vault(&vault)
}

pub fn remove_token() -> AppResult<()> {
    let mut vault = read_vault()?;
    vault.token = None;
    write_vault(&vault)
}

pub fn set_pin_hash(user_id: &str, pin_hash: &str) -> AppResult<()> {
    let mut vault = read_vault()?;
    vault.pin_hashes.insert(user_id.to_string(), pin_hash.to_string());
    write_vault(&vault)
}

pub fn get_pin_hash(user_id: &str) -> AppResult<Option<String>> {
    let vault = read_vault()?;
    Ok(vault.pin_hashes.get(user_id).cloned())
}

pub fn set_last_user_id(user_id: &str) -> AppResult<()> {
    let mut vault = read_vault()?;
    vault.last_user_id = Some(user_id.to_string());
    write_vault(&vault)
}

pub fn get_last_user_id() -> AppResult<Option<String>> {
    let vault = read_vault()?;
    Ok(vault.last_user_id)
}

pub fn clear_vault() -> AppResult<()> {
    write_vault(&SecureVault::default())
}

const MAX_PIN_ATTEMPTS: i32 = 5;
const PIN_LOCKOUT_SECONDS: i64 = 300; // 5 minutes

pub fn check_pin_rate_limit() -> AppResult<()> {
    let vault = read_vault()?;
    if let Some(locked_until) = vault.pin_locked_until {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_secs() as i64)
            .unwrap_or(0);
        if now < locked_until {
            let remaining = locked_until - now;
            return Err(AppError::Auth(format!("Too many attempts. Try again in {} seconds.", remaining)));
        }
    }
    Ok(())
}

pub fn record_pin_failure() -> AppResult<()> {
    let mut vault = read_vault()?;
    vault.pin_failed_attempts += 1;
    if vault.pin_failed_attempts >= MAX_PIN_ATTEMPTS {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_secs() as i64)
            .unwrap_or(0);
        vault.pin_locked_until = Some(now + PIN_LOCKOUT_SECONDS);
    }
    write_vault(&vault)
}

pub fn reset_pin_attempts() -> AppResult<()> {
    let mut vault = read_vault()?;
    vault.pin_failed_attempts = 0;
    vault.pin_locked_until = None;
    write_vault(&vault)
}
