"""Core OAuth 2.1 authorization code flow orchestrator.

This module contains the two public-facing classes:

- :class:`OAuthFlowResult` — immutable result returned by
  :meth:`UnifiedOAuthFlow.handle_callback`.
- :class:`UnifiedOAuthFlow` — the main orchestrator that both the MCP Gateway
  and CRUD Service delegate to.

Both the MCP Gateway and the CRUD Service implement ~300 lines of duplicated
OAuth mechanics independently.  This module consolidates them into a single,
tested implementation so each service only needs to supply service-specific
concerns (MCP elicitation notifications, CRUD popup nonces, etc.).
"""

from __future__ import annotations

import base64
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Awaitable, Callable
from urllib.parse import urlencode

import httpx

from empowernow_common.oauth.par import generate_pkce_challenge
from empowernow_common.oauth.security import generate_secure_token
from empowernow_common.oauth.user_flow.config import OAuthFlowConfig, OAuthStateStore
from empowernow_common.oauth.user_flow.enrichment import MetadataEnrichmentService
from empowernow_common.oauth.user_flow.exceptions import (
    OAuthProviderError,
    OAuthRateLimitError,
    OAuthSecretResolutionError,
    OAuthStateError,
)
from empowernow_common.vault.base import ReadableVaultProvider
from empowernow_common.vault.user_oauth.base import BaseUserOAuthProvider
from empowernow_common.vault.user_oauth.metadata import OAuthMetadataExtractor

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# OAuthFlowResult
# ---------------------------------------------------------------------------


@dataclass
class OAuthFlowResult:
    """Immutable result from a completed OAuth authorization code flow.

    Returned by :meth:`UnifiedOAuthFlow.handle_callback` after the code has
    been exchanged, metadata enriched, and the credential stored.

    Attributes:
        credential_id:  Opaque ID returned by the credential provider
                        (e.g. a UUID from PostgresUserOAuthProvider).
        user_arn:       ARN of the user who authorized the OAuth flow.
        system_key:     External system key (e.g. ``"jira_cloud"``).
        metadata:       Provider-specific metadata extracted from the token
                        response and enrichment calls (e.g. ``cloud_id``).
        tenant_id:      Optional tenant identifier for multi-tenancy.
        extra_state:    Service-specific data passed through from
                        :meth:`UnifiedOAuthFlow.generate_authorization_url`
                        (e.g. ``{"elicitation_id": "..."}`` for MCP Gateway).
    """

    credential_id: str
    user_arn: str
    system_key: str
    metadata: dict[str, Any]
    tenant_id: str | None = None
    extra_state: dict[str, Any] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# UnifiedOAuthFlow
# ---------------------------------------------------------------------------


class UnifiedOAuthFlow:
    """Core OAuth 2.1 authorization code flow with PKCE.

    Handles the shared mechanics that are otherwise duplicated across services:

    - PKCE challenge generation (S256) — via :func:`~empowernow_common.oauth.par.generate_pkce_challenge`
    - State token generation + Redis storage (one-time use) — via :func:`~empowernow_common.oauth.security.generate_secure_token`
    - Rate limiting (per-user, configurable window)
    - Authorization URL construction (with ``extra_auth_params`` support)
    - Token exchange (``client_secret_basic`` and ``client_secret_post``)
    - Just-in-time secret resolution via :class:`~empowernow_common.vault.base.ReadableVaultProvider`
    - Metadata enrichment (config-driven post-token-exchange API calls)
    - Provider-specific metadata extraction via :class:`~empowernow_common.vault.user_oauth.metadata.OAuthMetadataExtractor`
    - Encrypted credential storage via :class:`~empowernow_common.vault.user_oauth.base.BaseUserOAuthProvider`

    **Does NOT handle** (remains service-specific):

    - HTTP endpoint definitions or HTML callback pages
    - MCP elicitation completion notifications
    - CRUD popup nonce generation
    - User-facing error message formatting

    **Security**: ``client_secret`` is never stored on this class or on
    :class:`~empowernow_common.oauth.user_flow.config.OAuthFlowConfig`.
    The config holds a vault URI reference (``client_secret_ref``) which is
    resolved just-in-time via ``ReadableVaultProvider`` during
    :meth:`_exchange_code` only.  The secret exists in memory only for the
    duration of the token-exchange HTTP POST.
    """

    STATE_TTL: int = 300       # seconds — 5 minutes
    RATE_LIMIT_MAX: int = 10   # callbacks per window
    RATE_LIMIT_WINDOW: int = 60  # seconds

    def __init__(
        self,
        state_store: OAuthStateStore,
        vault_provider: ReadableVaultProvider,
        credential_provider: BaseUserOAuthProvider,
        http_client: httpx.AsyncClient,
        enrichment_service: MetadataEnrichmentService | None = None,
        rate_limit_max: int = RATE_LIMIT_MAX,
        rate_limit_window: int = RATE_LIMIT_WINDOW,
    ) -> None:
        self._state = state_store
        self._vault = vault_provider
        self._creds = credential_provider
        self._http = http_client
        self._enrichment = enrichment_service or MetadataEnrichmentService(http_client)
        self._rate_limit_max = rate_limit_max
        self._rate_limit_window = rate_limit_window

    # ------------------------------------------------------------------
    # Public API: generate_authorization_url
    # ------------------------------------------------------------------

    async def generate_authorization_url(
        self,
        config: OAuthFlowConfig,
        user_arn: str,
        system_key: str,
        tenant_id: str | None = None,
        extra_state: dict[str, Any] | None = None,
    ) -> tuple[str, str]:
        """Build an authorization URL with PKCE and store one-time state.

        Steps:

        1. Generate a PKCE ``(code_verifier, code_challenge)`` pair (S256).
        2. Generate a cryptographically secure ``state_token``.
        3. Persist state in Redis via ``state_store.store_state()``.
        4. Build the authorization URL with query parameters.

        The vault provider is **not** called here — secrets are resolved only
        during :meth:`handle_callback` → :meth:`_exchange_code`.

        Args:
            config:      OAuth client configuration.
            user_arn:    ARN of the user initiating the flow.
            system_key:  External system key (e.g. ``"jira_cloud"``).
            tenant_id:   Optional tenant for multi-tenancy.
            extra_state: Arbitrary service-specific data to carry through the
                         flow and return in
                         :attr:`OAuthFlowResult.extra_state` on callback
                         (e.g. ``{"elicitation_id": "...", "session_id": "..."}``.

        Returns:
            A ``(authorization_url, state_token)`` tuple.  The service layer
            redirects the user to ``authorization_url`` and stores
            ``state_token`` to match the callback.
        """
        # 1. PKCE — reuse existing SDK utility
        code_verifier, code_challenge = generate_pkce_challenge(method="S256")

        # 2. State token — reuse existing SDK utility
        state_token = generate_secure_token(length=32)

        # 3. Persist state (one-time use, 5-minute TTL)
        state_data: dict[str, Any] = {
            "user_arn": user_arn,
            "system_key": system_key,
            "code_verifier": code_verifier,
            "redirect_uri": config.redirect_uri,
            "tenant_id": tenant_id,
            "extra_state": extra_state or {},
            "created_at": datetime.now(timezone.utc).isoformat(),
        }
        await self._state.store_state(
            f"oauth_state:{state_token}", state_data, self.STATE_TTL
        )

        # 4. Build authorization URL
        params: dict[str, str] = {
            "response_type": "code",
            "client_id": config.client_id,
            "redirect_uri": config.redirect_uri,
            "state": state_token,
        }
        if config.scopes:
            params["scope"] = " ".join(config.scopes)
        if config.pkce_required:
            params["code_challenge"] = code_challenge
            params["code_challenge_method"] = "S256"
        if config.audience:
            params["audience"] = config.audience
        if config.extra_auth_params:
            params.update(config.extra_auth_params)

        auth_url = f"{config.authorization_url}?{urlencode(params)}"
        logger.info(
            "Authorization URL generated",
            extra={"system_key": system_key, "user_arn": user_arn},
        )
        return auth_url, state_token

    # ------------------------------------------------------------------
    # Public API: handle_callback
    # ------------------------------------------------------------------

    async def handle_callback(
        self,
        code: str,
        state: str,
        config_resolver: Callable[[str], Awaitable[OAuthFlowConfig]],
    ) -> OAuthFlowResult:
        """Complete the OAuth flow: validate state, exchange code, enrich, store.

        Steps:

        1. Consume state (one-time use) — raises :class:`~empowernow_common.oauth.user_flow.exceptions.OAuthStateError` if invalid/expired.
        2. Rate-limit check per user ARN — raises :class:`~empowernow_common.oauth.user_flow.exceptions.OAuthRateLimitError` if exceeded.
        3. Resolve :class:`~empowernow_common.oauth.user_flow.config.OAuthFlowConfig` via *config_resolver*.
        4. Token exchange via :meth:`_exchange_code` (secret resolved just-in-time).
        5. Optional metadata enrichment (config-driven API calls).
        6. Provider-specific metadata extraction via :class:`~empowernow_common.vault.user_oauth.metadata.OAuthMetadataExtractor`.
        7. Merge ``config.static_metadata`` if present.
        8. Store encrypted credential via ``credential_provider``.
        9. Return :class:`OAuthFlowResult`.

        Args:
            code:            Authorization code from the provider callback.
            state:           State token echoed back by the provider.
            config_resolver: Async callable ``(system_key) -> OAuthFlowConfig``.
                             Called after state is consumed so the config is
                             fetched only for valid requests.

        Returns:
            :class:`OAuthFlowResult` with ``credential_id``, ``metadata``,
            and ``extra_state`` populated.

        Raises:
            OAuthStateError:           Invalid or expired state token.
            OAuthRateLimitError:       Too many callbacks for this user.
            OAuthSecretResolutionError: Vault cannot resolve the client secret.
            OAuthProviderError:        Token exchange failed.
            OAuthEnrichmentError:      Required enrichment step failed.
        """
        # 1. Consume state — one-time use, atomic GET+DELETE
        state_data = await self._state.consume_state(f"oauth_state:{state}")
        if state_data is None:
            raise OAuthStateError("Invalid or expired OAuth state token")

        user_arn: str = state_data["user_arn"]
        system_key: str = state_data["system_key"]
        code_verifier: str = state_data["code_verifier"]
        redirect_uri: str = state_data["redirect_uri"]
        tenant_id: str | None = state_data.get("tenant_id")
        extra_state: dict[str, Any] = state_data.get("extra_state", {})

        # 2. Rate-limit check (per user ARN)
        within_limit = await self._state.check_rate_limit(
            f"oauth_rate:{user_arn}",
            self._rate_limit_max,
            self._rate_limit_window,
        )
        if not within_limit:
            raise OAuthRateLimitError(
                f"Rate limit exceeded: {self._rate_limit_max} callbacks "
                f"per {self._rate_limit_window}s for user {user_arn}"
            )

        # 3. Resolve OAuth config for this system
        config = await config_resolver(system_key)

        # 4. Token exchange (client_secret resolved just-in-time inside)
        token_data = await self._exchange_code(
            config=config,
            system_key=system_key,
            code=code,
            code_verifier=code_verifier,
            redirect_uri=redirect_uri,
        )

        # 5. Optional metadata enrichment (config-driven API calls)
        if config.metadata_enrichment:
            token_data = await self._enrichment.enrich(
                token_data, config.metadata_enrichment
            )

        # 6. Provider-specific metadata extraction
        provider_type = config.provider_type or system_key
        metadata: dict[str, Any] = OAuthMetadataExtractor.extract_metadata(
            provider_type=provider_type,
            token_response=token_data,
        )

        # 7. Merge static metadata (e.g. GitHub Enterprise base_url)
        if config.static_metadata:
            metadata.update(config.static_metadata)

        # 8. Store encrypted credential
        credential_id: str = await self._creds.store_credential(
            user_arn=user_arn,
            system_key=system_key,
            tokens=token_data,
            metadata=metadata,
            tenant_id=tenant_id,
        )

        logger.info(
            "OAuth flow completed",
            extra={
                "system_key": system_key,
                "user_arn": user_arn,
                "credential_id": credential_id,
            },
        )

        return OAuthFlowResult(
            credential_id=credential_id,
            user_arn=user_arn,
            system_key=system_key,
            metadata=metadata,
            tenant_id=tenant_id,
            extra_state=extra_state,
        )

    # ------------------------------------------------------------------
    # Private: _exchange_code
    # ------------------------------------------------------------------

    async def _exchange_code(
        self,
        config: OAuthFlowConfig,
        system_key: str,
        code: str,
        code_verifier: str,
        redirect_uri: str,
    ) -> dict[str, Any]:
        """Exchange an authorization code for tokens via HTTP POST.

        **Security**: ``client_secret`` is resolved just-in-time from the vault
        via ``config.client_secret_ref`` (a vault URI).  The secret exists in
        memory only as a local variable within this method — it is never stored
        on ``self``, ``config``, or any other persistent object.  After the HTTP
        POST completes the local variable goes out of scope and is eligible for
        garbage collection.

        Args:
            config:        OAuth client configuration.
            system_key:    External system key (used for logging only).
            code:          Authorization code from the provider callback.
            code_verifier: PKCE verifier generated during
                           :meth:`generate_authorization_url`.
            redirect_uri:  Callback URL (must match the one used in the
                           authorization request).

        Returns:
            The parsed JSON token response dict from the provider.

        Raises:
            OAuthSecretResolutionError: Vault cannot resolve ``client_secret_ref``.
            OAuthProviderError: Network error or non-200 response from the
                                token endpoint.
        """
        # 1. Resolve client secret just-in-time (secret never stored on self)
        try:
            client_secret = await self._vault.get_secret(config.client_secret_ref)
        except Exception as exc:
            raise OAuthSecretResolutionError(
                f"Failed to resolve client secret from {config.client_secret_ref}: {exc}"
            ) from exc

        # 2. Build POST body
        body: dict[str, str] = {
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": redirect_uri,
            "client_id": config.client_id,
        }
        if config.pkce_required:
            body["code_verifier"] = code_verifier

        headers: dict[str, str] = {"Accept": "application/json"}

        # 3. Apply client authentication method
        if config.client_auth_method == "client_secret_basic":
            credentials = base64.b64encode(
                f"{config.client_id}:{client_secret}".encode()
            ).decode()
            headers["Authorization"] = f"Basic {credentials}"
        else:
            # client_secret_post (default) — secret in POST body, then discarded
            body["client_secret"] = client_secret

        # 4. POST to token endpoint
        # client_secret local var is no longer referenced after this point
        # and becomes eligible for GC once the function returns
        try:
            response = await self._http.post(
                config.token_url,
                data=body,
                headers=headers,
                timeout=30.0,
            )
        except httpx.RequestError as exc:
            raise OAuthProviderError(
                f"Token exchange network error for {system_key}: {exc}"
            ) from exc

        # 5. Handle non-200 responses
        if response.status_code != 200:
            error_detail = self._parse_error_response(response)
            raise OAuthProviderError(
                f"Token exchange failed ({response.status_code}) for {system_key}: {error_detail}"
            )

        logger.info("Token exchange successful", extra={"system_key": system_key})
        return response.json()  # type: ignore[no-any-return]

    # ------------------------------------------------------------------
    # Private: _parse_error_response
    # ------------------------------------------------------------------

    @staticmethod
    def _parse_error_response(response: httpx.Response) -> str:
        """Extract a human-readable error description from a provider error response.

        Tries to parse the body as JSON and returns ``error_description`` or
        ``error``; falls back to the first 200 characters of the raw text.

        Args:
            response: The non-200 HTTP response from the token endpoint.

        Returns:
            A short string describing the error.
        """
        try:
            data = response.json()
            return str(data.get("error_description", data.get("error", "Unknown error")))
        except Exception:
            return response.text[:200]
