"""Module for working with real KMS/HSM systems."""

import base64
import logging
import os
from abc import ABC, abstractmethod
from typing import Any, Dict

import boto3
import requests
from botocore.exceptions import ClientError

logger = logging.getLogger(__name__)


class KMSProvider(ABC):
    """Abstract class for KMS providers."""

    @abstractmethod
    def encrypt(self, plaintext: bytes, key_id: str) -> bytes:
        """
        Encrypt data using KMS.

        Args:
            plaintext: Data to encrypt.
            key_id: Key identifier.

        Returns:
            Encrypted data.
        """
        pass

    @abstractmethod
    def decrypt(self, ciphertext: bytes, key_id: str) -> bytes:
        """
        Decrypt data using KMS.

        Args:
            ciphertext: Encrypted data.
            key_id: Key identifier.

        Returns:
            Decrypted data.
        """
        pass

    @abstractmethod
    def generate_data_key(
        self, key_id: str, key_spec: str = "AES_256"
    ) -> Dict[str, bytes]:
        """
        Generate data encryption key.

        Args:
            key_id: Key identifier.
            key_spec: Key specification.

        Returns:
            Dictionary with 'plaintext' and 'ciphertext' keys.
        """
        pass

class AWSKMSProvider(KMSProvider):
    """AWS KMS provider."""

    def __init__(self, region_name: str = "us-east-1") -> None:
        """
        Initialize AWS KMS provider.

        Args:
            region_name: AWS region name.
        """
        self.kms_client = boto3.client("kms", region_name=region_name)
        self.region_name = region_name

    def encrypt(self, plaintext: bytes, key_id: str) -> bytes:
        """Encrypt data using AWS KMS."""
        try:
            response = self.kms_client.encrypt(KeyId=key_id, Plaintext=plaintext)
            return response["CiphertextBlob"]
        except ClientError as e:
            logger.error("AWS KMS encryption failed: %s", e)
            raise

    def decrypt(self, ciphertext: bytes, key_id: str) -> bytes:
        """Decrypt data using AWS KMS."""
        try:
            response = self.kms_client.decrypt(
                CiphertextBlob=ciphertext, KeyId=key_id
            )
            return response["Plaintext"]
        except ClientError as e:
            logger.error("AWS KMS decryption failed: %s", e)
            raise

    def generate_data_key(
        self, key_id: str, key_spec: str = "AES_256"
    ) -> Dict[str, bytes]:
        """Generate data encryption key."""
        try:
            response = self.kms_client.generate_data_key(
                KeyId=key_id, KeySpec=key_spec
            )
            return {
                "plaintext": response["Plaintext"],
                "ciphertext": response["CiphertextBlob"],
            }
        except ClientError as e:
            logger.error("AWS KMS data key generation failed: %s", e)
            raise

class GoogleCloudKMSProvider(KMSProvider):
    """Google Cloud KMS provider."""

    def __init__(self, project_id: str, location: str = "global") -> None:
        """
        Initialize Google Cloud KMS provider.

        Args:
            project_id: Google Cloud project ID.
            location: Key location.
        """
        self.project_id = project_id
        self.location = location
        self.base_url = (
            f"https://cloudkms.googleapis.com/v1/projects/{project_id}"
            f"/locations/{location}"
        )

    def encrypt(self, plaintext: bytes, key_id: str) -> bytes:
        """Encrypt data using Google Cloud KMS."""
        try:
            url = (
                f"{self.base_url}/keyRings/default/cryptoKeys/{key_id}:encrypt"
            )

            response = requests.post(
                url,
                json={"plaintext": base64.b64encode(plaintext).decode()},
                headers={
                    "Authorization": f"Bearer {self._get_access_token()}"
                },
                timeout=30,
            )
            response.raise_for_status()

            return base64.b64decode(response.json()["ciphertext"])

        except requests.RequestException as e:
            logger.error("Google Cloud KMS encryption failed: %s", e)
            raise RuntimeError(
                f"Google Cloud KMS encryption failed: {e}"
            ) from e

    def decrypt(self, ciphertext: bytes, key_id: str) -> bytes:
        """Decrypt data using Google Cloud KMS."""
        try:
            url = (
                f"{self.base_url}/keyRings/default/cryptoKeys/{key_id}:decrypt"
            )

            response = requests.post(
                url,
                json={"ciphertext": base64.b64encode(ciphertext).decode()},
                headers={
                    "Authorization": f"Bearer {self._get_access_token()}"
                },
                timeout=30,
            )
            response.raise_for_status()

            return base64.b64decode(response.json()["plaintext"])

        except requests.RequestException as e:
            logger.error("Google Cloud KMS decryption failed: %s", e)
            raise RuntimeError(
                f"Google Cloud KMS decryption failed: {e}"
            ) from e

    def generate_data_key(
        self, key_id: str, key_spec: str = "AES_256"
    ) -> Dict[str, bytes]:
        """Generate data encryption key."""
        try:
            url = (
                f"{self.base_url}/keyRings/default/cryptoKeys/{key_id}"
                ":generateDataKey"
            )

            response = requests.post(
                url,
                json={"keySpec": key_spec},
                headers={
                    "Authorization": f"Bearer {self._get_access_token()}"
                },
                timeout=30,
            )
            response.raise_for_status()

            data = response.json()
            return {
                "plaintext": base64.b64decode(data["plaintext"]),
                "ciphertext": base64.b64decode(data["ciphertext"]),
            }

        except requests.RequestException as e:
            logger.error("Google Cloud KMS data key generation failed: %s", e)
            raise RuntimeError(
                f"Google Cloud KMS data key generation failed: {e}"
            ) from e

    def _get_access_token(self) -> str:
        """
        Get access token for Google Cloud API.

        Note: In production, use service account or metadata server.
        """
        return os.getenv("GOOGLE_CLOUD_ACCESS_TOKEN", "")

class YubiHSMProvider(KMSProvider):
    """YubiHSM provider (hardware security module)."""

    def __init__(self, hsm_url: str, auth_key_id: int) -> None:
        """
        Initialize YubiHSM provider.

        Args:
            hsm_url: YubiHSM URL.
            auth_key_id: Authentication key ID.
        """
        self.hsm_url = hsm_url
        self.auth_key_id = auth_key_id

    def encrypt(self, plaintext: bytes, key_id: str) -> bytes:
        """Encrypt data using YubiHSM."""
        try:
            response = requests.post(
                f"{self.hsm_url}/api/v1/encrypt",
                json={"key_id": key_id, "plaintext": plaintext.hex()},
                headers={"Authorization": f"Bearer {self._get_hsm_token()}"},
                timeout=30,
            )
            response.raise_for_status()

            return bytes.fromhex(response.json()["ciphertext"])

        except requests.RequestException as e:
            logger.error("YubiHSM encryption failed: %s", e)
            raise RuntimeError(f"YubiHSM encryption failed: {e}") from e

    def decrypt(self, ciphertext: bytes, key_id: str) -> bytes:
        """Decrypt data using YubiHSM."""
        try:
            response = requests.post(
                f"{self.hsm_url}/api/v1/decrypt",
                json={"key_id": key_id, "ciphertext": ciphertext.hex()},
                headers={"Authorization": f"Bearer {self._get_hsm_token()}"},
                timeout=30,
            )
            response.raise_for_status()

            return bytes.fromhex(response.json()["plaintext"])

        except requests.RequestException as e:
            logger.error("YubiHSM decryption failed: %s", e)
            raise RuntimeError(f"YubiHSM decryption failed: {e}") from e

    def generate_data_key(
        self, key_id: str, key_spec: str = "AES_256"
    ) -> Dict[str, bytes]:
        """Generate data encryption key."""
        try:
            response = requests.post(
                f"{self.hsm_url}/api/v1/generate-data-key",
                json={"key_id": key_id, "key_spec": key_spec},
                headers={"Authorization": f"Bearer {self._get_hsm_token()}"},
                timeout=30,
            )
            response.raise_for_status()

            data = response.json()
            return {
                "plaintext": bytes.fromhex(data["plaintext"]),
                "ciphertext": bytes.fromhex(data["ciphertext"]),
            }

        except requests.RequestException as e:
            logger.error("YubiHSM data key generation failed: %s", e)
            raise RuntimeError(
                f"YubiHSM data key generation failed: {e}"
            ) from e

    def _get_hsm_token(self) -> str:
        """
        Get token for YubiHSM.

        Note: In production, use proper YubiHSM authentication.
        """
        return os.getenv("YUBIHSM_TOKEN", "")

class SecureKeyManager:
    """Key manager using real KMS/HSM systems."""

    def __init__(self, kms_provider: KMSProvider, master_key_id: str) -> None:
        """
        Initialize secure key manager.

        Args:
            kms_provider: KMS provider instance.
            master_key_id: Master key identifier.
        """
        self.kms_provider = kms_provider
        self.master_key_id = master_key_id
        self.key_cache: Dict[str, bytes] = {}

    def encrypt_session_key(self, session_key: bytes, session_id: str) -> bytes:
        """
        Encrypt session key using KMS/HSM.

        Args:
            session_key: Session key to encrypt.
            session_id: Session ID for context.

        Returns:
            Encrypted session key.
        """
        try:
            encrypted_key = self.kms_provider.encrypt(
                session_key, self.master_key_id
            )

            logger.info(
                "Session key encrypted with KMS/HSM",
                extra={
                    "session_id": session_id,
                    "key_length": len(session_key),
                    "encrypted_length": len(encrypted_key),
                },
            )

            return encrypted_key

        except Exception as e:
            logger.error(
                "Failed to encrypt session key with KMS/HSM",
                extra={"session_id": session_id, "error": str(e)},
            )
            raise

    def decrypt_session_key(
        self, encrypted_session_key: bytes, session_id: str
    ) -> bytes:
        """
        Decrypt session key using KMS/HSM.

        Args:
            encrypted_session_key: Encrypted session key.
            session_id: Session ID for context.

        Returns:
            Decrypted session key.
        """
        try:
            decrypted_key = self.kms_provider.decrypt(
                encrypted_session_key, self.master_key_id
            )

            logger.info(
                "Session key decrypted with KMS/HSM",
                extra={"session_id": session_id, "key_length": len(decrypted_key)},
            )

            return decrypted_key

        except Exception as e:
            logger.error(
                "Failed to decrypt session key with KMS/HSM",
                extra={"session_id": session_id, "error": str(e)},
            )
            raise

    def generate_encryption_key(self, session_id: str) -> Dict[str, bytes]:
        """
        Generate encryption key using KMS/HSM.

        Args:
            session_id: Session ID.

        Returns:
            Dictionary with 'plaintext' and 'ciphertext' keys.
        """
        try:
            key_data = self.kms_provider.generate_data_key(self.master_key_id)

            logger.info(
                "Encryption key generated with KMS/HSM",
                extra={
                    "session_id": session_id,
                    "key_length": len(key_data["plaintext"]),
                },
            )

            return key_data

        except Exception as e:
            logger.error(
                "Failed to generate encryption key with KMS/HSM",
                extra={"session_id": session_id, "error": str(e)},
            )
            raise


class KMSProviderFactory:
    """Factory for creating KMS providers."""

    @staticmethod
    def create_provider(provider_type: str, **kwargs: Any) -> KMSProvider:
        """
        Create KMS provider by type.

        Args:
            provider_type: Provider type ('aws', 'gcp', 'yubihsm').
            **kwargs: Provider-specific arguments.

        Returns:
            KMS provider instance.

        Raises:
            ValueError: If provider type is unsupported.
        """
        if provider_type == "aws":
            return AWSKMSProvider(**kwargs)
        if provider_type == "gcp":
            return GoogleCloudKMSProvider(**kwargs)
        if provider_type == "yubihsm":
            return YubiHSMProvider(**kwargs)

        raise ValueError(f"Unsupported KMS provider: {provider_type}")


def get_secure_key_manager() -> SecureKeyManager:
    """
    Get configured secure key manager.

    Returns:
        Configured SecureKeyManager instance.

    Raises:
        ValueError: If provider type is unsupported.
    """
    provider_type = os.getenv("KMS_PROVIDER", "aws")
    master_key_id = os.getenv("KMS_MASTER_KEY_ID", "alias/session-keys")

    if provider_type == "aws":
        provider = AWSKMSProvider(
            region_name=os.getenv("AWS_REGION", "us-east-1")
        )
    elif provider_type == "gcp":
        provider = GoogleCloudKMSProvider(
            project_id=os.getenv("GCP_PROJECT_ID", ""),
            location=os.getenv("GCP_LOCATION", "global"),
        )
    elif provider_type == "yubihsm":
        provider = YubiHSMProvider(
            hsm_url=os.getenv("YUBIHSM_URL", ""),
            auth_key_id=int(os.getenv("YUBIHSM_AUTH_KEY_ID", "0")),
        )
    else:
        raise ValueError(f"Unsupported KMS provider: {provider_type}")

    return SecureKeyManager(provider, master_key_id)


secure_key_manager = get_secure_key_manager()
