"""KeyManager -- encrypted BYOK API key storage.

Uses enhanced machine-bound encryption via :class:`KeyEncryption`.  The
encryption key is derived from hostname + MAC + install-id via PBKDF2.
Keys are persisted as an encrypted JSON blob in ``~/.llmhosts/keys.enc``.
"""

from __future__ import annotations

import json
import logging
import os
import stat
from datetime import datetime, timezone
from pathlib import Path

from cryptography.fernet import InvalidToken
from pydantic import ValidationError

from llmhosts.constants import DATA_DIR_NAME
from llmhosts.keys.encryption import KeyEncryption
from llmhosts.keys.models import (
    SUPPORTED_PROVIDERS,
    EncryptedKey,
    KeyInfo,
    KeyValidation,
    _mask_key,
)
from llmhosts.keys.providers import ProviderValidator

logger = logging.getLogger(__name__)

_KEYS_FILENAME = "keys.enc"


class KeyManager:
    """Manages encrypted storage of BYOK API keys.

    Uses :class:`KeyEncryption` for machine-bound Fernet encryption derived
    from PBKDF2 (hostname + MAC + install-id).  Supports key rotation via
    :meth:`rotate_encryption_key`.
    """

    def __init__(self, data_dir: Path | None = None) -> None:
        """Initialize with data directory (default ``~/.llmhosts``)."""
        self._data_dir = data_dir or (Path.home() / DATA_DIR_NAME)
        self._keys_file = self._data_dir / _KEYS_FILENAME
        self._encryption = KeyEncryption(self._data_dir)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _ensure_dirs(self) -> None:
        """Create data directory if it doesn't exist, with 700 permissions."""
        if not self._data_dir.exists():
            self._data_dir.mkdir(parents=True, exist_ok=True)
            os.chmod(self._data_dir, stat.S_IRWXU)  # 0o700
            logger.debug("Created data directory: %s", self._data_dir)

    def _get_fernet(self) -> KeyEncryption:
        """Return the :class:`KeyEncryption` instance.

        Ensures the data directory exists before returning.

        Returns:
            The ``KeyEncryption`` used for all encrypt/decrypt operations.
        """
        self._ensure_dirs()
        return self._encryption

    def _load_keys(self) -> dict[str, EncryptedKey]:
        """Load and decrypt the keys file.

        Uses :class:`MultiFernet` so that keys encrypted with a previous
        install-id (during rotation) can still be read.

        Returns:
            A dict mapping provider name -> ``EncryptedKey``.
            Returns empty dict if the file doesn't exist or is corrupted.
        """
        if not self._keys_file.exists():
            return {}

        enc = self._get_fernet()
        encrypted_blob = self._keys_file.read_bytes()

        if not encrypted_blob:
            return {}

        try:
            multi = enc.get_multi_fernet()
            decrypted = multi.decrypt(encrypted_blob)
        except InvalidToken:
            logger.error("Failed to decrypt keys file (key mismatch or corruption). Returning empty key store.")
            return {}

        try:
            raw: dict[str, dict] = json.loads(decrypted.decode("utf-8"))
        except (json.JSONDecodeError, UnicodeDecodeError):
            logger.error("Keys file contains invalid JSON after decryption. Returning empty key store.")
            return {}

        keys: dict[str, EncryptedKey] = {}
        for provider, data in raw.items():
            try:
                # Use strict=False to allow ISO datetime string coercion from JSON
                keys[provider] = EncryptedKey.model_validate(data, strict=False)
            except ValidationError as exc:
                logger.warning("Skipping corrupt entry for provider '%s': %s", provider, exc)
        return keys

    def _save_keys(self, keys: dict[str, EncryptedKey]) -> None:
        """Encrypt and save all keys to disk using the current derived key."""
        self._ensure_dirs()
        enc = self._get_fernet()
        fernet = enc.get_fernet()

        payload = {p: k.model_dump(mode="json") for p, k in keys.items()}
        plaintext = json.dumps(payload, default=str).encode("utf-8")
        encrypted = fernet.encrypt(plaintext)

        self._keys_file.write_bytes(encrypted)
        os.chmod(self._keys_file, stat.S_IRUSR | stat.S_IWUSR)  # 0o600
        logger.debug("Saved %d key(s) to %s", len(keys), self._keys_file)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def add_key(self, provider: str, api_key: str) -> None:
        """Add or update an API key for a provider.

        Args:
            provider: Provider name (must be in ``SUPPORTED_PROVIDERS``).
            api_key: The raw API key string.

        Raises:
            ValueError: If ``provider`` is not supported or ``api_key`` is empty.
        """
        provider = provider.strip().lower()
        if provider not in SUPPORTED_PROVIDERS:
            raise ValueError(f"Unsupported provider '{provider}'. Supported: {', '.join(SUPPORTED_PROVIDERS)}")
        if not api_key or not api_key.strip():
            raise ValueError("API key cannot be empty")

        enc = self._get_fernet()
        fernet = enc.get_fernet()
        encrypted_value = fernet.encrypt(api_key.strip().encode("utf-8")).decode("ascii")

        keys = self._load_keys()
        keys[provider] = EncryptedKey(
            provider=provider,
            encrypted_value=encrypted_value,
            added_at=datetime.now(timezone.utc),
        )
        self._save_keys(keys)
        logger.info("Stored key for provider '%s'", provider)

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

        Args:
            provider: Provider name to remove.

        Returns:
            ``True`` if the key existed and was removed, ``False`` otherwise.
        """
        provider = provider.strip().lower()
        keys = self._load_keys()
        if provider not in keys:
            return False
        del keys[provider]
        self._save_keys(keys)
        logger.info("Removed key for provider '%s'", provider)
        return True

    def get_key(self, provider: str) -> str | None:
        """Get a decrypted API key for a provider.

        Also updates the ``last_used`` timestamp.

        Args:
            provider: Provider name.

        Returns:
            The decrypted key string, or ``None`` if not found.
        """
        provider = provider.strip().lower()
        keys = self._load_keys()
        entry = keys.get(provider)
        if entry is None:
            return None

        enc = self._get_fernet()
        multi = enc.get_multi_fernet()
        try:
            decrypted = multi.decrypt(entry.encrypted_value.encode("ascii")).decode("utf-8")
        except InvalidToken:
            logger.error("Failed to decrypt key for provider '%s'", provider)
            return None

        # Update last_used timestamp
        entry.last_used = datetime.now(timezone.utc)
        self._save_keys(keys)
        return decrypted

    def list_providers(self) -> list[KeyInfo]:
        """List all stored providers WITHOUT exposing actual keys.

        Returns:
            A list of ``KeyInfo`` with provider name, masked key, dates, and
            validation status.
        """
        keys = self._load_keys()
        enc = self._get_fernet()
        multi = enc.get_multi_fernet()
        result: list[KeyInfo] = []

        for provider, entry in sorted(keys.items()):
            # Decrypt just enough to produce a masked version
            try:
                raw_key = multi.decrypt(entry.encrypted_value.encode("ascii")).decode("utf-8")
                masked = _mask_key(raw_key)
            except InvalidToken:
                masked = "***decrypt-error***"

            result.append(
                KeyInfo(
                    provider=provider,
                    masked_key=masked,
                    added_at=entry.added_at,
                    last_used=entry.last_used,
                    is_valid=entry.is_valid,
                )
            )
        return result

    def has_key(self, provider: str) -> bool:
        """Check if a key exists for a provider."""
        provider = provider.strip().lower()
        keys = self._load_keys()
        return provider in keys

    def rotate_encryption_key(self) -> None:
        """Rotate the machine-bound encryption key.

        Generates a new install-id and re-encrypts all stored keys.
        Both the outer blob AND each individual ``encrypted_value`` are
        re-encrypted with the new derived key.
        """
        # 1. Decrypt everything with the CURRENT key
        keys = self._load_keys()
        enc = self._get_fernet()
        old_multi = enc.get_multi_fernet()

        # Decrypt all individual key values
        decrypted_values: dict[str, str] = {}
        for provider, entry in keys.items():
            try:
                decrypted_values[provider] = old_multi.decrypt(entry.encrypted_value.encode("ascii")).decode("utf-8")
            except InvalidToken:
                logger.warning("Cannot decrypt key for '%s' during rotation -- it will be lost", provider)

        # 2. Rotate the underlying install-id (generates new key material)
        self._encryption.rotate_key(self._keys_file)

        # 3. Re-encrypt all individual key values with the NEW key
        new_fernet = enc.get_fernet()
        for provider, plaintext in decrypted_values.items():
            if provider in keys:
                keys[provider].encrypted_value = new_fernet.encrypt(plaintext.encode("utf-8")).decode("ascii")

        # 4. Save the full blob (also encrypted with new key)
        self._save_keys(keys)
        logger.info("Encryption key rotated -- %d key(s) re-encrypted", len(decrypted_values))

    async def validate_key(self, provider: str) -> KeyValidation:
        """Validate a stored key by making a test API call.

        The validation result is persisted in the key store.

        Args:
            provider: Provider name to validate.

        Returns:
            ``KeyValidation`` with ``is_valid``, provider, error detail.

        Raises:
            ValueError: If no key is stored for the provider.
        """
        provider = provider.strip().lower()
        api_key = self.get_key(provider)
        if api_key is None:
            raise ValueError(f"No key stored for provider '{provider}'")

        result = await ProviderValidator.validate(provider, api_key)

        # Persist validation result
        keys = self._load_keys()
        if provider in keys:
            keys[provider].last_validated = result.validated_at
            keys[provider].is_valid = result.is_valid
            self._save_keys(keys)

        return result
