"""Job payload encryption module for Oclawma.

This module provides AES-256-GCM encryption for sensitive job payloads
stored in the SQLite queue. It supports:
- Key management (environment variable, file, or generated)
- Key rotation (versioned keys)
- Opt-in encryption per job type
- Minimal performance overhead
"""

from __future__ import annotations

import base64
import json
import os
import secrets
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Any

try:
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM

    HAS_CRYPTOGRAPHY = True
except ImportError:
    HAS_CRYPTOGRAPHY = False


class EncryptionError(Exception):
    """Base exception for encryption errors."""

    pass


class DecryptionError(EncryptionError):
    """Raised when decryption fails."""

    pass


class KeyNotFoundError(EncryptionError):
    """Raised when a key version is not found."""

    pass


def generate_key() -> bytes:
    """Generate a new 256-bit (32-byte) encryption key.

    Returns:
        Random 32-byte key suitable for AES-256-GCM.
    """
    return secrets.token_bytes(32)


def _encode_key(key: bytes) -> str:
    """Encode key to base64 string."""
    return base64.b64encode(key).decode("ascii")


def _decode_key(key_b64: str) -> bytes:
    """Decode key from base64 string."""
    return base64.b64decode(key_b64)


@dataclass
class EncryptionMetadata:
    """Metadata stored with encrypted payloads.

    Attributes:
        encrypted: Whether the payload is encrypted
        key_version: Version identifier for the encryption key used
        algorithm: Encryption algorithm used
    """

    encrypted: bool = True
    key_version: str = "1"
    algorithm: str = "AES-256-GCM"

    def to_dict(self) -> dict[str, Any]:
        """Convert metadata to dictionary."""
        return {
            "encrypted": self.encrypted,
            "key_version": self.key_version,
            "algorithm": self.algorithm,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> EncryptionMetadata:
        """Create metadata from dictionary."""
        return cls(
            encrypted=data.get("encrypted", True),
            key_version=data.get("key_version", "1"),
            algorithm=data.get("algorithm", "AES-256-GCM"),
        )


class EncryptionManager:
    """Manages encryption/decryption of job payloads.

    Uses AES-256-GCM for authenticated encryption. Supports key rotation
    through versioned keys and opt-in encryption per job type.

    Key sources (in order of priority):
    1. Explicit key parameter
    2. OCLAWMA_ENCRYPTION_KEY environment variable
    3. Key file (default: ~/.oclawma/encryption.key)
    4. Auto-generated key

    Example:
        >>> manager = EncryptionManager()
        >>> ciphertext = manager.encrypt(b"sensitive data")
        >>> plaintext = manager.decrypt(ciphertext)

        >>> # With key rotation
        >>> manager.add_key("v2", generate_key())
        >>> manager.current_key_version = "v2"
        >>> # Old data can still be decrypted
        >>> plaintext = manager.decrypt(old_ciphertext)
    """

    DEFAULT_KEY_FILE = Path.home() / ".oclawma" / "encryption.key"
    KEYS_FILE = Path.home() / ".oclawma" / "encryption_keys.json"

    def __init__(
        self,
        key: bytes | None = None,
        key_file: str | Path | None = None,
        encrypted_types: set[str] | None = None,
    ) -> None:
        """Initialize the encryption manager.

        Args:
            key: Explicit 32-byte encryption key (highest priority)
            key_file: Path to key file or keys JSON file
            encrypted_types: Set of job type names that should be encrypted

        Raises:
            ValueError: If key is provided but not 32 bytes
            ImportError: If cryptography library is not installed
        """
        if not HAS_CRYPTOGRAPHY:
            raise ImportError(
                "cryptography library is required for encryption. "
                "Install with: pip install cryptography"
            )

        self._lock = threading.RLock()
        self.keys: dict[str, bytes] = {}
        self.current_key_version: str = "1"
        self.encrypted_types: set[str] = encrypted_types or set()

        # Determine key source
        if key is not None:
            # Convert string key to bytes if needed
            if isinstance(key, str):
                key = key.encode("utf-8")
            self._validate_key(key)
            self.keys["1"] = key
        elif os.environ.get("OCLAWMA_ENCRYPTION_KEY"):
            key_b64 = os.environ["OCLAWMA_ENCRYPTION_KEY"]
            key = _decode_key(key_b64)
            self._validate_key(key)
            self.keys["1"] = key
        elif key_file:
            self._load_key_from_file(key_file)
        else:
            # Try default key file, generate if not exists
            if self.DEFAULT_KEY_FILE.exists():
                self._load_key_from_file(self.DEFAULT_KEY_FILE)
            elif self.KEYS_FILE.exists():
                self._load_keys_from_json(self.KEYS_FILE)
            else:
                # Generate new key and save
                key = generate_key()
                self.keys["1"] = key
                self._save_default_key()

    def _validate_key(self, key: bytes) -> None:
        """Validate that key is correct size for AES-256."""
        if len(key) != 32:
            raise ValueError(f"Key must be 32 bytes for AES-256, got {len(key)}")

    def _load_key_from_file(self, key_file: str | Path) -> None:
        """Load key from file (base64 encoded single key or JSON with multiple keys)."""
        path = Path(key_file)
        if path.exists():
            content = path.read_text().strip()
            # Check if it's a JSON file with multiple keys
            if content.startswith("{"):
                self._load_keys_from_json(path)
            else:
                # Single base64-encoded key
                key = _decode_key(content)
                self._validate_key(key)
                self.keys["1"] = key
        else:
            # Generate and save new key
            key = generate_key()
            self.keys["1"] = key
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_text(_encode_key(key))

    def _load_keys_from_json(self, keys_file: str | Path) -> None:
        """Load multiple versioned keys from JSON file."""
        path = Path(keys_file)
        data = json.loads(path.read_text())
        for version, key_b64 in data.items():
            key = _decode_key(key_b64)
            self._validate_key(key)
            self.keys[version] = key

    def _save_default_key(self) -> None:
        """Save the default key to the default location."""
        self.DEFAULT_KEY_FILE.parent.mkdir(parents=True, exist_ok=True)
        self.DEFAULT_KEY_FILE.write_text(_encode_key(self.keys["1"]))

    def add_key(self, version: str, key: bytes) -> None:
        """Add a new key version for rotation.

        Args:
            version: Key version identifier (e.g., "v2", "2024-01")
            key: 32-byte encryption key

        Raises:
            ValueError: If version already exists or key is invalid
        """
        with self._lock:
            if version in self.keys:
                raise ValueError(f"Key version '{version}' already exists")
            self._validate_key(key)
            self.keys[version] = key

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

        Args:
            version: Key version to remove

        Returns:
            True if removed, False if not found

        Raises:
            ValueError: If trying to remove the last key
        """
        with self._lock:
            if version not in self.keys:
                return False
            if len(self.keys) <= 1:
                raise ValueError("Cannot remove the last key")
            del self.keys[version]
            return True

    def save_keys(self, path: str | Path | None = None) -> None:
        """Save all keys to JSON file.

        Args:
            path: File path (default: ~/.oclawma/encryption_keys.json)
        """
        save_path = Path(path) if path else self.KEYS_FILE
        save_path.parent.mkdir(parents=True, exist_ok=True)
        data = {version: _encode_key(key) for version, key in self.keys.items()}
        save_path.write_text(json.dumps(data, indent=2))

    def should_encrypt(self, job_type: str) -> bool:
        """Check if a job type should be encrypted.

        Args:
            job_type: Type/name of the job

        Returns:
            True if jobs of this type should be encrypted
        """
        return job_type in self.encrypted_types

    def add_encrypted_type(self, job_type: str) -> None:
        """Add a job type to the encryption list.

        Args:
            job_type: Type/name of the job to encrypt
        """
        self.encrypted_types.add(job_type)

    def remove_encrypted_type(self, job_type: str) -> None:
        """Remove a job type from the encryption list.

        Args:
            job_type: Type/name of the job
        """
        self.encrypted_types.discard(job_type)

    def encrypt(
        self,
        plaintext: bytes,
        key_version: str | None = None,
    ) -> bytes:
        """Encrypt plaintext using AES-256-GCM.

        The returned ciphertext includes metadata for decryption:
        - key_version: Which key was used
        - nonce: AES-GCM nonce (12 bytes)
        - ciphertext: Encrypted data
        - tag: Authentication tag (16 bytes)

        Args:
            plaintext: Data to encrypt
            key_version: Key version to use (default: current_key_version)

        Returns:
            JSON-encoded ciphertext with metadata

        Raises:
            KeyNotFoundError: If specified key version doesn't exist
        """
        with self._lock:
            version = key_version or self.current_key_version
            if version not in self.keys:
                raise KeyNotFoundError(f"Key version '{version}' not found")

            key = self.keys[version]
            nonce = secrets.token_bytes(12)  # 96-bit nonce for GCM

            aesgcm = AESGCM(key)
            ciphertext = aesgcm.encrypt(nonce, plaintext, None)

            # ciphertext includes tag (last 16 bytes)
            ct_data = ciphertext[:-16]
            tag = ciphertext[-16:]

            # Package as JSON for storage
            data = {
                "key_version": version,
                "nonce": base64.b64encode(nonce).decode("ascii"),
                "ciphertext": base64.b64encode(ct_data).decode("ascii"),
                "tag": base64.b64encode(tag).decode("ascii"),
            }
            return json.dumps(data).encode("utf-8")

    def decrypt(self, ciphertext: bytes) -> bytes:
        """Decrypt ciphertext using AES-256-GCM.

        Args:
            ciphertext: JSON-encoded ciphertext from encrypt()

        Returns:
            Decrypted plaintext

        Raises:
            DecryptionError: If decryption fails (invalid data, tampering)
            KeyNotFoundError: If the key version is not available
        """
        with self._lock:
            try:
                data = json.loads(ciphertext.decode("utf-8"))
            except (json.JSONDecodeError, UnicodeDecodeError) as e:
                raise DecryptionError(f"Invalid ciphertext format: {e}") from e

            version = data.get("key_version", "1")
            if version not in self.keys:
                raise KeyNotFoundError(f"Key version '{version}' not found")

            try:
                nonce = base64.b64decode(data["nonce"])
                ct_data = base64.b64decode(data["ciphertext"])
                tag = base64.b64decode(data["tag"])

                # Reconstruct ciphertext with tag
                full_ciphertext = ct_data + tag

                key = self.keys[version]
                aesgcm = AESGCM(key)
                return aesgcm.decrypt(nonce, full_ciphertext, None)
            except Exception as e:
                raise DecryptionError(f"Decryption failed: {e}") from e

    def encrypt_string(self, plaintext: str, key_version: str | None = None) -> str:
        """Encrypt a string and return base64-encoded result.

        Args:
            plaintext: String to encrypt
            key_version: Key version to use

        Returns:
            Base64-encoded ciphertext
        """
        ciphertext = self.encrypt(plaintext.encode("utf-8"), key_version)
        return base64.b64encode(ciphertext).decode("ascii")

    def decrypt_string(self, ciphertext: str) -> str:
        """Decrypt a base64-encoded ciphertext to string.

        Args:
            ciphertext: Base64-encoded ciphertext

        Returns:
            Decrypted string
        """
        ct_bytes = base64.b64decode(ciphertext)
        return self.decrypt(ct_bytes).decode("utf-8")

    def encrypt_dict(self, plaintext: dict[str, Any], key_version: str | None = None) -> str:
        """Encrypt a dictionary (JSON-serialized).

        Args:
            plaintext: Dictionary to encrypt
            key_version: Key version to use

        Returns:
            Base64-encoded ciphertext
        """
        json_data = json.dumps(plaintext)
        return self.encrypt_string(json_data, key_version)

    def decrypt_dict(self, ciphertext: str) -> dict[str, Any]:
        """Decrypt ciphertext to dictionary.

        Args:
            ciphertext: Base64-encoded ciphertext

        Returns:
            Decrypted dictionary
        """
        json_data = self.decrypt_string(ciphertext)
        return json.loads(json_data)

    def __enter__(self) -> EncryptionManager:
        """Context manager entry."""
        return self

    def __exit__(self, *args: Any) -> None:
        """Context manager exit."""
        pass
