# backend/services/usb_vault.py
import os
import sqlite3
import logging
import base64
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from backend.services.usb_monitor import get_db_path
from enum import Enum

class MissingDomainKeyError(Exception):
    pass

class KeyDomain(str, Enum):
    CHAT = "chat"
    CLOUD = "cloud"
    GAMING_COACH = "gaming_coach"
    AUTO_UPGRADE_PC = "auto_upgrade_pc"
    AUTO_UPGRADE_CLOUD = "auto_upgrade_cloud"
    RESEARCH_ENGINE = "research_engine"
    CAPTCHA_SOLVER_PC = "captcha_solver_pc"
    CAPTCHA_SOLVER_CLOUD = "captcha_solver_cloud"
    IMAGE_GENERATION_PC = "image_generation_pc"
    IMAGE_GENERATION_CLOUD = "image_generation_cloud"
    OSINT_PROTOCOL_PC = "osint_protocol_pc"
    OSINT_PROTOCOL_CLOUD = "osint_protocol_cloud"
    SUPERVISOR_HEAL_PC = "supervisor_heal_pc"
    SUPERVISOR_HEAL_CLOUD = "supervisor_heal_cloud"
    HEAD_SUPERVISOR_PC = "head_supervisor_pc"
    HEAD_SUPERVISOR_CLOUD = "head_supervisor_cloud"
    CLOUD_SYNC_DOMAIN = "cloud_sync_domain"

KEY_DOMAIN_ENV_MAP = {
    KeyDomain.CHAT: "GOOGLE_API_KEY_CHAT",
    KeyDomain.CLOUD: "GOOGLE_API_KEY_CLOUD",
    KeyDomain.GAMING_COACH: "GOOGLE_API_KEY_GAMING_COACH",
    KeyDomain.AUTO_UPGRADE_PC: "GOOGLE_API_KEY_AUTO_UPGRADE_PC",
    KeyDomain.AUTO_UPGRADE_CLOUD: "GOOGLE_API_KEY_AUTO_UPGRADE_CLOUD",
    KeyDomain.RESEARCH_ENGINE: "GOOGLE_API_KEY_RESEARCH_ENGINE",
    KeyDomain.CAPTCHA_SOLVER_PC: "GOOGLE_API_KEY_CAPTCHA_SOLVER_PC",
    KeyDomain.CAPTCHA_SOLVER_CLOUD: "GOOGLE_API_KEY_CAPTCHA_SOLVER_CLOUD",
    KeyDomain.IMAGE_GENERATION_PC: "GOOGLE_API_KEY_IMAGE_GENERATION_PC",
    KeyDomain.IMAGE_GENERATION_CLOUD: "GOOGLE_API_KEY_IMAGE_GENERATION_CLOUD",
    KeyDomain.OSINT_PROTOCOL_PC: "GOOGLE_API_KEY_OSINT_PROTOCOL_PC",
    KeyDomain.OSINT_PROTOCOL_CLOUD: "GOOGLE_API_KEY_OSINT_PROTOCOL_CLOUD",
    KeyDomain.SUPERVISOR_HEAL_PC: "GOOGLE_API_KEY_SUPERVISOR_HEAL_PC",
    KeyDomain.SUPERVISOR_HEAL_CLOUD: "GOOGLE_API_KEY_SUPERVISOR_HEAL_CLOUD",
    KeyDomain.HEAD_SUPERVISOR_PC: "GOOGLE_API_KEY_HEAD_SUPERVISOR_PC",
    KeyDomain.HEAD_SUPERVISOR_CLOUD: "GOOGLE_API_KEY_HEAD_SUPERVISOR_CLOUD",
    KeyDomain.CLOUD_SYNC_DOMAIN: "GOOGLE_API_KEY_CLOUD_SYNC_DOMAIN",
}

CLOUD_ONLY_DOMAINS = {
    KeyDomain.CLOUD, KeyDomain.AUTO_UPGRADE_CLOUD, KeyDomain.RESEARCH_ENGINE,
    KeyDomain.CAPTCHA_SOLVER_CLOUD, KeyDomain.IMAGE_GENERATION_CLOUD,
    KeyDomain.OSINT_PROTOCOL_CLOUD, KeyDomain.SUPERVISOR_HEAL_CLOUD,
    KeyDomain.HEAD_SUPERVISOR_CLOUD, KeyDomain.CLOUD_SYNC_DOMAIN
}

# Legacy defaults. These are the values the vault has always used, and they are
# kept ONLY so existing installs stay decryptable — changing them would make
# every already-encrypted secret unreadable.
#
# They are not secret: both the passphrase and the salt were hardcoded in this
# file, and this file ships to the Space. Anyone who could read the source could
# re-derive the master key with PBKDF2 and decrypt the whole vault, with or
# without the vault_keys row. That is why a public repo containing memory.db was
# a total compromise rather than merely an encrypted-blob exposure.
#
# Set VAULT_MASTER_PASSWORD (and ideally PBKDF2_SALT) as a Space Secret and the
# key is no longer derivable from anything published — the repo can then hold the
# vault safely, because the KDF inputs live only in the environment.
_LEGACY_PASSPHRASE = "always iron man"
_LEGACY_SALT = b"JARVIS_OMEGA_SALT_V15"


def _vault_passphrase() -> str:
    return os.environ.get("VAULT_MASTER_PASSWORD", "").strip() or _LEGACY_PASSPHRASE


def _vault_salt() -> bytes:
    configured = os.environ.get("PBKDF2_SALT", "").strip()
    return configured.encode("utf-8") if configured else _LEGACY_SALT


def generate_key_from_password(password: str, salt: bytes | None = None) -> str:
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt if salt is not None else _vault_salt(),
        iterations=480000,
    )
    return base64.urlsafe_b64encode(kdf.derive(password.encode())).decode('utf-8')

def _init_vault_db() -> str:
    """Initializes the vault keys table and retrieves or generates the AES Master Key."""
    try:
        db_path = get_db_path()
        with sqlite3.connect(db_path) as conn:
            conn.execute('PRAGMA journal_mode=WAL')
            conn.execute("CREATE TABLE IF NOT EXISTS vault_keys (id INTEGER PRIMARY KEY, key_data TEXT)")
            conn.execute("CREATE TABLE IF NOT EXISTS vault_secrets (key_name TEXT PRIMARY KEY, encrypted_value BLOB)")
            conn.execute("CREATE TABLE IF NOT EXISTS vault_credentials (id INTEGER PRIMARY KEY AUTOINCREMENT, site TEXT, username TEXT, password_enc BLOB, captured_at TEXT, source TEXT CHECK(source IN ('sentinel','manual','internet')))")
            cursor = conn.cursor()
            cursor.execute("SELECT key_data FROM vault_keys WHERE id = 1")
            row = cursor.fetchone()
            
            target_key = generate_key_from_password(_vault_passphrase())

            # When the passphrase comes from the environment, do NOT persist the
            # derived key back into the database. Storing it there would put the
            # key and the ciphertext in the same file again and undo the whole
            # point of moving the passphrase into a Space Secret.
            if os.environ.get("VAULT_MASTER_PASSWORD", "").strip():
                return target_key

            # No VAULT_MASTER_PASSWORD in the environment, so the derived key has to be
            # persisted — into the SAME SQLite file that holds the ciphertext it
            # protects. That combination means the file alone is enough to decrypt every
            # secret in it: verified in Part 28 by decrypting real vault_secrets rows
            # using nothing but a copy of memory.db (no password, no env var, no source).
            # The deploy guard in scripts/deploy_hf_space.py does keep this file out of
            # the public Space (name + suffix + "SQLite format 3" header sniff), so the
            # exposure is local rather than published — but it is still real, and it is
            # silent. Say so loudly instead of letting "the vault is encrypted" stand
            # unqualified.
            logging.warning(
                "VAULT: VAULT_MASTER_PASSWORD is not set, so the vault key is being "
                "stored in vault_keys inside the same database as vault_secrets. "
                "Anyone with a copy of that file can decrypt every secret in it. "
                "Set VAULT_MASTER_PASSWORD (env or Space secret) so the key is derived "
                "at runtime and never written to disk beside the ciphertext."
            )
            if not row:
                conn.execute("INSERT INTO vault_keys (id, key_data) VALUES (1, ?)", (target_key,))
                conn.commit()
                return target_key
            elif row[0] != target_key:
                conn.execute("UPDATE vault_keys SET key_data = ? WHERE id = 1", (target_key,))
                conn.commit()
                return target_key
                
            return row[0]
    except Exception as e:
        logging.error(f"Vault DB Error: {e}")
        return generate_key_from_password(_vault_passphrase())

MASTER_KEY = _init_vault_db()
cipher_suite = Fernet(MASTER_KEY.encode('utf-8'))

def set_secret(key_name: str, plain_value: str):
    """Encrypts and stores a string secret in the memory.db vault."""
    encrypted_value = cipher_suite.encrypt(plain_value.encode('utf-8'))
    db_path = get_db_path()
    with sqlite3.connect(db_path) as conn:
        conn.execute('PRAGMA journal_mode=WAL')
        conn.execute("INSERT OR REPLACE INTO vault_secrets (key_name, encrypted_value) VALUES (?, ?)", (key_name, encrypted_value))
        conn.commit()

def _legacy_cipher() -> Fernet:
    # Derivation frozen to the pre-env-passphrase constants: secrets written before
    # VAULT_MASTER_PASSWORD/PBKDF2_SALT were configured are only readable this way.
    return Fernet(generate_key_from_password(_LEGACY_PASSPHRASE, _LEGACY_SALT).encode('utf-8'))


def get_secret(key_name: str) -> str:
    """Retrieves and decrypts a string secret from the memory.db vault. Returns None if not found."""
    # First check actual environment variables for cloud deployment
    env_val = os.environ.get(key_name)
    if env_val:
        return env_val

    db_path = get_db_path()
    try:
        with sqlite3.connect(db_path) as conn:
            conn.execute('PRAGMA journal_mode=WAL')
            cursor = conn.cursor()
            cursor.execute("SELECT encrypted_value FROM vault_secrets WHERE key_name = ?", (key_name,))
            row = cursor.fetchone()
            if row and row[0]:
                try:
                    return cipher_suite.decrypt(row[0]).decode('utf-8')
                except Exception:
                    # Entry predates the configured passphrase (encrypted under the
                    # legacy key). Setting VAULT_MASTER_PASSWORD used to silently
                    # lose every such secret — fall back, and migrate it forward so
                    # the next read succeeds under the configured key directly.
                    value = _legacy_cipher().decrypt(row[0]).decode('utf-8')
                    if os.environ.get("VAULT_MASTER_PASSWORD", "").strip():
                        conn.execute(
                            "UPDATE vault_secrets SET encrypted_value = ? WHERE key_name = ?",
                            (cipher_suite.encrypt(value.encode('utf-8')), key_name))
                        conn.commit()
                        logging.info(f"Vault: migrated legacy-encrypted secret {key_name} to the configured passphrase.")
                    return value
    except Exception as e:
        logging.error(f"Error decrypting secret {key_name}: {e}")
    return None

def resolve_vault_key(domain: KeyDomain) -> str:
    """Dynamically resolves a key based on the environment and domain constraints."""
    import os
    env_var_name = KEY_DOMAIN_ENV_MAP[domain]
    
    # HF Spaces use actual environment variables natively
    key = os.environ.get(env_var_name)
    if not key:
        # Fallback to local SQLite vault
        key = get_secret(env_var_name)
        
    # BACKWARD COMPATIBILITY / SINGLE-KEY FALLBACK: every domain here is a Google
    # (Gemini) key slot. If a caller hasn't provisioned the dedicated per-domain
    # key, fall back to the generic GEMINI_API_KEY rather than crashing the
    # feature. The dedicated key is always preferred (checked above) so the
    # multi-key rate-limit partitioning still applies whenever it's configured;
    # this only rescues the common single-key desktop user, for whom otherwise
    # only chat/cloud worked and gaming-coach / image-gen / self-heal / etc.
    # raised MissingDomainKeyError.
    if not key:
        key = os.environ.get("GEMINI_API_KEY") or get_secret("GEMINI_API_KEY")
        
    if not key:
        location = "HF Space Secrets" if domain in CLOUD_ONLY_DOMAINS else "local .env/vault"
        raise MissingDomainKeyError(
            f"No API key configured for domain '{domain.value}'. "
            f"Add {env_var_name} to your {location} before this feature will work."
        )
    return key

def get_recovery_key() -> str:
    """Returns the Master Key. This should be kept strictly secret."""
    return MASTER_KEY

def encrypt_file(source_path: str, dest_path: str):
    """Encrypts a file from the PC and writes it to the Pendrive target."""
    with open(source_path, 'rb') as f:
        file_data = f.read()
    
    encrypted_data = cipher_suite.encrypt(file_data)
    
    with open(dest_path, 'wb') as f:
        f.write(encrypted_data)
    
    logging.info(f"VAULT: Successfully encrypted {source_path} to {dest_path}")

def decrypt_file(source_path: str, dest_path: str):
    """Decrypts a file from the Pendrive target and restores it to the PC."""
    with open(source_path, 'rb') as f:
        encrypted_data = f.read()
        
    try:
        decrypted_data = cipher_suite.decrypt(encrypted_data)
    except Exception as e:
        logging.error(f"VAULT DECRYPTION ERROR: Invalid key or corrupted file. {e}")
        raise ValueError("Decryption failed. The file is corrupted or you are using the wrong JARVIS PC.")
        
    with open(dest_path, 'wb') as f:
        f.write(decrypted_data)

    logging.info(f"VAULT: Successfully decrypted {source_path} to {dest_path}")

import shutil
import tempfile

def encrypt_directory(source_dir: str, dest_vault_file: str):
    """Zips an entire directory, encrypts it, and writes it to the pendrive."""
    with tempfile.TemporaryDirectory() as tmpdirname:
        zip_path = os.path.join(tmpdirname, "archive")
        shutil.make_archive(zip_path, 'zip', source_dir)
        zip_file = zip_path + ".zip"
        
        encrypt_file(zip_file, dest_vault_file)
        logging.info(f"VAULT: Successfully zipped and encrypted directory {source_dir} to {dest_vault_file}")

def decrypt_directory(source_vault_file: str, dest_dir: str):
    """Decrypts a vault file from the pendrive and extracts the directory to the PC."""
    with tempfile.TemporaryDirectory() as tmpdirname:
        zip_file = os.path.join(tmpdirname, "decrypted.zip")
        decrypt_file(source_vault_file, zip_file)
        
        shutil.unpack_archive(zip_file, dest_dir, 'zip')
        logging.info(f"VAULT: Successfully decrypted and unpacked {source_vault_file} to {dest_dir}")