"""OAuth 2.0 flow management for MCP authentication.

This module implements the OAuth 2.0 flows:
- Authorization Code with PKCE
- Client Credentials
- Refresh Token
"""

import base64
import hashlib
import secrets
import urllib.parse
import webbrowser
from typing import Optional

import aiohttp

from .oauth import (
    OAuthAuthorizationError,
    OAuthConfig,
    OAuthError,
    OAuthGrantType,
    OAuthTimeoutError,
    OAuthToken,
    OAuthTokenError,
)
from .oauth_callback import OAuthCallbackServer


class OAuthFlowManager:
    """Manages OAuth 2.0 authentication flows for MCP servers."""

    def __init__(self, timeout: int = 300):
        """Initialize OAuth flow manager.

        Args:
            timeout: Timeout in seconds for interactive flows (default 5 minutes)
        """
        self.timeout = timeout
        self._callback_server: Optional[OAuthCallbackServer] = None

    async def authenticate(self, config: OAuthConfig) -> OAuthToken:
        """Authenticate using the configured grant type.

        Args:
            config: OAuth configuration

        Returns:
            OAuth token

        Raises:
            OAuthError: If authentication fails
        """
        if config.grant_type == OAuthGrantType.AUTHORIZATION_CODE:
            return await self.authorize_authorization_code(config)
        elif config.grant_type == OAuthGrantType.CLIENT_CREDENTIALS:
            return await self.authorize_client_credentials(config)
        else:
            raise OAuthError(f"Unsupported grant type: {config.grant_type}")

    async def authorize_authorization_code(self, config: OAuthConfig) -> OAuthToken:
        """Perform Authorization Code flow with PKCE.

        Flow:
        1. Generate PKCE code verifier and challenge
        2. Build authorization URL with challenge
        3. Open browser to authorization URL
        4. Start local callback server
        5. Wait for authorization callback
        6. Exchange authorization code for token

        Args:
            config: OAuth configuration

        Returns:
            OAuth token

        Raises:
            OAuthAuthorizationError: If user denies or authorization fails
            OAuthTimeoutError: If flow times out
            OAuthTokenError: If token exchange fails
        """
        # Generate PKCE parameters
        code_verifier = self._generate_code_verifier()
        code_challenge = self._generate_code_challenge(code_verifier)
        state = secrets.token_urlsafe(32)

        # Build authorization URL
        auth_params = {
            "client_id": config.client_id,
            "response_type": "code",
            "redirect_uri": config.redirect_uri,
            "state": state,
            "code_challenge": code_challenge,
            "code_challenge_method": "S256",
        }

        if config.scope:
            auth_params["scope"] = config.scope

        auth_url = self._build_url(config.authorization_endpoint, auth_params)

        # Start callback server
        port = int(urllib.parse.urlparse(config.redirect_uri).port or 31415)
        self._callback_server = OAuthCallbackServer(port=port)

        try:
            # Open browser
            print(f"\n🔐 Opening browser for authentication...")
            print(f"If browser doesn't open, visit: {auth_url}\n")

            if not webbrowser.open(auth_url):
                print(f"⚠️  Could not open browser automatically.")
                print(f"Please visit this URL: {auth_url}\n")

            # Wait for callback
            code, error = await self._callback_server.start(timeout=self.timeout)

            if error:
                raise OAuthAuthorizationError(
                    error, "User denied authorization or an error occurred"
                )

            if not code:
                raise OAuthTimeoutError(self.timeout)

            # Exchange code for token
            return await self._exchange_code_for_token(
                config, code, code_verifier
            )

        finally:
            if self._callback_server:
                await self._callback_server.stop()
                self._callback_server = None

    async def authorize_client_credentials(self, config: OAuthConfig) -> OAuthToken:
        """Perform Client Credentials flow.

        Flow:
        1. POST to token endpoint with client_id and client_secret
        2. Return access token

        Args:
            config: OAuth configuration

        Returns:
            OAuth token

        Raises:
            OAuthTokenError: If token request fails
        """
        if not config.client_secret:
            raise OAuthError("client_secret is required for client_credentials flow")

        token_params = {
            "grant_type": "client_credentials",
            "client_id": config.client_id,
            "client_secret": config.client_secret,
        }

        if config.scope:
            token_params["scope"] = config.scope

        return await self._request_token(config.token_endpoint, token_params)

    async def refresh_token(
        self, config: OAuthConfig, refresh_token: str
    ) -> OAuthToken:
        """Refresh an expired access token.

        Args:
            config: OAuth configuration
            refresh_token: Refresh token from previous authentication

        Returns:
            New OAuth token

        Raises:
            OAuthTokenError: If refresh fails
        """
        token_params = {
            "grant_type": "refresh_token",
            "refresh_token": refresh_token,
            "client_id": config.client_id,
        }

        if config.client_secret:
            token_params["client_secret"] = config.client_secret

        return await self._request_token(config.token_endpoint, token_params)

    async def _exchange_code_for_token(
        self, config: OAuthConfig, code: str, code_verifier: str
    ) -> OAuthToken:
        """Exchange authorization code for access token.

        Args:
            config: OAuth configuration
            code: Authorization code from callback
            code_verifier: PKCE code verifier

        Returns:
            OAuth token

        Raises:
            OAuthTokenError: If token exchange fails
        """
        token_params = {
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": config.redirect_uri,
            "client_id": config.client_id,
            "code_verifier": code_verifier,
        }

        if config.client_secret:
            token_params["client_secret"] = config.client_secret

        return await self._request_token(config.token_endpoint, token_params)

    async def _request_token(self, token_url: str, params: dict) -> OAuthToken:
        """Make token request to OAuth server.

        Args:
            token_url: Token endpoint URL
            params: Token request parameters

        Returns:
            OAuth token

        Raises:
            OAuthTokenError: If request fails
        """
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    token_url,
                    data=params,
                    headers={
                        "Content-Type": "application/x-www-form-urlencoded",
                        "Accept": "application/json",
                    },
                ) as response:
                    data = await response.json()

                    if response.status != 200:
                        error = data.get("error", "unknown_error")
                        error_description = data.get("error_description")
                        raise OAuthTokenError(error, error_description)

                    return OAuthToken.from_dict(data)

        except aiohttp.ClientError as e:
            raise OAuthTokenError("network_error", str(e))
        except Exception as e:
            if isinstance(e, (OAuthTokenError, OAuthError)):
                raise
            raise OAuthTokenError("unexpected_error", str(e))

    @staticmethod
    def _generate_code_verifier() -> str:
        """Generate PKCE code verifier.

        Returns:
            URL-safe base64 encoded random string (43-128 characters)
        """
        # Generate 32 random bytes = 43 base64 characters
        random_bytes = secrets.token_bytes(32)
        return base64.urlsafe_b64encode(random_bytes).decode("utf-8").rstrip("=")

    @staticmethod
    def _generate_code_challenge(code_verifier: str) -> str:
        """Generate PKCE code challenge from verifier.

        Args:
            code_verifier: Code verifier

        Returns:
            URL-safe base64 encoded SHA256 hash of verifier
        """
        verifier_bytes = code_verifier.encode("utf-8")
        challenge_bytes = hashlib.sha256(verifier_bytes).digest()
        return base64.urlsafe_b64encode(challenge_bytes).decode("utf-8").rstrip("=")

    @staticmethod
    def _build_url(base_url: str, params: dict) -> str:
        """Build URL with query parameters.

        Args:
            base_url: Base URL
            params: Query parameters

        Returns:
            Complete URL with encoded parameters
        """
        query_string = urllib.parse.urlencode(params)
        separator = "&" if "?" in base_url else "?"
        return f"{base_url}{separator}{query_string}"
