"""
Runtime Key Vault — Production-grade local storage for user API keys.

Encryption: Fernet (AES-128-CBC + HMAC-SHA256) with machine-specific key
derivation.  The vault file at ~/.brynq/vault.json is useless without the
exact machine (hostname + MAC) and optional user passphrase.

Key derivation chain:
    machine_id = SHA-256(hostname || first-MAC)
    master     = PBKDF2-HMAC-SHA256(machine_id || passphrase, salt, 480_000)
    fernet_key = urlsafe_b64encode(master)

Security guarantees:
    - Keys are NEVER logged (even at DEBUG level)
    - Keys are NEVER sent to brynq.ai servers
    - Memory is zeroed after use where possible (best-effort in Python)
    - Vault file permissions: 0600 (user-only read/write)
    - Salt is per-vault (stored alongside encrypted data)
    - Without the exact machine + passphrase, the vault file is inert
"""

from __future__ import annotations

import base64
import hashlib
import json
import logging
import os
import secrets
import shutil
import socket
import ssl
import urllib.error
import urllib.request
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

from cryptography.fernet import Fernet, InvalidToken

log = logging.getLogger("brynq.runtime.vault")

# ── Constants ────────────────────────────────────────────────────────────

VAULT_FILE = Path.home() / ".brynq" / "vault.json"
VAULT_VERSION = 2

_KDF_ITERATIONS = 480_000  # OWASP 2023 recommendation for PBKDF2-SHA256
_SALT_LEN = 32  # 256-bit salt
_KEY_LEN = 32  # 256-bit derived key (Fernet uses first 128 for AES, rest for HMAC)

# Providers we know how to test
_KNOWN_PROVIDERS = frozenset({
    "claude", "anthropic",
    "gemini", "google",
    "gpt", "openai",
    "ollama",
    "cohere",
    "mistral",
    "groq",
    "deepseek",
    "perplexity",
    "together",
    "fireworks",
    "replicate",
    "huggingface",
})


# ── Data Classes ─────────────────────────────────────────────────────────


@dataclass(frozen=True)
class KeyInfo:
    """Read-only view of a stored key — NEVER contains the full key."""

    provider: str
    masked_key: str
    added_at: str
    last_used: Optional[str] = None
    provider_label: Optional[str] = None  # e.g. "anthropic", "google"


# ── Machine Fingerprint ─────────────────────────────────────────────────


def _get_machine_id() -> bytes:
    """Derive a stable machine identifier from hostname + MAC address.

    This makes the encryption key physically bound to this machine.
    Moving the vault file to another machine renders it undecryptable.
    """
    hostname = socket.gethostname()
    # Get the first real MAC address (not all-zeros loopback)
    mac = uuid.getnode()
    mac_hex = format(mac, "012x")
    raw = f"brynq-vault-v2:{hostname}:{mac_hex}"
    return hashlib.sha256(raw.encode("utf-8")).digest()


def _derive_fernet_key(salt: bytes, passphrase: Optional[str] = None) -> bytes:
    """Derive a Fernet-compatible key from machine ID + optional passphrase.

    Returns url-safe base64-encoded 32-byte key as Fernet requires.
    """
    machine_id = _get_machine_id()
    # Combine machine ID with passphrase (or empty bytes)
    ikm = machine_id + (passphrase or "").encode("utf-8")
    derived = hashlib.pbkdf2_hmac(
        "sha256",
        ikm,
        salt,
        iterations=_KDF_ITERATIONS,
        dklen=_KEY_LEN,
    )
    return base64.urlsafe_b64encode(derived)


# ── Key Masking ──────────────────────────────────────────────────────────


def _mask_key(api_key: str) -> str:
    """Mask an API key for safe display. NEVER returns the full key."""
    if len(api_key) <= 6:
        return "***"
    if len(api_key) <= 8:
        return api_key[:3] + "..." + api_key[-3:]
    return api_key[:4] + "..." + api_key[-4:]


# ── Provider Metadata ────────────────────────────────────────────────────

_PROVIDER_LABELS: dict[str, str] = {
    "claude": "anthropic",
    "anthropic": "anthropic",
    "gemini": "google",
    "google": "google",
    "gpt": "openai",
    "openai": "openai",
    "ollama": "ollama",
    "cohere": "cohere",
    "mistral": "mistral",
    "groq": "groq",
    "deepseek": "deepseek",
    "perplexity": "perplexity",
    "together": "together",
    "fireworks": "fireworks",
    "replicate": "replicate",
    "huggingface": "huggingface",
}

# Friendly aliases → canonical provider names
_PROVIDER_ALIASES: dict[str, str] = {
    "claude": "anthropic",
    "gpt": "openai",
    "gemini": "google",
}


# ── Key Vault ────────────────────────────────────────────────────────────


class KeyVault:
    """Secure local key storage.  Keys are encrypted at rest with
    machine-specific Fernet encryption.

    Usage::

        vault = KeyVault()
        vault.store_key("claude", "sk-ant-api03-...")
        key = vault.get_key("claude")
        vault.list_keys()       # [KeyInfo(..., masked_key="sk-a...xxxx")]
        vault.delete_key("claude")

    With passphrase (extra security)::

        vault = KeyVault()
        vault.store_key("claude", "sk-ant-...", passphrase="my-secret")
        key = vault.get_key("claude", passphrase="my-secret")
    """

    def __init__(self, vault_path: Optional[Path] = None) -> None:
        self._vault_path = vault_path or VAULT_FILE

    # ── Public API ───────────────────────────────────────────────────

    def store_key(
        self,
        provider: str,
        api_key: str,
        passphrase: Optional[str] = None,
    ) -> None:
        """Encrypt and store an API key for a provider.

        Args:
            provider: Provider name (e.g. "claude", "gemini", "gpt").
            api_key: The raw API key to store.
            passphrase: Optional passphrase for extra encryption strength.

        Raises:
            ValueError: If provider is empty or api_key is empty.
        """
        provider = provider.strip().lower()
        if not provider:
            raise ValueError("Provider name cannot be empty")

        # Normalize friendly aliases to canonical provider names
        provider = _PROVIDER_ALIASES.get(provider, provider)

        api_key = api_key.strip()
        if not api_key:
            raise ValueError("API key cannot be empty")

        # Warn on unexpected key format (still store regardless)
        _KEY_PREFIX_RULES: dict[str, tuple[str, str]] = {
            "claude": ("sk-ant-", "Claude keys typically start with 'sk-ant-'"),
            "anthropic": ("sk-ant-", "Anthropic keys typically start with 'sk-ant-'"),
            "gemini": ("AIza", "Google keys typically start with 'AIza'"),
            "google": ("AIza", "Google keys typically start with 'AIza'"),
            "gpt": ("sk-", "OpenAI keys typically start with 'sk-'"),
            "openai": ("sk-", "OpenAI keys typically start with 'sk-'"),
        }
        prefix_rule = _KEY_PREFIX_RULES.get(provider)
        if prefix_rule and not api_key.startswith(prefix_rule[0]):
            log.warning(
                "Key for '%s' may be invalid: %s", provider, prefix_rule[1]
            )

        data = self._load()
        salt = base64.b64decode(data["salt"])

        # Duplicate detection — skip re-encrypting if same key already stored
        existing_b64 = data["keys"].get(provider)
        if existing_b64:
            try:
                fernet_key_check = _derive_fernet_key(salt, passphrase)
                f_check = Fernet(fernet_key_check)
                existing_plain = f_check.decrypt(
                    base64.b64decode(existing_b64)
                ).decode("utf-8")
                if existing_plain == api_key:
                    log.info("Key already stored for %s", provider)
                    return
            except (InvalidToken, Exception):
                pass  # Can't decrypt (wrong passphrase?) — proceed to overwrite

        # Encrypt the key
        fernet_key = _derive_fernet_key(salt, passphrase)
        f = Fernet(fernet_key)
        encrypted = f.encrypt(api_key.encode("utf-8"))
        data["keys"][provider] = base64.b64encode(encrypted).decode("utf-8")

        # Update metadata
        now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
        if provider not in data["metadata"]:
            data["metadata"][provider] = {}
        data["metadata"][provider]["added_at"] = now
        data["metadata"][provider]["provider"] = _PROVIDER_LABELS.get(
            provider, provider
        )

        self._save(data)
        # Do NOT log the key value — ever
        log.info("Stored key for provider: %s", provider)

    def get_key(
        self,
        provider: str,
        passphrase: Optional[str] = None,
    ) -> Optional[str]:
        """Decrypt and return a stored API key.

        Args:
            provider: Provider name.
            passphrase: Passphrase used when the key was stored.

        Returns:
            The decrypted API key string, or None if not found / decryption fails.
        """
        provider = provider.strip().lower()
        provider = _PROVIDER_ALIASES.get(provider, provider)
        data = self._load()

        encrypted_b64 = data["keys"].get(provider)
        if not encrypted_b64:
            return None

        salt = base64.b64decode(data["salt"])
        fernet_key = _derive_fernet_key(salt, passphrase)
        f = Fernet(fernet_key)

        try:
            encrypted = base64.b64decode(encrypted_b64)
            decrypted = f.decrypt(encrypted).decode("utf-8")
        except InvalidToken:
            log.error(
                "Vault corruption or wrong passphrase for provider '%s' "
                "— Fernet decryption failed. The key entry may be corrupted.",
                provider,
            )
            return None
        except Exception:
            log.error(
                "Failed to decrypt key for provider '%s': unexpected error",
                provider,
            )
            return None

        # Update last_used timestamp
        now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
        if provider in data["metadata"]:
            data["metadata"][provider]["last_used"] = now
            self._save(data)

        return decrypted

    def list_keys(self) -> list[KeyInfo]:
        """List all stored keys with masked values and metadata.

        Returns a list of KeyInfo objects.  The full key is NEVER exposed.
        """
        data = self._load()
        result: list[KeyInfo] = []

        for provider in sorted(data["keys"]):
            meta = data["metadata"].get(provider, {})
            # Attempt decryption (no passphrase) to get a masked preview
            encrypted_b64 = data["keys"][provider]
            masked = "***"
            try:
                salt = base64.b64decode(data["salt"])
                fernet_key = _derive_fernet_key(salt, None)
                f = Fernet(fernet_key)
                encrypted = base64.b64decode(encrypted_b64)
                plain = f.decrypt(encrypted).decode("utf-8")
                masked = _mask_key(plain)
            except Exception:
                # Key was stored with a passphrase — cannot mask without it
                masked = "****...****"

            result.append(
                KeyInfo(
                    provider=provider,
                    masked_key=masked,
                    added_at=meta.get("added_at", "unknown"),
                    last_used=meta.get("last_used"),
                    provider_label=meta.get("provider"),
                )
            )

        return result

    def delete_key(self, provider: str) -> bool:
        """Remove a stored key.

        Args:
            provider: Provider name.

        Returns:
            True if a key was deleted, False if the provider was not found.
        """
        provider = provider.strip().lower()
        provider = _PROVIDER_ALIASES.get(provider, provider)
        data = self._load()

        if provider not in data["keys"]:
            return False

        del data["keys"][provider]
        data["metadata"].pop(provider, None)
        self._save(data)
        log.info("Deleted key for provider: %s", provider)
        return True

    def test_key(
        self, provider: str, passphrase: Optional[str] = None
    ) -> bool:
        """Make a lightweight API call to verify a stored key works.

        Args:
            provider: Provider name.
            passphrase: Passphrase if the key was stored with one.

        Returns:
            True if the API call succeeds, False otherwise.
        """
        api_key = self.get_key(provider, passphrase)
        if not api_key:
            return False

        try:
            result = _test_provider_key(provider, api_key)
        finally:
            # Best-effort memory clearing
            _zero_string(api_key)
        return result

    def rotate_encryption(
        self,
        old_passphrase: Optional[str] = None,
        new_passphrase: Optional[str] = None,
    ) -> int:
        """Re-encrypt all keys with a new passphrase.

        Decrypts every key with *old_passphrase*, generates a new salt,
        and re-encrypts with *new_passphrase*.

        Args:
            old_passphrase: Current passphrase (None if none was set).
            new_passphrase: New passphrase (None to remove passphrase protection).

        Returns:
            Number of keys successfully rotated.

        Raises:
            ValueError: If any key fails to decrypt with the old passphrase.
        """
        data = self._load()
        old_salt = base64.b64decode(data["salt"])
        old_fernet_key = _derive_fernet_key(old_salt, old_passphrase)
        old_f = Fernet(old_fernet_key)

        # Decrypt all keys first — fail fast if old passphrase is wrong
        plaintext_keys: dict[str, str] = {}
        for provider, encrypted_b64 in data["keys"].items():
            try:
                encrypted = base64.b64decode(encrypted_b64)
                plaintext_keys[provider] = old_f.decrypt(encrypted).decode("utf-8")
            except (InvalidToken, Exception) as exc:
                raise ValueError(
                    f"Failed to decrypt key for '{provider}' with the provided passphrase"
                ) from exc

        # Generate new salt and re-encrypt everything
        new_salt = secrets.token_bytes(_SALT_LEN)
        new_fernet_key = _derive_fernet_key(new_salt, new_passphrase)
        new_f = Fernet(new_fernet_key)

        data["salt"] = base64.b64encode(new_salt).decode("utf-8")
        for provider, plaintext in plaintext_keys.items():
            encrypted = new_f.encrypt(plaintext.encode("utf-8"))
            data["keys"][provider] = base64.b64encode(encrypted).decode("utf-8")

        self._save(data)

        # Best-effort clear plaintext keys from memory
        for key in plaintext_keys.values():
            _zero_string(key)
        plaintext_keys.clear()

        count = len(data["keys"])
        log.info("Rotated encryption for %d keys", count)
        return count

    def export_encrypted(self, passphrase: str) -> bytes:
        """Export the vault as a portable encrypted blob.

        Decrypts all keys (using no passphrase — if keys were stored with a
        passphrase, call ``rotate_encryption`` first to remove it), then wraps
        everything in a passphrase-only Fernet envelope that is **not** bound
        to any machine.

        Args:
            passphrase: Passphrase to encrypt the export with.

        Returns:
            Encrypted bytes suitable for writing to a backup file.

        Raises:
            ValueError: If passphrase is empty.
        """
        if not passphrase:
            raise ValueError("Export passphrase cannot be empty")

        data = self._load()
        salt = base64.b64decode(data["salt"])
        fernet_key = _derive_fernet_key(salt, None)
        f = Fernet(fernet_key)

        # Decrypt all keys to plaintext for portable export
        plaintext_keys: dict[str, str] = {}
        for provider, encrypted_b64 in data["keys"].items():
            try:
                encrypted_bytes = base64.b64decode(encrypted_b64)
                plaintext_keys[provider] = f.decrypt(encrypted_bytes).decode("utf-8")
            except Exception:
                log.warning(
                    "Could not decrypt key for %s during export — skipping", provider
                )

        export_payload = {
            "version": VAULT_VERSION,
            "plaintext_keys": plaintext_keys,
            "metadata": data.get("metadata", {}),
        }

        payload_bytes = json.dumps(export_payload, indent=2).encode("utf-8")

        # Derive a portable key from passphrase alone (no machine ID)
        export_salt = secrets.token_bytes(_SALT_LEN)
        derived = hashlib.pbkdf2_hmac(
            "sha256",
            passphrase.encode("utf-8"),
            export_salt,
            iterations=_KDF_ITERATIONS,
            dklen=_KEY_LEN,
        )
        export_fernet_key = base64.urlsafe_b64encode(derived)
        ef = Fernet(export_fernet_key)
        encrypted_blob = ef.encrypt(payload_bytes)

        # Clear plaintext keys from memory
        for v in plaintext_keys.values():
            _zero_string(v)
        plaintext_keys.clear()

        header = b"BRYNQ_EXPORT_V1:"
        return header + bytes([_SALT_LEN]) + export_salt + encrypted_blob

    def import_encrypted(self, data: bytes, passphrase: str) -> int:
        """Import a vault from an encrypted export blob.

        Decrypts the portable export, then re-encrypts all keys for *this*
        machine.  Existing keys for the same providers are **overwritten**.

        Args:
            data: Encrypted bytes from :meth:`export_encrypted`.
            passphrase: Passphrase used during export.

        Returns:
            Number of keys imported.

        Raises:
            ValueError: If decryption fails or format is invalid.
        """
        if not passphrase:
            raise ValueError("Import passphrase cannot be empty")

        header = b"BRYNQ_EXPORT_V1:"
        if not data.startswith(header):
            raise ValueError("Invalid export format — not a Brynq vault export")

        rest = data[len(header):]
        if len(rest) < 2:
            raise ValueError("Invalid export format — truncated")

        salt_len = rest[0]
        if len(rest) < 1 + salt_len + 1:
            raise ValueError("Invalid export format — truncated salt")

        export_salt = rest[1 : 1 + salt_len]
        encrypted_blob = rest[1 + salt_len :]

        # Derive the portable key from passphrase
        derived = hashlib.pbkdf2_hmac(
            "sha256",
            passphrase.encode("utf-8"),
            export_salt,
            iterations=_KDF_ITERATIONS,
            dklen=_KEY_LEN,
        )
        fernet_key = base64.urlsafe_b64encode(derived)
        f = Fernet(fernet_key)

        try:
            decrypted = f.decrypt(encrypted_blob)
        except (InvalidToken, Exception) as exc:
            raise ValueError(
                "Failed to decrypt export — wrong passphrase?"
            ) from exc

        try:
            imported_data = json.loads(decrypted.decode("utf-8"))
        except (json.JSONDecodeError, UnicodeDecodeError) as exc:
            raise ValueError("Corrupted export data") from exc

        plaintext_keys = imported_data.get("plaintext_keys", {})
        if not plaintext_keys:
            raise ValueError(
                "Export contains no keys (or incompatible format)"
            )

        # Re-encrypt each key for this machine
        count = 0
        for prov, plain_key in plaintext_keys.items():
            self.store_key(prov, plain_key)
            count += 1

        # Merge metadata (preserve imported added_at if older)
        if "metadata" in imported_data:
            current = self._load()
            for prov, meta in imported_data["metadata"].items():
                if prov in current["metadata"]:
                    if "added_at" in meta:
                        current["metadata"][prov].setdefault(
                            "added_at", meta["added_at"]
                        )
                    if "provider" in meta:
                        current["metadata"][prov].setdefault(
                            "provider", meta["provider"]
                        )
            self._save(current)

        log.info("Imported %d keys", count)
        return count

    # ── Private Helpers ──────────────────────────────────────────────

    def _load(self) -> dict[str, Any]:
        """Load the vault from disk, creating a new empty vault if needed."""
        if not self._vault_path.exists():
            return self._create_empty_vault()

        try:
            raw = self._vault_path.read_text("utf-8")
            data = json.loads(raw)
        except (json.JSONDecodeError, OSError):
            log.warning("Corrupted vault file — creating new vault")
            return self._create_empty_vault()

        # Ensure required fields exist
        if "version" not in data or "salt" not in data or "keys" not in data:
            log.warning("Invalid vault structure — creating new vault")
            return self._create_empty_vault()

        data.setdefault("metadata", {})
        return data

    def _save(self, data: dict[str, Any]) -> None:
        """Save the vault to disk with restrictive permissions."""
        self._vault_path.parent.mkdir(parents=True, exist_ok=True)

        # Backup existing vault before writing
        if self._vault_path.exists():
            backup_path = self._vault_path.with_suffix(".bak")
            try:
                shutil.copy2(self._vault_path, backup_path)
            except OSError:
                log.warning("Failed to create vault backup at %s", backup_path)
        content = json.dumps(data, indent=2) + "\n"
        self._vault_path.write_text(content, encoding="utf-8")

        # Set file permissions to 0600 (user-only) on Unix
        try:
            if os.name != "nt":
                self._vault_path.chmod(0o600)
        except OSError:
            pass

    def _create_empty_vault(self) -> dict[str, Any]:
        """Create and persist a new empty vault with a fresh salt."""
        salt = secrets.token_bytes(_SALT_LEN)
        data: dict[str, Any] = {
            "version": VAULT_VERSION,
            "salt": base64.b64encode(salt).decode("utf-8"),
            "keys": {},
            "metadata": {},
        }
        self._save(data)
        return data


# ── Key Testing (Lightweight API Calls) ──────────────────────────────────


def _test_provider_key(provider: str, api_key: str) -> bool:
    """Make a minimal API call to verify the key works.

    Uses stdlib ``urllib`` only — no extra dependencies.
    """
    provider = provider.strip().lower()

    test_configs: dict[str, dict[str, Any]] = {
        "claude": {
            "url": "https://api.anthropic.com/v1/models",
            "headers": {
                "x-api-key": api_key,
                "anthropic-version": "2023-06-01",
            },
        },
        "anthropic": {
            "url": "https://api.anthropic.com/v1/models",
            "headers": {
                "x-api-key": api_key,
                "anthropic-version": "2023-06-01",
            },
        },
        "gemini": {
            "url": f"https://generativelanguage.googleapis.com/v1beta/models?key={api_key}",
            "headers": {},
        },
        "google": {
            "url": f"https://generativelanguage.googleapis.com/v1beta/models?key={api_key}",
            "headers": {},
        },
        "gpt": {
            "url": "https://api.openai.com/v1/models",
            "headers": {"Authorization": f"Bearer {api_key}"},
        },
        "openai": {
            "url": "https://api.openai.com/v1/models",
            "headers": {"Authorization": f"Bearer {api_key}"},
        },
        "ollama": {
            "url": "http://localhost:11434/api/tags",
            "headers": {},
        },
    }

    config = test_configs.get(provider)
    if not config:
        log.warning("No test endpoint configured for provider: %s", provider)
        return False

    try:
        req = urllib.request.Request(config["url"], method="GET")
        for header_name, header_val in config["headers"].items():
            req.add_header(header_name, header_val)

        # Create an SSL context that validates certificates
        ctx = ssl.create_default_context()
        with urllib.request.urlopen(req, timeout=10, context=ctx) as resp:
            return resp.status == 200
    except urllib.error.HTTPError as e:
        # 401/403 = bad key, anything else might be transient
        if e.code in (401, 403):
            return False
        log.warning("HTTP %d testing %s key", e.code, provider)
        return False
    except Exception as exc:
        log.warning("Failed to test %s key: %s", provider, type(exc).__name__)
        return False


# ── Memory Safety ────────────────────────────────────────────────────────


def _zero_string(s: str) -> None:
    """Best-effort attempt to reduce key lifetime in memory.

    Python strings are immutable so true zeroing is not reliably possible.
    We avoid unsafe ctypes.memset (which causes access violations on
    Python 3.14+ during GC) and instead rely on:
      - Deleting references so GC can reclaim the memory promptly
      - Caller clearing dict entries after use

    This is defense-in-depth.  The real protection is Fernet encryption
    at rest and never logging key values.
    """
    # Intentionally a no-op beyond documentation.  The caller is expected
    # to ``del`` its reference and clear any dicts holding the value.
    # ctypes.memset on string internals is too fragile across CPython
    # versions and causes segfaults during garbage collection.
    pass
