"""
FIPS 140-3 Compliant AES-256-GCM Encryption

This module provides authenticated encryption for sensitive data such as
OAuth tokens, API keys, and other credentials.

Ciphertext Layout:
    |1B version|12B nonce|AES-256-GCM(ciphertext + 16B auth-tag)|

Key Source (checked in order):
    1. CREDENTIAL_ENCRYPTION_KEY - Base64-encoded 32-byte key
    2. CREDENTIAL_ENCRYPTION_KEY_POINTER - Vault path to retrieve key

Security Notes:
    - AES-256-GCM is FIPS 140-3 approved (NIST SP 800-38D)
    - 12-byte nonce is NIST recommended for GCM
    - Each encryption uses a unique random nonce
    - Authentication tag prevents tampering
"""

import base64
import json
import logging
import os
import secrets
from functools import lru_cache
from typing import Union

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

logger = logging.getLogger(__name__)

# Constants
_VERSION = 1  # Ciphertext format version (increment if format changes)
_NONCE_SIZE = 12  # 96 bits - NIST recommended for GCM
_KEY_SIZE = 32  # 256 bits for AES-256


class EncryptionKeyError(Exception):
    """Raised when encryption key cannot be loaded or is invalid."""
    pass


@lru_cache(maxsize=1)
def _load_key() -> bytes:
    """
    Load encryption key from environment or vault.

    Checks in order:
    1. CREDENTIAL_ENCRYPTION_KEY - Base64 encoded 32-byte key
    2. CREDENTIAL_ENCRYPTION_KEY_POINTER - Vault path

    Returns:
        32-byte encryption key

    Raises:
        EncryptionKeyError: If no key configured or key is invalid
    """
    # Option 1: Direct key from environment (Base64 encoded)
    key_b64 = os.environ.get("CREDENTIAL_ENCRYPTION_KEY")
    if key_b64:
        try:
            key = base64.b64decode(key_b64)
            if len(key) != _KEY_SIZE:
                raise EncryptionKeyError(
                    f"CREDENTIAL_ENCRYPTION_KEY must be {_KEY_SIZE} bytes, got {len(key)}"
                )
            return key
        except EncryptionKeyError:
            raise
        except Exception as e:
            raise EncryptionKeyError(f"Invalid CREDENTIAL_ENCRYPTION_KEY: {e}") from e

    # Option 2: Key from vault via pointer
    pointer = os.environ.get("CREDENTIAL_ENCRYPTION_KEY_POINTER")
    if pointer:
        try:
            from empowernow_common.vault import get_vault_provider
            provider = get_vault_provider()
            secret = provider.get_secret(pointer)
            key_value = secret.get("key") or secret.get("value")
            if not key_value:
                raise EncryptionKeyError(
                    f"Vault secret at '{pointer}' missing 'key' or 'value' field"
                )
            key = base64.b64decode(key_value)
            if len(key) != _KEY_SIZE:
                raise EncryptionKeyError(
                    f"Vault key must be {_KEY_SIZE} bytes, got {len(key)}"
                )
            return key
        except EncryptionKeyError:
            raise
        except ImportError:
            raise EncryptionKeyError(
                "CREDENTIAL_ENCRYPTION_KEY_POINTER set but vault module not available"
            )
        except Exception as e:
            raise EncryptionKeyError(f"Failed to load key from vault: {e}") from e

    # No key configured
    raise EncryptionKeyError(
        "No encryption key configured. Set CREDENTIAL_ENCRYPTION_KEY (Base64) or "
        "CREDENTIAL_ENCRYPTION_KEY_POINTER (vault path) environment variable."
    )


def clear_key_cache() -> None:
    """
    Clear cached encryption key.

    Call this after key rotation to force reloading the key
    on the next encrypt/decrypt operation.
    """
    _load_key.cache_clear()


def encrypt(plaintext: Union[bytes, str, dict, list]) -> bytes:
    """
    Encrypt plaintext using AES-256-GCM (FIPS 140-3 compliant).

    Args:
        plaintext: Data to encrypt. Accepts:
            - bytes: Used directly
            - str: Encoded to UTF-8
            - dict/list: JSON serialized then encoded

    Returns:
        Ciphertext bytes with layout: |1B version|12B nonce|ciphertext+tag|

    Raises:
        EncryptionKeyError: If encryption key not configured
        TypeError: If plaintext type is not supported

    Example:
        >>> tokens = {"access_token": "abc123", "refresh_token": "xyz789"}
        >>> ciphertext = encrypt(tokens)
        >>> # Store ciphertext in database BLOB column
    """
    # Normalize input to bytes
    if isinstance(plaintext, (dict, list)):
        plaintext = json.dumps(plaintext, separators=(",", ":")).encode("utf-8")
    elif isinstance(plaintext, str):
        plaintext = plaintext.encode("utf-8")
    elif isinstance(plaintext, bytes):
        pass  # Already bytes
    else:
        raise TypeError(
            f"plaintext must be bytes, str, dict, or list, got {type(plaintext).__name__}"
        )

    # Load encryption key
    key = _load_key()

    # Create AES-GCM cipher
    aesgcm = AESGCM(key)

    # Generate random 12-byte nonce (NIST recommended for GCM)
    # IMPORTANT: Never reuse a nonce with the same key
    nonce = secrets.token_bytes(_NONCE_SIZE)

    # Encrypt (automatically appends 16-byte authentication tag)
    ciphertext = aesgcm.encrypt(nonce, plaintext, associated_data=None)

    # Build output: version + nonce + ciphertext
    return bytes([_VERSION]) + nonce + ciphertext


def decrypt(blob: bytes) -> Union[dict, list, str, bytes]:
    """
    Decrypt ciphertext produced by encrypt().

    Args:
        blob: Ciphertext bytes from encrypt()

    Returns:
        Decrypted data with automatic type detection:
            - Valid JSON object → dict
            - Valid JSON array → list
            - Valid UTF-8 → str
            - Otherwise → bytes

    Raises:
        ValueError: If ciphertext is too short or version is unsupported
        cryptography.exceptions.InvalidTag: If ciphertext was tampered with
        EncryptionKeyError: If encryption key not configured

    Example:
        >>> ciphertext = encrypt({"access_token": "abc123"})
        >>> plaintext = decrypt(ciphertext)
        >>> print(plaintext["access_token"])  # "abc123"
    """
    # Minimum length: 1 (version) + 12 (nonce) + 16 (auth tag) + 1 (min data)
    min_length = 1 + _NONCE_SIZE + 16
    if not blob or len(blob) < min_length:
        raise ValueError(
            f"Ciphertext too short: {len(blob) if blob else 0} bytes, "
            f"minimum {min_length} bytes required"
        )

    # Parse layout: |1B version|12B nonce|ciphertext+tag|
    version = blob[0]

    if version != _VERSION:
        raise ValueError(
            f"Unsupported ciphertext version {version}, expected {_VERSION}. "
            "This may indicate key rotation or data migration is required."
        )

    nonce = blob[1:1 + _NONCE_SIZE]
    ciphertext = blob[1 + _NONCE_SIZE:]

    # Load key and decrypt
    key = _load_key()
    aesgcm = AESGCM(key)

    # Decrypt and verify authentication tag
    # Raises InvalidTag if ciphertext was tampered with
    plaintext_bytes = aesgcm.decrypt(nonce, ciphertext, associated_data=None)

    # Auto-detect output type
    try:
        text = plaintext_bytes.decode("utf-8")
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            return text
    except UnicodeDecodeError:
        return plaintext_bytes
