"""
Core OAuth flow manager for Brynq Runtime.

Orchestrates the full OAuth 2.0 authorization code flow for all supported
providers (Anthropic, Google, OpenAI). Handles:
    - PKCE (Proof Key for Code Exchange) for providers that require it
    - State parameter for CSRF protection
    - Token exchange, refresh, and revocation
    - Encrypted local token storage (tokens NEVER leave the user's machine)

Usage::

    manager = OAuthManager()
    url = await manager.start_login("anthropic")
    # Browser opens to url, user consents, redirect hits callback server
    tokens = await manager.handle_callback("anthropic", code, state)
    # Later:
    token = await manager.get_token("anthropic")  # auto-refreshes if expired
"""

from __future__ import annotations

import base64
import hashlib
import json
import logging
import os
import secrets
import ssl
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import webbrowser
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional

from .callback_server import OAuthCallbackServer
from .providers import ProviderConfig, get_provider, list_provider_names
from .token_store import TokenSet, TokenStore

log = logging.getLogger("brynq.runtime.auth.oauth_manager")


# ── Data Classes ─────────────────────────────────────────────────────────


@dataclass(frozen=True)
class Connection:
    """Status of a single provider connection."""

    provider: str
    provider_name: str
    status: str  # "connected", "expired", "not_connected"
    user_email: Optional[str] = None
    scopes: list[str] = field(default_factory=list)


# ── PKCE Helpers ─────────────────────────────────────────────────────────


def generate_pkce_pair() -> tuple[str, str]:
    """Generate a PKCE code verifier and challenge.

    Returns:
        (code_verifier, code_challenge) tuple.

    The code verifier is a random 43-128 character URL-safe string.
    The code challenge is its SHA-256 hash, base64url-encoded (no padding).
    Per RFC 7636 Section 4.1-4.2.
    """
    # 32 bytes -> 43 base64url characters (within 43-128 range)
    verifier_bytes = secrets.token_bytes(32)
    code_verifier = (
        base64.urlsafe_b64encode(verifier_bytes).rstrip(b"=").decode("ascii")
    )

    # RFC 7636 Section 4.1: code_verifier must be 43-128 characters
    if not (43 <= len(code_verifier) <= 128):
        raise ValueError(
            f"PKCE code verifier length {len(code_verifier)} is outside "
            f"the 43-128 character range required by RFC 7636"
        )

    # S256 challenge
    digest = hashlib.sha256(code_verifier.encode("ascii")).digest()
    code_challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii")

    return code_verifier, code_challenge


def generate_state() -> str:
    """Generate a cryptographically random state parameter for CSRF protection.

    Returns a 32-byte hex string (64 characters).
    """
    return secrets.token_hex(32)


# ── OAuth Manager ────────────────────────────────────────────────────────


class OAuthManager:
    """Manages OAuth 2.0 flows for all supported providers.

    Tokens are stored locally via TokenStore (Fernet encrypted,
    machine-bound). Tokens NEVER leave the user's machine.
    """

    def __init__(
        self,
        token_store: Optional[TokenStore] = None,
        callback_host: str = "127.0.0.1",
    ) -> None:
        """Initialize the OAuth manager.

        Args:
            token_store: Token storage backend. Uses default if not provided.
            callback_host: Host for the OAuth callback server.
        """
        self._store = token_store or TokenStore()
        self._callback_host = callback_host
        # Pending state -> (provider, code_verifier) for CSRF + PKCE
        self._pending_states: dict[str, tuple[str, Optional[str]]] = {}
        # Prevent two simultaneous OAuth login flows
        self._login_lock = threading.Lock()
        self._login_lock_held = False

    async def start_login(
        self,
        provider: str,
        open_browser: bool = True,
    ) -> str:
        """Initiate an OAuth login flow for a provider.

        Opens the user's browser to the provider's consent page and starts
        a localhost callback server to receive the redirect.

        Args:
            provider: Provider name ("anthropic", "google", "openai").
            open_browser: Whether to auto-open the browser.

        Returns:
            The authorization URL (also opened in browser if open_browser=True).

        Raises:
            ValueError: If the provider is unknown or client ID is not set.
        """
        if not self._login_lock.acquire(blocking=False):
            raise RuntimeError(
                "Another OAuth login flow is already in progress. "
                "Please wait for it to complete before starting a new one."
            )
        self._login_lock_held = True

        try:
            return await self._start_login_inner(provider, open_browser)
        except BaseException:
            self._login_lock_held = False
            self._login_lock.release()
            raise

    async def _start_login_inner(
        self,
        provider: str,
        open_browser: bool,
    ) -> str:
        """Inner login logic, called while holding _login_lock."""
        config = get_provider(provider)
        client_id = config.get_client_id()
        if not client_id:
            raise ValueError(
                f"No client ID configured for {config.name}. "
                f"Set the {config.client_id_env} environment variable."
            )

        # Set up callback server to get the redirect URI
        callback_server = OAuthCallbackServer(host=self._callback_host)
        redirect_uri = callback_server.redirect_uri

        # Generate state for CSRF protection
        state = generate_state()

        # Build authorization URL
        params: dict[str, str] = {
            "client_id": client_id,
            "redirect_uri": redirect_uri,
            "response_type": "code",
            "scope": " ".join(config.scopes),
            "state": state,
        }

        # PKCE if required
        code_verifier: Optional[str] = None
        if config.requires_pkce:
            code_verifier, code_challenge = generate_pkce_pair()
            params["code_challenge"] = code_challenge
            params["code_challenge_method"] = "S256"

        # Provider-specific extra params
        if config.extra_auth_params:
            params.update(config.extra_auth_params)

        auth_url = f"{config.auth_url}?{urllib.parse.urlencode(params)}"

        # Store pending state for CSRF validation
        self._pending_states[state] = (provider, code_verifier)

        log.info("Starting OAuth login for %s (state=%s...)", config.name, state[:8])

        if open_browser:
            webbrowser.open(auth_url)

        return auth_url

    async def handle_callback(
        self,
        provider: str,
        code: str,
        state: str,
    ) -> TokenSet:
        """Exchange an authorization code for tokens.

        Validates the state parameter against pending logins, exchanges
        the auth code at the provider's token endpoint, and stores the
        resulting tokens in the encrypted local store.

        Args:
            provider: Provider name.
            code: Authorization code from the callback.
            state: State parameter from the callback.

        Returns:
            The stored TokenSet.

        Raises:
            ValueError: If state doesn't match any pending login (CSRF).
            OAuthTokenError: If the token exchange fails.
        """
        try:
            # CSRF validation
            pending = self._pending_states.pop(state, None)
            if pending is None:
                raise ValueError(
                    "OAuth state mismatch — possible CSRF attack. "
                    "Please try logging in again."
                )

            expected_provider, code_verifier = pending
            if expected_provider.lower() != provider.lower():
                raise ValueError(
                    f"State mismatch: expected provider {expected_provider!r}, "
                    f"got {provider!r}"
                )

            config = get_provider(provider)
            client_id = config.get_client_id()
            if not client_id:
                raise ValueError(f"No client ID for {config.name}")

            # Build callback server redirect_uri (must match what was sent in auth request)
            # We use a standard format since the port may vary per login
            callback_server = OAuthCallbackServer(host=self._callback_host)
            redirect_uri = callback_server.redirect_uri

            # Exchange the authorization code for tokens
            token_data = await self._exchange_code(
                config=config,
                code=code,
                redirect_uri=redirect_uri,
                client_id=client_id,
                code_verifier=code_verifier,
            )

            # Build TokenSet from response
            tokens = TokenSet(
                access_token=token_data["access_token"],
                refresh_token=token_data.get("refresh_token"),
                expires_at=self._compute_expiry(token_data),
                scopes=token_data.get("scope", "").split() if token_data.get("scope") else config.scopes,
                user_email=token_data.get("email"),
            )

            # Store encrypted locally
            self._store.save_tokens(provider, tokens)

            log.info("OAuth login complete for %s", config.name)
            return tokens
        finally:
            if self._login_lock_held:
                self._login_lock_held = False
                self._login_lock.release()

    async def get_token(self, provider: str) -> Optional[str]:
        """Get a valid access token for a provider.

        Auto-refreshes if the token is expired and a refresh token exists.

        Args:
            provider: Provider name.

        Returns:
            The access token string, or None if not connected.
        """
        tokens = self._store.get_tokens(provider)
        if tokens is None:
            return None

        if tokens.is_expired:
            if tokens.refresh_token:
                try:
                    refreshed = await self.refresh_token(provider)
                    return refreshed
                except Exception:
                    log.warning(
                        "Token refresh failed for %s -- token may be expired",
                        provider,
                    )
                    return None
            else:
                log.warning(
                    "Token expired for %s and no refresh token available",
                    provider,
                )
                return None

        # Proactive refresh: if token expires within 5 minutes, refresh early
        if (
            tokens.expires_at is not None
            and tokens.expires_at - time.time() < 300
            and tokens.refresh_token
        ):
            log.warning(
                "Token for %s expires within 5 minutes, attempting early refresh",
                provider,
            )
            try:
                refreshed = await self.refresh_token(provider)
                return refreshed
            except Exception:
                log.warning(
                    "Early token refresh failed for %s -- returning current token",
                    provider,
                )

        return tokens.access_token

    def get_token_sync(self, provider: str) -> Optional[str]:
        """Synchronous version of get_token for non-async contexts.

        Returns the access token if available and not expired.
        Does NOT auto-refresh (refresh requires async HTTP).
        Returns None if no token or token is expired.
        """
        tokens = self._store.get_tokens(provider)
        if tokens is None:
            return None
        if tokens.is_expired:
            log.debug("OAuth token expired for %s (sync — cannot auto-refresh)", provider)
            return None
        return tokens.access_token

    async def refresh_token(self, provider: str) -> str:
        """Force-refresh an expired token using the refresh token.

        Args:
            provider: Provider name.

        Returns:
            The new access token.

        Raises:
            ValueError: If no refresh token is available.
            OAuthTokenError: If the refresh request fails.
        """
        tokens = self._store.get_tokens(provider)
        if tokens is None or tokens.refresh_token is None:
            raise ValueError(f"No refresh token available for {provider}")

        config = get_provider(provider)
        client_id = config.get_client_id()
        if not client_id:
            raise ValueError(f"No client ID for {config.name}")

        data: dict[str, str] = {
            "grant_type": "refresh_token",
            "refresh_token": tokens.refresh_token,
            "client_id": client_id,
        }

        client_secret = config.get_client_secret()
        if client_secret:
            data["client_secret"] = client_secret

        token_data = await self._post_token_request(config.token_url, data)

        # Update stored tokens
        new_tokens = TokenSet(
            access_token=token_data["access_token"],
            # Some providers return a new refresh token; keep the old one if not
            refresh_token=token_data.get("refresh_token", tokens.refresh_token),
            expires_at=self._compute_expiry(token_data),
            scopes=tokens.scopes,
            user_email=tokens.user_email,
        )
        self._store.save_tokens(provider, new_tokens)

        log.info("Refreshed access token for %s", config.name)
        return new_tokens.access_token

    async def logout(self, provider: str) -> None:
        """Revoke tokens and remove them from local storage.

        Args:
            provider: Provider name.
        """
        tokens = self._store.get_tokens(provider)

        # Attempt revocation if the provider supports it
        if tokens:
            config = get_provider(provider)
            if config.revoke_url:
                try:
                    await self._revoke_token(config, tokens)
                except Exception:
                    # Revocation is best-effort; always delete locally
                    log.warning("Token revocation failed for %s", provider)

        # Always remove from local store
        self._store.delete_tokens(provider)
        log.info("Logged out from %s", provider)

    async def list_connections(self) -> list[Connection]:
        """List all supported providers with their connection status.

        Returns:
            List of Connection objects for every supported provider.
        """
        connections: list[Connection] = []

        for provider_key in list_provider_names():
            config = get_provider(provider_key)
            tokens = self._store.get_tokens(provider_key)

            if tokens is None:
                status = "not_connected"
                email = None
                scopes: list[str] = []
            elif tokens.is_expired:
                status = "expired"
                email = tokens.user_email
                scopes = tokens.scopes
            else:
                status = "connected"
                email = tokens.user_email
                scopes = tokens.scopes

            connections.append(
                Connection(
                    provider=provider_key,
                    provider_name=config.name,
                    status=status,
                    user_email=email,
                    scopes=scopes,
                )
            )

        return connections

    # ── Private Helpers ──────────────────────────────────────────────

    async def _exchange_code(
        self,
        config: ProviderConfig,
        code: str,
        redirect_uri: str,
        client_id: str,
        code_verifier: Optional[str] = None,
    ) -> dict[str, Any]:
        """Exchange an authorization code for tokens at the token endpoint."""
        data: dict[str, str] = {
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": redirect_uri,
            "client_id": client_id,
        }

        # Google (and others) require client_secret for desktop app token exchange
        client_secret = config.get_client_secret()
        if client_secret:
            data["client_secret"] = client_secret

        if code_verifier:
            data["code_verifier"] = code_verifier

        return await self._post_token_request(config.token_url, data, timeout=10)

    async def _post_token_request(
        self, url: str, data: dict[str, str], timeout: int = 30
    ) -> dict[str, Any]:
        """POST a form-encoded request to a token endpoint.

        Uses stdlib urllib -- no extra dependencies needed.

        Args:
            url: Token endpoint URL.
            data: Form data to POST.
            timeout: Request timeout in seconds.

        Raises:
            OAuthTokenError: If the request fails or times out.
        """
        encoded = urllib.parse.urlencode(data).encode("utf-8")
        req = urllib.request.Request(
            url,
            data=encoded,
            method="POST",
            headers={
                "Content-Type": "application/x-www-form-urlencoded",
                "Accept": "application/json",
            },
        )

        ctx = ssl.create_default_context()
        try:
            with urllib.request.urlopen(req, timeout=timeout, context=ctx) as resp:
                body = resp.read()
                return json.loads(body)
        except urllib.error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")
            raise OAuthTokenError(
                f"Token request failed (HTTP {exc.code}): {body}"
            ) from exc
        except (TimeoutError, urllib.error.URLError) as exc:
            if isinstance(exc, urllib.error.URLError) and not isinstance(exc.reason, TimeoutError):
                raise OAuthTokenError(
                    f"Token request failed: {type(exc).__name__}: {exc}"
                ) from exc
            raise OAuthTokenError(
                f"Token exchange timed out after {timeout}s"
            ) from exc
        except Exception as exc:
            raise OAuthTokenError(
                f"Token request failed: {type(exc).__name__}: {exc}"
            ) from exc

    async def _revoke_token(
        self, config: ProviderConfig, tokens: TokenSet
    ) -> None:
        """Attempt to revoke tokens at the provider's revocation endpoint."""
        if not config.revoke_url:
            return

        # Revoke access token (and refresh token if separate)
        for token_value in [tokens.access_token, tokens.refresh_token]:
            if not token_value:
                continue

            data = urllib.parse.urlencode({"token": token_value}).encode("utf-8")
            req = urllib.request.Request(
                config.revoke_url,
                data=data,
                method="POST",
                headers={"Content-Type": "application/x-www-form-urlencoded"},
            )

            ctx = ssl.create_default_context()
            try:
                with urllib.request.urlopen(req, timeout=10, context=ctx):
                    pass
            except Exception:
                # Revocation is best-effort
                log.debug(
                    "Revocation request failed for %s", config.name
                )

    @staticmethod
    def _compute_expiry(token_data: dict[str, Any]) -> Optional[float]:
        """Compute the absolute expiry timestamp from a token response."""
        expires_in = token_data.get("expires_in")
        if expires_in is not None:
            return time.time() + float(expires_in)
        return None


# ── Exceptions ───────────────────────────────────────────────────────────


class OAuthTokenError(Exception):
    """Raised when a token exchange or refresh request fails."""
