"""
KMS Envelope Encryption Service for BYO Provider Locker.

Implements envelope encryption pattern:
1. Generate random 256-bit DEK (data encryption key)
2. Encrypt secret with DEK using AES-256-GCM
3. Encrypt DEK with Cloud KMS (store encrypted_dek + kms_key_version)
4. To decrypt: KMS decrypt DEK -> AES-GCM decrypt secret
5. Rotation: rewrap DEKs with new KMS version without touching ciphertext
"""

from __future__ import annotations

import os
import secrets
from dataclasses import dataclass
from typing import Optional

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
import logging

logger = logging.getLogger(__name__)

# Try to import Google Cloud KMS - graceful fallback for local dev
try:
    from google.cloud import kms
    HAS_KMS = True
except ImportError:
    HAS_KMS = False
    kms = None

# Constants
DEK_SIZE_BYTES = 32  # 256-bit key for AES-256
NONCE_SIZE_BYTES = 12  # 96-bit nonce for GCM (recommended)


@dataclass
class EncryptedSecret:
    """Container for envelope-encrypted secret data."""
    encrypted_value: bytes  # AES-256-GCM encrypted secret
    encryption_nonce: bytes  # 12-byte GCM nonce
    encrypted_dek: bytes  # DEK wrapped by KMS
    kms_key_version: str  # KMS key version for rotation tracking


@dataclass
class KMSConfig:
    """Configuration for Google Cloud KMS."""
    project_id: str
    location: str
    key_ring: str
    crypto_key: str

    @property
    def key_name(self) -> str:
        """Full resource name for the KMS key."""
        return (
            f"projects/{self.project_id}/locations/{self.location}/"
            f"keyRings/{self.key_ring}/cryptoKeys/{self.crypto_key}"
        )

    @classmethod
    def from_env(cls) -> "KMSConfig":
        """Load KMS config from environment variables."""
        return cls(
            project_id=os.getenv("KMS_PROJECT_ID") or os.getenv("GOOGLE_CLOUD_PROJECT", ""),
            location=os.getenv("KMS_LOCATION", "us-central1"),
            key_ring=os.getenv("KMS_KEY_RING", "foundrymatch-secrets"),
            crypto_key=os.getenv("KMS_CRYPTO_KEY", "workspace-secrets"),
        )


class KMSEnvelopeEncryption:
    """
    Envelope encryption service using Google Cloud KMS.

    For local development without KMS, uses a fallback mode that stores
    DEKs directly (encrypted with local key). This is NOT secure for production.
    """

    def __init__(
        self,
        config: Optional[KMSConfig] = None,
        allow_local_fallback: bool = False,
    ):
        """
        Initialize KMS envelope encryption.

        Args:
            config: KMS configuration. If None, loads from environment.
            allow_local_fallback: If True, use local encryption when KMS unavailable.
                                  Only enable for development/testing.
        """
        self.config = config or KMSConfig.from_env()
        self.allow_local_fallback = allow_local_fallback
        self._kms_client: Optional["kms.KeyManagementServiceClient"] = None

        # Local fallback key (only for development)
        self._local_fallback_key: Optional[bytes] = None
        if allow_local_fallback:
            fallback_key = os.getenv("ENCRYPTION_KEY", "")
            if fallback_key:
                # Derive 32-byte key from whatever key is provided
                from hashlib import sha256
                self._local_fallback_key = sha256(fallback_key.encode()).digest()

    @property
    def kms_client(self) -> "kms.KeyManagementServiceClient":
        """Lazy-load KMS client."""
        if self._kms_client is None:
            if not HAS_KMS:
                raise RuntimeError(
                    "google-cloud-kms is not installed. "
                    "Install with: pip install google-cloud-kms"
                )
            self._kms_client = kms.KeyManagementServiceClient()
        return self._kms_client

    def _use_local_fallback(self) -> bool:
        """Check if we should use local fallback mode."""
        if not self.allow_local_fallback:
            return False
        if not self._local_fallback_key:
            return False
        # Use fallback if KMS isn't available or configured
        if not HAS_KMS:
            return True
        if not self.config.project_id:
            return True
        return False

    def _local_wrap_dek(self, dek: bytes) -> tuple[bytes, str]:
        """
        Wrap DEK using local fallback (development only).

        WARNING: This is NOT secure for production. DEK is encrypted with
        a local key derived from ENCRYPTION_KEY, not KMS.
        """
        if not self._local_fallback_key:
            raise RuntimeError("Local fallback key not configured")

        nonce = secrets.token_bytes(NONCE_SIZE_BYTES)
        aesgcm = AESGCM(self._local_fallback_key)
        encrypted_dek = nonce + aesgcm.encrypt(nonce, dek, None)
        return encrypted_dek, "local-fallback-v1"

    def _local_unwrap_dek(self, encrypted_dek: bytes, key_version: str) -> bytes:
        """
        Unwrap DEK using local fallback (development only).
        """
        if not self._local_fallback_key:
            raise RuntimeError("Local fallback key not configured")
        if key_version != "local-fallback-v1":
            raise ValueError(f"Unknown local key version: {key_version}")

        nonce = encrypted_dek[:NONCE_SIZE_BYTES]
        ciphertext = encrypted_dek[NONCE_SIZE_BYTES:]
        aesgcm = AESGCM(self._local_fallback_key)
        return aesgcm.decrypt(nonce, ciphertext, None)

    def _kms_wrap_dek(self, dek: bytes) -> tuple[bytes, str]:
        """
        Wrap DEK using Google Cloud KMS.

        Returns:
            Tuple of (encrypted_dek, kms_key_version)
        """
        response = self.kms_client.encrypt(
            request={
                "name": self.config.key_name,
                "plaintext": dek,
            }
        )
        # Extract version from the response name
        # Format: projects/.../cryptoKeys/.../cryptoKeyVersions/1
        version = response.name.rsplit("/", 1)[-1] if "/" in response.name else "1"
        return response.ciphertext, f"kms-v{version}"

    def _kms_unwrap_dek(self, encrypted_dek: bytes, key_version: str) -> bytes:
        """
        Unwrap DEK using Google Cloud KMS.
        """
        # Cloud KMS symmetric decrypt expects the CryptoKey resource name, not
        # CryptoKeyVersion. The ciphertext carries version metadata internally.
        # We still persist key_version for audit/rotation reporting.
        response = self.kms_client.decrypt(
            request={
                "name": self.config.key_name,
                "ciphertext": encrypted_dek,
            }
        )
        return response.plaintext

    def encrypt(self, plaintext: str) -> EncryptedSecret:
        """
        Encrypt a secret using envelope encryption.

        Args:
            plaintext: The secret to encrypt (e.g., API key)

        Returns:
            EncryptedSecret containing all components needed for storage
        """
        # Step 1: Generate random DEK
        dek = secrets.token_bytes(DEK_SIZE_BYTES)

        # Step 2: Encrypt plaintext with DEK using AES-256-GCM
        nonce = secrets.token_bytes(NONCE_SIZE_BYTES)
        aesgcm = AESGCM(dek)
        encrypted_value = aesgcm.encrypt(nonce, plaintext.encode("utf-8"), None)

        # Step 3: Wrap DEK with KMS (or local fallback)
        if self._use_local_fallback():
            logger.warning(
                "Using local fallback encryption - NOT SECURE FOR PRODUCTION"
            )
            encrypted_dek, kms_key_version = self._local_wrap_dek(dek)
        else:
            encrypted_dek, kms_key_version = self._kms_wrap_dek(dek)

        # Step 4: Clear DEK from memory
        # Note: Python doesn't guarantee immediate clearing, but we try
        dek = b"\x00" * DEK_SIZE_BYTES  # noqa: F841

        return EncryptedSecret(
            encrypted_value=encrypted_value,
            encryption_nonce=nonce,
            encrypted_dek=encrypted_dek,
            kms_key_version=kms_key_version,
        )

    def decrypt(self, secret: EncryptedSecret) -> str:
        """
        Decrypt an envelope-encrypted secret.

        Args:
            secret: EncryptedSecret from storage

        Returns:
            Decrypted plaintext string
        """
        # Step 1: Unwrap DEK
        if secret.kms_key_version.startswith("local-fallback"):
            dek = self._local_unwrap_dek(secret.encrypted_dek, secret.kms_key_version)
        else:
            dek = self._kms_unwrap_dek(secret.encrypted_dek, secret.kms_key_version)

        # Step 2: Decrypt secret with DEK
        aesgcm = AESGCM(dek)
        plaintext_bytes = aesgcm.decrypt(
            secret.encryption_nonce,
            secret.encrypted_value,
            None,
        )

        # Step 3: Clear DEK from memory
        dek = b"\x00" * DEK_SIZE_BYTES  # noqa: F841

        return plaintext_bytes.decode("utf-8")

    def rewrap_dek(self, secret: EncryptedSecret) -> EncryptedSecret:
        """
        Re-wrap DEK with current KMS key version (for key rotation).

        This allows rotating KMS keys without re-encrypting the actual secrets.
        Only the DEK wrapper changes.

        Args:
            secret: Existing EncryptedSecret to rewrap

        Returns:
            New EncryptedSecret with same encrypted_value but new DEK wrapper
        """
        # Unwrap DEK with old key version
        if secret.kms_key_version.startswith("local-fallback"):
            dek = self._local_unwrap_dek(secret.encrypted_dek, secret.kms_key_version)
        else:
            dek = self._kms_unwrap_dek(secret.encrypted_dek, secret.kms_key_version)

        # Wrap with current key version
        if self._use_local_fallback():
            encrypted_dek, kms_key_version = self._local_wrap_dek(dek)
        else:
            encrypted_dek, kms_key_version = self._kms_wrap_dek(dek)

        # Clear DEK
        dek = b"\x00" * DEK_SIZE_BYTES  # noqa: F841

        return EncryptedSecret(
            encrypted_value=secret.encrypted_value,
            encryption_nonce=secret.encryption_nonce,
            encrypted_dek=encrypted_dek,
            kms_key_version=kms_key_version,
        )


# Global instance (lazy initialization)
_envelope_encryption: Optional[KMSEnvelopeEncryption] = None


def get_envelope_encryption(allow_local_fallback: Optional[bool] = None) -> KMSEnvelopeEncryption:
    """
    Get or create the global envelope encryption instance.

    Args:
        allow_local_fallback: Override for local fallback setting.
                              If None, uses ENV != "production" as default.
    """
    global _envelope_encryption
    if _envelope_encryption is None:
        if allow_local_fallback is None:
            # Default: allow fallback in non-production environments
            allow_local_fallback = os.getenv("ENV", "development") != "production"
        _envelope_encryption = KMSEnvelopeEncryption(
            allow_local_fallback=allow_local_fallback
        )
    return _envelope_encryption


def encrypt_secret(plaintext: str) -> EncryptedSecret:
    """Convenience function to encrypt a secret."""
    return get_envelope_encryption().encrypt(plaintext)


def decrypt_secret(secret: EncryptedSecret) -> str:
    """Convenience function to decrypt a secret."""
    return get_envelope_encryption().decrypt(secret)


def get_key_suffix(plaintext: str, length: int = 4) -> str:
    """
    Get the last N characters of a secret for display purposes.

    Args:
        plaintext: The secret value
        length: Number of characters to show (default 4)

    Returns:
        Last N characters of the secret
    """
    if len(plaintext) <= length:
        return "*" * len(plaintext)
    return plaintext[-length:]
