"""
PKCE (Proof Key for Code Exchange) utilities for OAuth 2.1 flows.

This module provides utilities for generating and validating PKCE parameters
as defined in RFC 7636, enhancing security for OAuth flows.

═══════════════════════════════════════════════════════════════════════════════
                         🔐 THE VERIFIER'S OATH 🔐
═══════════════════════════════════════════════════════════════════════════════

    A secret born from entropy's well,
    128 characters no one can tell.
    SHA256 transforms the phrase
    into a challenge through the haze.

    The verifier waits, patient, still—
    while challenge travels, testing will.
    When token comes back to the door,
    it proves you are who you were before.

    No interceptor along the wire
    can steal what they did not acquire.
    The secret never leaves your hand,
    only its shadow crossed the land.

    Consume and forget, the session dies.
    One proof, one use—that's how trust flies.
    In the space between request and grant,
    PKCE guards what hackers can't.

                                        — Field Notes, Jan 2026

═══════════════════════════════════════════════════════════════════════════════
"""

import base64
import hashlib
import secrets

from typing_extensions import Any, Dict, Tuple

from config.enhanced_logging import setup_logger

logger = setup_logger()


def generate_code_verifier(length: int = 128) -> str:
    """
    Generate a cryptographically random code verifier for PKCE.

    Args:
        length: Length of the code verifier (43-128 characters)

    Returns:
        URL-safe base64-encoded code verifier string

    Raises:
        ValueError: If length is not between 43 and 128
    """
    if not 43 <= length <= 128:
        raise ValueError("Code verifier length must be between 43 and 128 characters")

    # Generate random bytes (length * 3/4 bytes for base64 encoding)
    random_bytes = secrets.token_bytes(length * 3 // 4)

    # Base64 URL-safe encode and remove padding
    code_verifier = base64.urlsafe_b64encode(random_bytes).decode("utf-8").rstrip("=")

    # Ensure exact length
    code_verifier = code_verifier[:length]

    logger.debug(
        f"Generated PKCE code verifier: {code_verifier[:10]}... (length: {len(code_verifier)})"
    )
    return code_verifier


def generate_code_challenge(code_verifier: str, method: str = "S256") -> str:
    """
    Generate a code challenge from a code verifier.

    Args:
        code_verifier: The code verifier string
        method: Challenge method ("S256" or "plain")

    Returns:
        Code challenge string

    Raises:
        ValueError: If method is not supported
    """
    if method == "S256":
        # SHA256 hash and base64 URL-safe encode
        digest = hashlib.sha256(code_verifier.encode("utf-8")).digest()
        code_challenge = base64.urlsafe_b64encode(digest).decode("utf-8").rstrip("=")
    elif method == "plain":
        # Plain text (not recommended for security)
        code_challenge = code_verifier
    else:
        raise ValueError(f"Unsupported code challenge method: {method}")

    logger.debug(
        f"Generated PKCE code challenge using {method}: {code_challenge[:10]}..."
    )
    return code_challenge


def generate_pkce_pair(length: int = 128, method: str = "S256") -> Tuple[str, str]:
    """
    Generate a complete PKCE code verifier and challenge pair.

    Args:
        length: Length of the code verifier (43-128 characters)
        method: Challenge method ("S256" or "plain")

    Returns:
        Tuple of (code_verifier, code_challenge)
    """
    code_verifier = generate_code_verifier(length)
    code_challenge = generate_code_challenge(code_verifier, method)

    logger.info(f"Generated PKCE pair using {method} method")
    return code_verifier, code_challenge


def validate_code_verifier(
    code_verifier: str, code_challenge: str, method: str = "S256"
) -> bool:
    """
    Validate that a code verifier matches the code challenge.

    Args:
        code_verifier: The code verifier to validate
        code_challenge: The expected code challenge
        method: Challenge method used ("S256" or "plain")

    Returns:
        True if the code verifier is valid, False otherwise
    """
    try:
        expected_challenge = generate_code_challenge(code_verifier, method)
        is_valid = secrets.compare_digest(expected_challenge, code_challenge)

        logger.debug(f"PKCE validation result: {'VALID' if is_valid else 'INVALID'}")
        return is_valid
    except Exception as e:
        logger.error(f"PKCE validation error: {e}")
        return False


def create_authorization_params(
    client_id: str,
    redirect_uri: str,
    scope: str,
    state: str = None,
    use_pkce: bool = True,
) -> Dict[str, Any]:
    """
    Create authorization URL parameters with optional PKCE.

    Args:
        client_id: OAuth client ID
        redirect_uri: Redirect URI for callback
        scope: OAuth scopes (space-separated)
        state: Optional state parameter
        use_pkce: Whether to include PKCE parameters

    Returns:
        Dictionary of authorization parameters including PKCE if enabled
    """
    params = {
        "client_id": client_id,
        "response_type": "code",
        "redirect_uri": redirect_uri,
        "scope": scope,
    }

    if state:
        params["state"] = state

    pkce_data = None
    if use_pkce:
        code_verifier, code_challenge = generate_pkce_pair()
        params["code_challenge"] = code_challenge
        params["code_challenge_method"] = "S256"

        # Return PKCE data separately for storage
        pkce_data = {
            "code_verifier": code_verifier,
            "code_challenge": code_challenge,
            "code_challenge_method": "S256",
        }

        logger.info(f"Created authorization params with PKCE for client: {client_id}")
    else:
        logger.info(
            f"Created authorization params without PKCE for client: {client_id}"
        )

    return {"params": params, "pkce_data": pkce_data}


class PKCEManager:
    """
    Manager class for handling PKCE parameters throughout OAuth flows.

    This class provides a convenient interface for generating, storing,
    and retrieving PKCE parameters during OAuth authorization flows.
    """

    def __init__(self):
        self._storage = {}

    def create_pkce_session(
        self, session_id: str, method: str = "S256"
    ) -> Dict[str, str]:
        """
        Create a new PKCE session with generated parameters.

        Args:
            session_id: Unique session identifier
            method: PKCE challenge method

        Returns:
            Dictionary with code_challenge and method for authorization URL
        """
        code_verifier, code_challenge = generate_pkce_pair(method=method)

        self._storage[session_id] = {
            "code_verifier": code_verifier,
            "code_challenge": code_challenge,
            "code_challenge_method": method,
        }

        logger.info(f"Created PKCE session: {session_id}")

        return {"code_challenge": code_challenge, "code_challenge_method": method}

    def get_code_verifier(self, session_id: str) -> str:
        """
        Retrieve the code verifier for a session.

        Args:
            session_id: Session identifier

        Returns:
            Code verifier string

        Raises:
            KeyError: If session not found
        """
        if session_id not in self._storage:
            raise KeyError(f"PKCE session not found: {session_id}")

        return self._storage[session_id]["code_verifier"]

    def validate_and_consume(self, session_id: str, code_challenge: str) -> str:
        """
        Validate a session and consume the PKCE parameters.

        Args:
            session_id: Session identifier
            code_challenge: Expected code challenge

        Returns:
            Code verifier for token exchange

        Raises:
            KeyError: If session not found
            ValueError: If validation fails
        """
        if session_id not in self._storage:
            raise KeyError(f"PKCE session not found: {session_id}")

        session_data = self._storage[session_id]
        stored_challenge = session_data["code_challenge"]

        if not secrets.compare_digest(stored_challenge, code_challenge):
            raise ValueError("PKCE code challenge validation failed")

        # Consume the session (remove it)
        code_verifier = session_data["code_verifier"]
        del self._storage[session_id]

        logger.info(f"PKCE session validated and consumed: {session_id}")
        return code_verifier

    def cleanup_expired_sessions(self, max_age_seconds: int = 3600):
        """
        Clean up expired PKCE sessions.

        Args:
            max_age_seconds: Maximum age of sessions to keep
        """
        # Note: This is a simple implementation. In production, you'd want
        # to store creation timestamps and implement proper cleanup.
        logger.debug(
            f"PKCE session cleanup not implemented (sessions: {len(self._storage)})"
        )


# Global PKCE manager instance
pkce_manager = PKCEManager()
