"""
Key Management Utilities for AI Data Encryption.

This module provides utilities for managing encryption keys stored in AWS SSM
Parameter Store, implementing the key hierarchy from:
    lightwave://schemas/governance/security/ai_data_protection

Key Hierarchy:
    /lightwave/encryption/system              (system master key)
    /lightwave/encryption/tenants/{tenant_id} (per-tenant keys)

Note: User keys are derived from tenant keys using HKDF, not stored separately.

Usage:
    from lightwave.security.keys import KeyManager

    # Get the system key
    key_manager = KeyManager()
    system_key = key_manager.get_system_key()

    # Get or create a tenant key
    tenant_key = key_manager.get_tenant_key("tenant-123")
"""

from __future__ import annotations

import base64
import logging
import os
import secrets
from dataclasses import dataclass
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from mypy_boto3_ssm import SSMClient

logger = logging.getLogger(__name__)


# =============================================================================
# Constants from ai_data_protection.yaml
# =============================================================================

SYSTEM_KEY_PATH = "/lightwave/encryption/system"
TENANT_KEY_PATH_TEMPLATE = "/lightwave/encryption/tenants/{tenant_id}"
DEFAULT_REGION = "us-east-1"

# Key rotation periods (in days)
SYSTEM_KEY_ROTATION_DAYS = 90
TENANT_KEY_ROTATION_DAYS = 30


# =============================================================================
# Exceptions
# =============================================================================


class KeyManagementError(Exception):
    """Base exception for key management errors."""


class KeyNotFoundError(KeyManagementError):
    """Raised when a key is not found in SSM."""


class KeyCreationError(KeyManagementError):
    """Raised when key creation fails."""


# =============================================================================
# Key Generation
# =============================================================================


def generate_fernet_key() -> str:
    """
    Generate a new Fernet-compatible encryption key.

    Returns:
        URL-safe base64-encoded 32-byte key
    """
    # Fernet requires exactly 32 bytes, url-safe base64 encoded
    key_bytes = secrets.token_bytes(32)
    return base64.urlsafe_b64encode(key_bytes).decode()


# =============================================================================
# Key Info
# =============================================================================


@dataclass
class KeyInfo:
    """Information about an encryption key."""

    path: str
    version: int
    last_modified: str | None
    description: str | None


# =============================================================================
# Key Manager
# =============================================================================


class KeyManager:
    """
    Manages encryption keys in AWS SSM Parameter Store.

    Implements the key hierarchy from ai_data_protection.yaml:
    - System key: Master key for system-level encryption
    - Tenant keys: Per-tenant derived keys for tenant isolation

    All keys are stored as SecureString parameters in SSM.

    Usage:
        key_manager = KeyManager()

        # Get system key (from environment or SSM)
        system_key = key_manager.get_system_key()

        # Get tenant key (creates if not exists)
        tenant_key = key_manager.get_tenant_key("tenant-123")

        # Rotate keys
        key_manager.rotate_system_key()
        key_manager.rotate_tenant_key("tenant-123")
    """

    def __init__(
        self,
        region: str = DEFAULT_REGION,
        ssm_client: "SSMClient | None" = None,
    ) -> None:
        """
        Initialize the key manager.

        Args:
            region: AWS region for SSM
            ssm_client: Optional SSM client (for testing)
        """
        self._region = region
        self._ssm_client = ssm_client
        self._cached_system_key: str | None = None

    @property
    def ssm_client(self) -> "SSMClient":
        """Get or create the SSM client."""
        if self._ssm_client is None:
            import boto3

            self._ssm_client = boto3.client("ssm", region_name=self._region)
        return self._ssm_client

    # -------------------------------------------------------------------------
    # System Key Management
    # -------------------------------------------------------------------------

    def get_system_key(self, use_env: bool = True) -> str:
        """
        Get the system encryption key.

        Checks in order:
        1. Cached value
        2. FERNET_KEY environment variable (if use_env=True)
        3. AWS SSM Parameter Store

        Args:
            use_env: Whether to check environment variable first

        Returns:
            The system encryption key

        Raises:
            KeyNotFoundError: If key is not found
        """
        # Check cache
        if self._cached_system_key:
            return self._cached_system_key

        # Check environment variable
        if use_env:
            env_key = os.environ.get("FERNET_KEY")
            if env_key:
                self._cached_system_key = env_key
                return env_key

        # Check SSM
        try:
            response = self.ssm_client.get_parameter(
                Name=SYSTEM_KEY_PATH,
                WithDecryption=True,
            )
            key = response["Parameter"]["Value"]
            self._cached_system_key = key
            return key
        except self.ssm_client.exceptions.ParameterNotFound as e:
            raise KeyNotFoundError(
                f"System key not found at {SYSTEM_KEY_PATH}. "
                "Either set FERNET_KEY environment variable or create the SSM parameter."
            ) from e
        except Exception as e:
            logger.error(f"Failed to retrieve system key from SSM: {e}")
            raise KeyManagementError(f"Failed to retrieve system key: {e}") from e

    def create_system_key(self, overwrite: bool = False) -> str:
        """
        Create a new system encryption key in SSM.

        Args:
            overwrite: Whether to overwrite existing key

        Returns:
            The newly created key

        Raises:
            KeyCreationError: If key creation fails
        """
        key = generate_fernet_key()

        try:
            self.ssm_client.put_parameter(
                Name=SYSTEM_KEY_PATH,
                Description="System-level encryption key for AI data (v1)",
                Value=key,
                Type="SecureString",
                Overwrite=overwrite,
                Tags=[
                    {"Key": "Purpose", "Value": "ai-data-encryption"},
                    {"Key": "RotationDays", "Value": str(SYSTEM_KEY_ROTATION_DAYS)},
                ],
            )
            self._cached_system_key = key
            logger.info(f"Created system key at {SYSTEM_KEY_PATH}")
            return key
        except self.ssm_client.exceptions.ParameterAlreadyExists as e:
            raise KeyCreationError(
                f"System key already exists at {SYSTEM_KEY_PATH}. " "Use overwrite=True to replace it."
            ) from e
        except Exception as e:
            logger.error(f"Failed to create system key: {e}")
            raise KeyCreationError(f"Failed to create system key: {e}") from e

    def rotate_system_key(self) -> str:
        """
        Rotate the system encryption key.

        Creates a new key and stores it in SSM, overwriting the old one.
        Note: After rotation, existing encrypted data must be re-encrypted.

        Returns:
            The new system key
        """
        logger.warning(
            "Rotating system key. All tenant keys will be invalidated. " "Ensure data re-encryption process is running."
        )
        self._cached_system_key = None
        return self.create_system_key(overwrite=True)

    # -------------------------------------------------------------------------
    # Tenant Key Management
    # -------------------------------------------------------------------------

    def _get_tenant_key_path(self, tenant_id: str) -> str:
        """Get the SSM path for a tenant key."""
        return TENANT_KEY_PATH_TEMPLATE.format(tenant_id=tenant_id)

    def get_tenant_key(self, tenant_id: str, create_if_missing: bool = True) -> str:
        """
        Get the encryption key for a specific tenant.

        Args:
            tenant_id: The tenant identifier
            create_if_missing: Whether to create the key if it doesn't exist

        Returns:
            The tenant encryption key

        Raises:
            KeyNotFoundError: If key not found and create_if_missing=False
        """
        path = self._get_tenant_key_path(tenant_id)

        try:
            response = self.ssm_client.get_parameter(
                Name=path,
                WithDecryption=True,
            )
            return response["Parameter"]["Value"]
        except self.ssm_client.exceptions.ParameterNotFound as err:
            if create_if_missing:
                return self.create_tenant_key(tenant_id)
            raise KeyNotFoundError(f"Tenant key not found for {tenant_id}") from err
        except Exception as e:
            logger.error(f"Failed to retrieve tenant key from SSM: {e}")
            raise KeyManagementError(f"Failed to retrieve tenant key: {e}") from e

    def create_tenant_key(self, tenant_id: str, overwrite: bool = False) -> str:
        """
        Create a new encryption key for a tenant.

        Args:
            tenant_id: The tenant identifier
            overwrite: Whether to overwrite existing key

        Returns:
            The newly created key
        """
        key = generate_fernet_key()
        path = self._get_tenant_key_path(tenant_id)

        try:
            self.ssm_client.put_parameter(
                Name=path,
                Description=f"Encryption key for tenant {tenant_id}",
                Value=key,
                Type="SecureString",
                Overwrite=overwrite,
                Tags=[
                    {"Key": "Purpose", "Value": "ai-data-encryption"},
                    {"Key": "TenantId", "Value": tenant_id},
                    {"Key": "RotationDays", "Value": str(TENANT_KEY_ROTATION_DAYS)},
                ],
            )
            logger.info(f"Created tenant key for {tenant_id}")
            return key
        except self.ssm_client.exceptions.ParameterAlreadyExists as e:
            raise KeyCreationError(
                f"Key already exists for tenant {tenant_id}. " "Use overwrite=True to replace it."
            ) from e
        except Exception as e:
            logger.error(f"Failed to create tenant key: {e}")
            raise KeyCreationError(f"Failed to create tenant key: {e}") from e

    def rotate_tenant_key(self, tenant_id: str) -> str:
        """
        Rotate the encryption key for a tenant.

        Note: After rotation, existing encrypted data must be re-encrypted.

        Args:
            tenant_id: The tenant identifier

        Returns:
            The new tenant key
        """
        logger.info(f"Rotating key for tenant {tenant_id}")
        return self.create_tenant_key(tenant_id, overwrite=True)

    def delete_tenant_key(self, tenant_id: str) -> None:
        """
        Delete the encryption key for a tenant.

        WARNING: This will make all encrypted data for this tenant unrecoverable.

        Args:
            tenant_id: The tenant identifier
        """
        path = self._get_tenant_key_path(tenant_id)

        try:
            self.ssm_client.delete_parameter(Name=path)
            logger.warning(f"Deleted encryption key for tenant {tenant_id}")
        except self.ssm_client.exceptions.ParameterNotFound:
            logger.warning(f"No key found to delete for tenant {tenant_id}")
        except Exception as e:
            logger.error(f"Failed to delete tenant key: {e}")
            raise KeyManagementError(f"Failed to delete tenant key: {e}") from e

    # -------------------------------------------------------------------------
    # Key Listing & Info
    # -------------------------------------------------------------------------

    def list_tenant_keys(self) -> list[str]:
        """
        List all tenant IDs that have encryption keys.

        Returns:
            List of tenant IDs
        """
        prefix = TENANT_KEY_PATH_TEMPLATE.replace("{tenant_id}", "")
        tenant_ids = []

        try:
            paginator = self.ssm_client.get_paginator("describe_parameters")
            for page in paginator.paginate(
                ParameterFilters=[{"Key": "Path", "Option": "Recursive", "Values": [prefix]}]
            ):
                for param in page.get("Parameters", []):
                    # Extract tenant_id from path
                    path = param["Name"]
                    tenant_id = path.replace(prefix, "")
                    if tenant_id:
                        tenant_ids.append(tenant_id)
        except Exception as e:
            logger.error(f"Failed to list tenant keys: {e}")
            raise KeyManagementError(f"Failed to list tenant keys: {e}") from e

        return tenant_ids

    def get_key_info(self, path: str) -> KeyInfo | None:
        """
        Get metadata about a key without retrieving its value.

        Args:
            path: The SSM parameter path

        Returns:
            KeyInfo or None if not found
        """
        try:
            response = self.ssm_client.describe_parameters(ParameterFilters=[{"Key": "Name", "Values": [path]}])
            params = response.get("Parameters", [])
            if not params:
                return None

            param = params[0]
            return KeyInfo(
                path=param["Name"],
                version=param.get("Version", 1),
                last_modified=str(param.get("LastModifiedDate", "")),
                description=param.get("Description"),
            )
        except Exception as e:
            logger.error(f"Failed to get key info: {e}")
            return None

    # -------------------------------------------------------------------------
    # Cache Management
    # -------------------------------------------------------------------------

    def clear_cache(self) -> None:
        """Clear cached keys (useful after rotation)."""
        self._cached_system_key = None
        logger.info("Key manager cache cleared")


# =============================================================================
# Module-level convenience functions
# =============================================================================

_default_manager: KeyManager | None = None


def get_key_manager(region: str = DEFAULT_REGION) -> KeyManager:
    """Get the default key manager instance."""
    global _default_manager
    if _default_manager is None:
        _default_manager = KeyManager(region=region)
    return _default_manager
