"""OAuth 2.0 authentication for MCP servers.

This module provides OAuth 2.0 authentication support for MCP servers,
including Authorization Code, Client Credentials, and Refresh Token flows.
"""

from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
from typing import Optional


class OAuthGrantType(Enum):
    """OAuth 2.0 grant types supported for MCP authentication."""

    AUTHORIZATION_CODE = "authorization_code"
    CLIENT_CREDENTIALS = "client_credentials"
    REFRESH_TOKEN = "refresh_token"


@dataclass
class OAuthConfig:
    """OAuth configuration for an MCP server.

    Attributes:
        client_id: OAuth client ID
        client_secret: OAuth client secret (optional for public clients)
        authorization_endpoint: URL for authorization (Authorization Code flow)
        token_endpoint: URL for token exchange
        grant_type: OAuth grant type to use
        scope: Optional space-separated list of scopes
        redirect_uri: Redirect URI for callbacks (defaults to localhost)
    """

    client_id: str
    client_secret: Optional[str]
    authorization_endpoint: str
    token_endpoint: str
    grant_type: OAuthGrantType
    scope: Optional[str] = None
    redirect_uri: str = "http://localhost:31415/oauth/callback"

    def __post_init__(self):
        """Validate OAuth configuration."""
        if isinstance(self.grant_type, str):
            self.grant_type = OAuthGrantType(self.grant_type)

        if self.grant_type == OAuthGrantType.CLIENT_CREDENTIALS:
            if not self.client_secret:
                raise ValueError(
                    "client_secret is required for client_credentials grant type"
                )

    @classmethod
    def from_dict(cls, data: dict) -> "OAuthConfig":
        """Create OAuthConfig from dictionary.

        Args:
            data: Configuration dictionary

        Returns:
            OAuthConfig instance
        """
        grant_type = data.get("grant_type", "authorization_code")
        if isinstance(grant_type, str):
            grant_type = OAuthGrantType(grant_type)

        return cls(
            client_id=data["client_id"],
            client_secret=data.get("client_secret"),
            authorization_endpoint=data["authorization_endpoint"],
            token_endpoint=data["token_endpoint"],
            grant_type=grant_type,
            scope=data.get("scope"),
            redirect_uri=data.get("redirect_uri", "http://localhost:31415/oauth/callback"),
        )

    def to_dict(self) -> dict:
        """Convert OAuthConfig to dictionary.

        Returns:
            Configuration dictionary
        """
        return {
            "client_id": self.client_id,
            "client_secret": self.client_secret,
            "authorization_endpoint": self.authorization_endpoint,
            "token_endpoint": self.token_endpoint,
            "grant_type": self.grant_type.value,
            "scope": self.scope,
            "redirect_uri": self.redirect_uri,
        }


@dataclass
class OAuthToken:
    """OAuth access token with metadata.

    Attributes:
        access_token: The access token string
        token_type: Token type (usually "Bearer")
        expires_in: Token lifetime in seconds
        refresh_token: Optional refresh token for getting new access tokens
        scope: Granted scopes (may differ from requested)
        issued_at: When the token was issued
    """

    access_token: str
    token_type: str
    expires_in: int
    refresh_token: Optional[str] = None
    scope: Optional[str] = None
    issued_at: datetime = field(default_factory=datetime.now)

    def is_expired(self, buffer_seconds: int = 300) -> bool:
        """Check if token is expired or will expire soon.

        Args:
            buffer_seconds: Consider token expired this many seconds early
                           (default 5 minutes to allow time for refresh)

        Returns:
            True if token is expired or expires within buffer_seconds
        """
        if self.expires_in <= 0:
            return True

        expiry_time = self.issued_at + timedelta(seconds=self.expires_in)
        buffer_time = datetime.now() + timedelta(seconds=buffer_seconds)

        return buffer_time >= expiry_time

    def expires_at(self) -> datetime:
        """Get the absolute expiry time of the token.

        Returns:
            datetime when token expires
        """
        return self.issued_at + timedelta(seconds=self.expires_in)

    @classmethod
    def from_dict(cls, data: dict) -> "OAuthToken":
        """Create OAuthToken from token response dictionary.

        Args:
            data: Token response from OAuth server

        Returns:
            OAuthToken instance
        """
        return cls(
            access_token=data["access_token"],
            token_type=data.get("token_type", "Bearer"),
            expires_in=data.get("expires_in", 3600),
            refresh_token=data.get("refresh_token"),
            scope=data.get("scope"),
            issued_at=datetime.now(),
        )

    def to_dict(self) -> dict:
        """Convert OAuthToken to dictionary for storage.

        Returns:
            Token dictionary
        """
        return {
            "access_token": self.access_token,
            "token_type": self.token_type,
            "expires_in": self.expires_in,
            "refresh_token": self.refresh_token,
            "scope": self.scope,
            "issued_at": self.issued_at.isoformat(),
        }

    @classmethod
    def from_stored_dict(cls, data: dict) -> "OAuthToken":
        """Create OAuthToken from stored dictionary.

        Args:
            data: Token dictionary from storage

        Returns:
            OAuthToken instance
        """
        issued_at = data.get("issued_at")
        if isinstance(issued_at, str):
            issued_at = datetime.fromisoformat(issued_at)
        elif issued_at is None:
            issued_at = datetime.now()

        return cls(
            access_token=data["access_token"],
            token_type=data.get("token_type", "Bearer"),
            expires_in=data.get("expires_in", 3600),
            refresh_token=data.get("refresh_token"),
            scope=data.get("scope"),
            issued_at=issued_at,
        )


class OAuthError(Exception):
    """Base exception for OAuth errors."""

    pass


class OAuthAuthorizationError(OAuthError):
    """User denied authorization or authorization failed."""

    def __init__(self, error: str, error_description: Optional[str] = None):
        self.error = error
        self.error_description = error_description
        message = f"Authorization failed: {error}"
        if error_description:
            message += f" - {error_description}"
        super().__init__(message)


class OAuthTokenError(OAuthError):
    """Token exchange or refresh failed."""

    def __init__(self, error: str, error_description: Optional[str] = None):
        self.error = error
        self.error_description = error_description
        message = f"Token error: {error}"
        if error_description:
            message += f" - {error_description}"
        super().__init__(message)


class OAuthTimeoutError(OAuthError):
    """OAuth flow timed out waiting for user interaction."""

    def __init__(self, timeout_seconds: int):
        self.timeout_seconds = timeout_seconds
        super().__init__(f"OAuth flow timed out after {timeout_seconds} seconds")
