"""
ID Wispera Vault
AES-256-GCM encrypted storage for passports
"""

import json
import os
import secrets
from datetime import datetime
from pathlib import Path
from typing import Optional

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
from cryptography.hazmat.backends import default_backend

from id_wispera.types import Passport, AuditEntry


# Vault file format version
VAULT_VERSION = 1

# Scrypt parameters
SCRYPT_N = 2**14
SCRYPT_R = 8
SCRYPT_P = 1


class Vault:
    """
    Encrypted vault for storing passports.
    Uses AES-256-GCM for encryption with Scrypt key derivation.
    """

    def __init__(self, path: str, key: bytes):
        self.path = path
        self._key = key
        self._passports: dict[str, Passport] = {}
        self._audit_log: list[AuditEntry] = []
        self._salt: bytes = b""
        self._loaded = False

    @property
    def is_unlocked(self) -> bool:
        return self._loaded

    @property
    def passport_count(self) -> int:
        return len(self._passports)

    def _load(self) -> None:
        """Load and decrypt vault data"""
        if self._loaded:
            return

        with open(self.path, "r") as f:
            data = json.load(f)

        self._salt = bytes.fromhex(data["salt"])
        iv = bytes.fromhex(data["iv"])

        # Decrypt passports
        aesgcm = AESGCM(self._key)
        encrypted_passports = bytes.fromhex(data["encryptedPassports"])
        decrypted = aesgcm.decrypt(iv, encrypted_passports, None)
        passports_data = json.loads(decrypted.decode("utf-8"))

        for p in passports_data:
            passport = Passport(**p)
            self._passports[passport.id] = passport

        # Decrypt audit log
        if data.get("encryptedAuditLog"):
            encrypted_audit = bytes.fromhex(data["encryptedAuditLog"])
            decrypted_audit = aesgcm.decrypt(iv, encrypted_audit, None)
            audit_data = json.loads(decrypted_audit.decode("utf-8"))
            self._audit_log = [AuditEntry(**e) for e in audit_data]

        self._loaded = True

    def save(self) -> None:
        """Encrypt and save vault data"""
        iv = secrets.token_bytes(12)
        aesgcm = AESGCM(self._key)

        # Serialize passports
        passports_json = json.dumps(
            [p.model_dump(by_alias=True, mode="json") for p in self._passports.values()],
            default=str
        )
        encrypted_passports = aesgcm.encrypt(iv, passports_json.encode("utf-8"), None)

        # Serialize audit log
        audit_json = json.dumps(
            [e.model_dump(by_alias=True, mode="json") for e in self._audit_log],
            default=str
        )
        encrypted_audit = aesgcm.encrypt(iv, audit_json.encode("utf-8"), None)

        data = {
            "version": VAULT_VERSION,
            "salt": self._salt.hex(),
            "iv": iv.hex(),
            "encryptedPassports": encrypted_passports.hex(),
            "encryptedAuditLog": encrypted_audit.hex(),
            "metadata": {
                "createdAt": datetime.utcnow().isoformat(),
                "updatedAt": datetime.utcnow().isoformat(),
                "passportCount": len(self._passports),
            },
        }

        # Write atomically
        temp_path = f"{self.path}.tmp"
        with open(temp_path, "w") as f:
            json.dump(data, f, indent=2)
        os.replace(temp_path, self.path)

    def get_passport(self, passport_id: str) -> Optional[Passport]:
        """Get a passport by ID"""
        self._load()
        return self._passports.get(passport_id)

    def list_passports(self) -> list[Passport]:
        """List all passports"""
        self._load()
        return list(self._passports.values())

    def add_passport(self, passport: Passport) -> None:
        """Add a passport to the vault"""
        self._load()
        self._passports[passport.id] = passport
        self.save()

    def update_passport(self, passport: Passport) -> None:
        """Update an existing passport"""
        self._load()
        if passport.id not in self._passports:
            raise ValueError(f"Passport {passport.id} not found")
        passport.updated_at = datetime.utcnow()
        self._passports[passport.id] = passport
        self.save()

    def delete_passport(self, passport_id: str) -> None:
        """Delete a passport"""
        self._load()
        if passport_id in self._passports:
            del self._passports[passport_id]
            self.save()

    def add_audit_entry(self, entry: AuditEntry) -> None:
        """Add an audit log entry"""
        self._load()
        self._audit_log.append(entry)
        self.save()

    def get_audit_log(self, passport_id: Optional[str] = None) -> list[AuditEntry]:
        """Get audit log entries, optionally filtered by passport"""
        self._load()
        if passport_id:
            return [e for e in self._audit_log if e.passport_id == passport_id]
        return self._audit_log.copy()


def _derive_key(passphrase: str, salt: bytes) -> bytes:
    """Derive encryption key from passphrase using Scrypt"""
    kdf = Scrypt(
        salt=salt,
        length=32,
        n=SCRYPT_N,
        r=SCRYPT_R,
        p=SCRYPT_P,
        backend=default_backend(),
    )
    return kdf.derive(passphrase.encode("utf-8"))


def create_vault(passphrase: str, path: Optional[str] = None) -> Vault:
    """
    Create a new encrypted vault.

    Args:
        passphrase: The passphrase to encrypt the vault
        path: Optional custom path (defaults to ~/.id-wispera/vault.json)

    Returns:
        The created Vault object
    """
    if path is None:
        path = get_default_vault_path()

    # Ensure directory exists
    vault_dir = os.path.dirname(path)
    os.makedirs(vault_dir, mode=0o700, exist_ok=True)

    if os.path.exists(path):
        raise FileExistsError(f"Vault already exists at {path}")

    # Generate salt and derive key
    salt = secrets.token_bytes(32)
    key = _derive_key(passphrase, salt)

    # Create vault instance
    vault = Vault(path, key)
    vault._salt = salt
    vault._loaded = True
    vault.save()

    return vault


def unlock_vault(passphrase: str, path: Optional[str] = None) -> Vault:
    """
    Unlock an existing vault.

    Args:
        passphrase: The passphrase to decrypt the vault
        path: Optional custom path (defaults to ~/.id-wispera/vault.json)

    Returns:
        The unlocked Vault object
    """
    if path is None:
        path = get_default_vault_path()

    if not os.path.exists(path):
        raise FileNotFoundError(f"Vault not found at {path}")

    # Read vault to get salt
    with open(path, "r") as f:
        data = json.load(f)

    salt = bytes.fromhex(data["salt"])
    key = _derive_key(passphrase, salt)

    # Create vault and verify key by loading
    vault = Vault(path, key)
    vault._salt = salt

    try:
        vault._load()
    except Exception as e:
        raise ValueError("Invalid passphrase or corrupted vault") from e

    return vault


def vault_exists(path: Optional[str] = None) -> bool:
    """Check if a vault exists at the given path"""
    if path is None:
        path = get_default_vault_path()
    return os.path.exists(path)


def get_default_vault_path() -> str:
    """Get the default vault path"""
    return os.path.join(os.path.expanduser("~"), ".id-wispera", "vault.json")
