"""Enhanced API key encryption with PBKDF2 key derivation and rotation.

Derives a Fernet key from a machine-specific seed (hostname + MAC address +
install-id UUID).  Keys are bound to the machine they were created on.

Supports key rotation via :meth:`rotate_key`, which generates a new install-id
and re-encrypts all stored keys using :class:`~cryptography.fernet.MultiFernet`
for a seamless transition window.
"""

from __future__ import annotations

import base64
import contextlib
import logging
import os
import platform
import stat
import uuid
from typing import TYPE_CHECKING

from cryptography.fernet import Fernet, InvalidToken, MultiFernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

if TYPE_CHECKING:
    from pathlib import Path

logger = logging.getLogger(__name__)

_PBKDF2_ITERATIONS = 480_000
_PBKDF2_SALT = b"llmhost-key-derivation-v1"  # static salt; entropy comes from seed
_INSTALL_ID_FILENAME = ".install-id"
_OLD_INSTALL_ID_FILENAME = ".install-id.old"


class KeyEncryption:
    """Enhanced key encryption with key derivation and rotation.

    Uses PBKDF2 to derive a Fernet key from a machine-specific seed
    (hostname + MAC + install-id).  This means keys are bound to the machine.
    """

    def __init__(self, data_dir: Path) -> None:
        self._data_dir = data_dir
        self._install_id_path = data_dir / _INSTALL_ID_FILENAME
        self._old_install_id_path = data_dir / _OLD_INSTALL_ID_FILENAME

    # ------------------------------------------------------------------
    # Machine seed
    # ------------------------------------------------------------------

    def _get_machine_seed(self) -> bytes:
        """Generate machine-specific seed from hostname + MAC + install-id.

        The install-id is a random UUID generated on first use and stored in
        ``~/.llmhosts/.install-id``.  Combined with hostname and MAC address,
        this produces a seed that is unique to this machine installation.

        Returns:
            Raw bytes suitable for PBKDF2 key derivation.
        """
        hostname = platform.node() or "unknown-host"
        mac = self._get_mac_address()
        install_id = self._get_or_create_install_id()

        seed = f"{hostname}:{mac}:{install_id}".encode()
        logger.debug("Machine seed components: hostname=%s, mac=%s, install_id=%s", hostname, mac, install_id[:8])
        return seed

    @staticmethod
    def _get_mac_address() -> str:
        """Get the primary MAC address as a hex string.

        Falls back to a deterministic placeholder if MAC detection fails.
        """
        try:
            mac_int = uuid.getnode()
            # uuid.getnode() returns a random value with bit 0 of first octet set
            # when it can't find a real MAC.  Check for this.
            if (mac_int >> 40) & 1:
                # Likely random -- fall back to hostname-based value
                return f"no-mac-{platform.node()}"
            return f"{mac_int:012x}"
        except Exception:
            return f"no-mac-{platform.node()}"

    def _get_or_create_install_id(self) -> str:
        """Load or generate the installation UUID.

        The install-id is stored at ``~/.llmhosts/.install-id`` with 600
        permissions.  A new random UUID is generated on first use.

        Returns:
            The install-id string (UUID hex).
        """
        self._ensure_dirs()

        if self._install_id_path.exists():
            try:
                raw = self._install_id_path.read_text(encoding="utf-8").strip()
                if raw:
                    return raw
            except (OSError, UnicodeDecodeError):
                logger.warning("Corrupt install-id file -- regenerating")

        install_id = uuid.uuid4().hex
        self._install_id_path.write_text(install_id, encoding="utf-8")
        os.chmod(self._install_id_path, stat.S_IRUSR | stat.S_IWUSR)  # 0o600
        logger.info("Generated new install-id: %s", self._install_id_path)
        return install_id

    # ------------------------------------------------------------------
    # Key derivation
    # ------------------------------------------------------------------

    def _derive_key(self, seed: bytes) -> bytes:
        """Derive a Fernet key from seed using PBKDF2HMAC (SHA256, 480000 iterations).

        Args:
            seed: The raw machine-specific seed bytes.

        Returns:
            A 32-byte URL-safe base64-encoded key suitable for Fernet.
        """
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=_PBKDF2_SALT,
            iterations=_PBKDF2_ITERATIONS,
        )
        derived = kdf.derive(seed)
        return base64.urlsafe_b64encode(derived)

    def get_fernet(self) -> Fernet:
        """Get a Fernet instance for encryption/decryption.

        Derives the key from the current machine seed on every call (cheap
        after first PBKDF2 derivation is cached by the OS).

        Returns:
            A :class:`~cryptography.fernet.Fernet` instance.
        """
        seed = self._get_machine_seed()
        key = self._derive_key(seed)
        return Fernet(key)

    def get_multi_fernet(self) -> MultiFernet:
        """Get a MultiFernet that can decrypt with both current and old keys.

        During key rotation, the old install-id is preserved temporarily.
        This MultiFernet can decrypt data encrypted with either key, and
        always encrypts with the current key.

        Returns:
            A :class:`~cryptography.fernet.MultiFernet` instance.
        """
        fernets: list[Fernet] = [self.get_fernet()]

        # If there's an old install-id, also add a Fernet for it
        if self._old_install_id_path.exists():
            try:
                old_id = self._old_install_id_path.read_text(encoding="utf-8").strip()
                if old_id:
                    hostname = platform.node() or "unknown-host"
                    mac = self._get_mac_address()
                    old_seed = f"{hostname}:{mac}:{old_id}".encode()
                    old_key = self._derive_key(old_seed)
                    fernets.append(Fernet(old_key))
                    logger.debug("MultiFernet includes old rotation key")
            except (OSError, UnicodeDecodeError):
                logger.warning("Could not read old install-id for MultiFernet")

        return MultiFernet(fernets)

    # ------------------------------------------------------------------
    # Key rotation
    # ------------------------------------------------------------------

    def rotate_key(self, keys_file: Path) -> None:
        """Generate new install-id and re-encrypt all keys.

        Uses :class:`~cryptography.fernet.MultiFernet` for a seamless
        transition.  The old install-id is preserved temporarily so that
        in-flight decryption calls still work.

        Steps:
            1. Save old install-id as ``.install-id.old``
            2. Generate fresh install-id
            3. Build MultiFernet (new key first, old key second)
            4. Re-encrypt the keys file with the new key
            5. Remove the old install-id file

        Args:
            keys_file: Path to the encrypted keys file (``keys.enc``).

        Raises:
            FileNotFoundError: If ``keys_file`` does not exist.
            cryptography.fernet.InvalidToken: If decryption fails.
        """
        if not keys_file.exists():
            logger.info("No keys file to rotate -- generating new install-id only")
            self._regenerate_install_id()
            return

        # 1. Read current encrypted data
        encrypted_blob = keys_file.read_bytes()
        if not encrypted_blob:
            self._regenerate_install_id()
            return

        # 2. Decrypt with current key
        old_fernet = self.get_fernet()
        try:
            plaintext = old_fernet.decrypt(encrypted_blob)
        except InvalidToken:
            logger.error("Cannot decrypt keys file during rotation -- aborting")
            raise

        # 3. Preserve old install-id
        if self._install_id_path.exists():
            old_id = self._install_id_path.read_text(encoding="utf-8").strip()
            self._old_install_id_path.write_text(old_id, encoding="utf-8")
            os.chmod(self._old_install_id_path, stat.S_IRUSR | stat.S_IWUSR)

        # 4. Generate new install-id
        self._regenerate_install_id()

        # 5. Re-encrypt with new key
        new_fernet = self.get_fernet()
        re_encrypted = new_fernet.encrypt(plaintext)
        keys_file.write_bytes(re_encrypted)
        os.chmod(keys_file, stat.S_IRUSR | stat.S_IWUSR)

        # 6. Clean up old install-id
        self.secure_delete(self._old_install_id_path)
        logger.info("Key rotation complete -- all keys re-encrypted with new derived key")

    def _regenerate_install_id(self) -> None:
        """Generate a fresh install-id UUID, overwriting the existing one."""
        self._ensure_dirs()
        new_id = uuid.uuid4().hex
        self._install_id_path.write_text(new_id, encoding="utf-8")
        os.chmod(self._install_id_path, stat.S_IRUSR | stat.S_IWUSR)
        logger.info("Generated new install-id for key rotation")

    # ------------------------------------------------------------------
    # Secure deletion
    # ------------------------------------------------------------------

    @staticmethod
    def secure_delete(path: Path) -> None:
        """Overwrite file with random bytes before deleting.

        Writes random data matching the file size, flushes to disk, then
        removes the file.  This is a best-effort mitigation -- SSDs with
        wear-levelling may retain old data regardless.

        Args:
            path: File to securely delete.
        """
        if not path.exists():
            return

        try:
            file_size = path.stat().st_size
            if file_size > 0:
                # Overwrite with random bytes (3 passes)
                for _ in range(3):
                    with open(path, "wb") as f:
                        f.write(os.urandom(file_size))
                        f.flush()
                        os.fsync(f.fileno())
            path.unlink()
            logger.debug("Securely deleted: %s", path)
        except OSError as exc:
            logger.warning("Secure delete failed for %s: %s -- attempting normal delete", path, exc)
            with contextlib.suppress(OSError):
                path.unlink()

    # ------------------------------------------------------------------
    # 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
