"""OAuth Authorization Server Provider for mssql-mcp.

Implements MCP SDK's OAuthAuthorizationServerProvider protocol,
bridging to external identity providers (Duo, Auth0, generic OIDC).
"""

import logging
import secrets
import time
from typing import TYPE_CHECKING

from mcp.server.auth.provider import (
    AccessToken,
    AuthorizationCode,
    AuthorizationParams,
    AuthorizeError,
    OAuthAuthorizationServerProvider,
    RefreshToken,
    RegistrationError,
    construct_redirect_uri,
)
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken

from .storage import (
    InMemoryAuthStorage,
    PendingAuthorization,
    StoredAuthCode,
    StoredClient,
    StoredToken,
)

if TYPE_CHECKING:
    from .idp.base import BaseIdPAdapter

logger = logging.getLogger(__name__)


class MSSQLOAuthProvider(
    OAuthAuthorizationServerProvider[AuthorizationCode, RefreshToken, AccessToken]
):
    """OAuth provider that bridges Claude.ai DCR to external identity providers.

    This provider acts as an OAuth Authorization Server for Claude.ai clients
    while delegating actual user authentication to an external IdP (Duo, Auth0, etc.).

    Flow:
        1. Claude.ai registers via DCR -> stored in memory
        2. Claude.ai initiates /authorize -> redirects to external IdP
        3. User authenticates with IdP -> callback to /oauth/callback
        4. Callback generates auth code -> redirects to Claude callback
        5. Claude.ai exchanges code at /token -> returns access token
    """

    def __init__(
        self,
        idp_adapter: "BaseIdPAdapter",
        issuer_url: str,
        callback_path: str = "/oauth/callback",
        token_expiry: int = 3600,
        refresh_token_expiry: int = 86400 * 30,  # 30 days
        storage: InMemoryAuthStorage | None = None,
    ):
        """Initialize the OAuth provider.

        Args:
            idp_adapter: Adapter for the external identity provider
            issuer_url: Public URL of this server (e.g., https://mssql-mcp.example.com)
            callback_path: Path for IdP callback (default /oauth/callback)
            token_expiry: Access token lifetime in seconds (default 1 hour)
            refresh_token_expiry: Refresh token lifetime in seconds (default 30 days)
            storage: Auth state storage (default InMemoryAuthStorage)
        """
        self.idp_adapter = idp_adapter
        self.issuer_url = issuer_url.rstrip("/")
        self.callback_path = callback_path
        self.token_expiry = token_expiry
        self.refresh_token_expiry = refresh_token_expiry
        self.storage = storage or InMemoryAuthStorage()

        # Internal callback URL for external IdP
        self.idp_callback_url = f"{self.issuer_url}{self.callback_path}"

    # -------------------------------------------------------------------------
    # Client Registration (DCR)
    # -------------------------------------------------------------------------

    async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
        """Retrieve registered client by ID."""
        stored = self.storage.get_client(client_id)
        if stored is None:
            return None

        return OAuthClientInformationFull(
            client_id=stored.client_id,
            client_secret=stored.client_secret,
            client_name=stored.client_name,
            redirect_uris=stored.redirect_uris,
            grant_types=stored.grant_types,
            response_types=stored.response_types,
            scope=stored.scope,
            token_endpoint_auth_method=stored.token_endpoint_auth_method,
            client_id_issued_at=int(stored.created_at),
        )

    async def register_client(self, client_info: OAuthClientInformationFull) -> None:
        """Register a new OAuth client (DCR).

        Claude.ai uses DCR to register itself as a client.
        """
        # Validate redirect URIs
        if not client_info.redirect_uris:
            raise RegistrationError(
                error="invalid_redirect_uri",
                error_description="At least one redirect_uri is required",
            )

        # Generate client credentials if not provided
        client_id = client_info.client_id or f"client_{secrets.token_urlsafe(16)}"
        client_secret = client_info.client_secret or secrets.token_urlsafe(32)

        stored = StoredClient(
            client_id=client_id,
            client_secret=client_secret,
            client_name=client_info.client_name,
            redirect_uris=[str(uri) for uri in client_info.redirect_uris],
            grant_types=list(client_info.grant_types),
            response_types=list(client_info.response_types),
            scope=client_info.scope,
            token_endpoint_auth_method=(
                client_info.token_endpoint_auth_method or "client_secret_post"
            ),
        )

        self.storage.store_client(stored)

        # Update client_info with generated credentials
        client_info.client_id = client_id
        client_info.client_secret = client_secret
        client_info.client_id_issued_at = int(stored.created_at)

        logger.info(f"Registered client: {client_id} ({client_info.client_name})")

    # -------------------------------------------------------------------------
    # Authorization
    # -------------------------------------------------------------------------

    async def authorize(
        self, client: OAuthClientInformationFull, params: AuthorizationParams
    ) -> str:
        """Start authorization flow by redirecting to external IdP.

        Returns a URL to redirect the user to the external IdP for authentication.
        After IdP auth, user is redirected to our callback which then redirects
        to Claude's callback with an auth code.
        """
        # Generate state for the IdP redirect (contains original request info)
        idp_state = secrets.token_urlsafe(32)

        # Store pending authorization state
        pending = PendingAuthorization(
            state=idp_state,
            client_id=client.client_id,
            redirect_uri=str(params.redirect_uri),
            scope=" ".join(params.scopes) if params.scopes else "",
            code_challenge=params.code_challenge,
            code_challenge_method="S256",  # We always use S256
            claude_redirect_uri=str(params.redirect_uri),
            claude_state=params.state,  # Store Claude's original state to return later
        )
        self.storage.store_pending_auth(pending)

        # Generate authorization URL for external IdP
        idp_auth_url = await self.idp_adapter.get_authorization_url(
            redirect_uri=self.idp_callback_url,
            state=idp_state,
            # We don't use PKCE with the IdP - we use it between Claude and us
        )

        logger.info(f"Redirecting to IdP for client {client.client_id}")
        return idp_auth_url

    # -------------------------------------------------------------------------
    # Authorization Code Exchange
    # -------------------------------------------------------------------------

    async def load_authorization_code(
        self, client: OAuthClientInformationFull, authorization_code: str
    ) -> AuthorizationCode | None:
        """Load an authorization code for validation."""
        stored = self.storage.get_auth_code(authorization_code)
        if stored is None:
            return None

        if stored.client_id != client.client_id:
            logger.warning(f"Auth code client mismatch: {stored.client_id} != {client.client_id}")
            return None

        return AuthorizationCode(
            code=stored.code,
            scopes=stored.scope.split() if stored.scope else [],
            expires_at=stored.expires_at,
            client_id=stored.client_id,
            code_challenge=stored.code_challenge or "",
            redirect_uri=stored.redirect_uri,
            redirect_uri_provided_explicitly=True,
        )

    async def exchange_authorization_code(
        self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode
    ) -> OAuthToken:
        """Exchange authorization code for tokens."""
        # Generate tokens
        access_token = secrets.token_urlsafe(32)
        refresh_token = secrets.token_urlsafe(32)

        now = time.time()
        access_expires = now + self.token_expiry
        refresh_expires = now + self.refresh_token_expiry

        # Store access token
        self.storage.store_access_token(
            StoredToken(
                token=access_token,
                token_type="access",
                client_id=client.client_id,
                user_subject=authorization_code.client_id,  # Will be set properly in callback
                scope=" ".join(authorization_code.scopes),
                expires_at=access_expires,
            )
        )

        # Store refresh token
        self.storage.store_refresh_token(
            StoredToken(
                token=refresh_token,
                token_type="refresh",
                client_id=client.client_id,
                user_subject=authorization_code.client_id,
                scope=" ".join(authorization_code.scopes),
                expires_at=refresh_expires,
            )
        )

        logger.info(f"Issued tokens for client {client.client_id}")

        return OAuthToken(
            access_token=access_token,
            token_type="Bearer",
            expires_in=self.token_expiry,
            refresh_token=refresh_token,
            scope=" ".join(authorization_code.scopes) if authorization_code.scopes else None,
        )

    # -------------------------------------------------------------------------
    # Refresh Token Exchange
    # -------------------------------------------------------------------------

    async def load_refresh_token(
        self, client: OAuthClientInformationFull, refresh_token: str
    ) -> RefreshToken | None:
        """Load a refresh token for validation."""
        stored = self.storage.get_refresh_token(refresh_token)
        if stored is None:
            return None

        if stored.client_id != client.client_id:
            logger.warning("Refresh token client mismatch")
            return None

        return RefreshToken(
            token=stored.token,
            client_id=stored.client_id,
            scopes=stored.scope.split() if stored.scope else [],
            expires_at=int(stored.expires_at) if stored.expires_at else None,
        )

    async def exchange_refresh_token(
        self,
        client: OAuthClientInformationFull,
        refresh_token: RefreshToken,
        scopes: list[str],
    ) -> OAuthToken:
        """Exchange refresh token for new tokens (token rotation)."""
        # Revoke old refresh token
        self.storage.revoke_refresh_token(refresh_token.token)

        # Generate new tokens
        new_access_token = secrets.token_urlsafe(32)
        new_refresh_token = secrets.token_urlsafe(32)

        now = time.time()
        access_expires = now + self.token_expiry
        refresh_expires = now + self.refresh_token_expiry

        scope_str = " ".join(scopes) if scopes else " ".join(refresh_token.scopes)

        # Store new access token
        self.storage.store_access_token(
            StoredToken(
                token=new_access_token,
                token_type="access",
                client_id=client.client_id,
                user_subject="",  # Preserved from original auth
                scope=scope_str,
                expires_at=access_expires,
            )
        )

        # Store new refresh token
        self.storage.store_refresh_token(
            StoredToken(
                token=new_refresh_token,
                token_type="refresh",
                client_id=client.client_id,
                user_subject="",
                scope=scope_str,
                expires_at=refresh_expires,
            )
        )

        logger.info(f"Rotated tokens for client {client.client_id}")

        return OAuthToken(
            access_token=new_access_token,
            token_type="Bearer",
            expires_in=self.token_expiry,
            refresh_token=new_refresh_token,
            scope=scope_str,
        )

    # -------------------------------------------------------------------------
    # Token Validation
    # -------------------------------------------------------------------------

    async def load_access_token(self, token: str) -> AccessToken | None:
        """Validate and load an access token."""
        stored = self.storage.get_access_token(token)
        if stored is None:
            return None

        return AccessToken(
            token=stored.token,
            client_id=stored.client_id,
            scopes=stored.scope.split() if stored.scope else [],
            expires_at=int(stored.expires_at) if stored.expires_at else None,
        )

    # -------------------------------------------------------------------------
    # Token Revocation
    # -------------------------------------------------------------------------

    async def revoke_token(self, token: AccessToken | RefreshToken) -> None:
        """Revoke an access or refresh token."""
        if isinstance(token, AccessToken):
            self.storage.revoke_access_token(token.token)
        else:
            self.storage.revoke_refresh_token(token.token)
        logger.info(f"Revoked token for client {token.client_id}")

    # -------------------------------------------------------------------------
    # IdP Callback Handling
    # -------------------------------------------------------------------------

    async def handle_idp_callback(self, code: str, state: str) -> str:
        """Handle callback from external IdP after user authentication.

        This is called from the /oauth/callback route handler.

        Args:
            code: Authorization code from the IdP
            state: State parameter (our internal state with pending auth info)

        Returns:
            Redirect URL to Claude's callback with our authorization code

        Raises:
            AuthorizeError: If callback is invalid
        """
        # Retrieve pending authorization
        pending = self.storage.get_pending_auth(state)
        if pending is None:
            raise AuthorizeError(
                error="invalid_request",
                error_description="Invalid or expired state parameter",
            )

        # Exchange code with IdP for tokens
        try:
            idp_tokens = await self.idp_adapter.exchange_code(
                code=code,
                redirect_uri=self.idp_callback_url,
            )
        except Exception as e:
            logger.error(f"IdP code exchange failed: {e}")
            raise AuthorizeError(
                error="access_denied",
                error_description="Authentication failed with identity provider",
            ) from e

        # Validate IdP token to get user info
        try:
            user_info = await self.idp_adapter.validate_token(idp_tokens.access_token)
            if user_info is None:
                raise AuthorizeError(
                    error="access_denied",
                    error_description="Failed to validate identity",
                )
        except AuthorizeError:
            raise  # Re-raise our own errors
        except Exception as e:
            logger.error(f"IdP token validation failed: {e}")
            raise AuthorizeError(
                error="access_denied",
                error_description="Failed to validate identity",
            ) from e

        # Generate authorization code for Claude
        auth_code = secrets.token_urlsafe(32)

        stored_code = StoredAuthCode(
            code=auth_code,
            client_id=pending.client_id,
            redirect_uri=pending.claude_redirect_uri,
            scope=pending.scope,
            state=state,
            code_challenge=pending.code_challenge,
            code_challenge_method=pending.code_challenge_method,
            user_subject=user_info.subject,
            user_claims=user_info.raw_claims or {},
        )
        self.storage.store_auth_code(stored_code)

        # Build redirect URL to Claude's callback (use Claude's original state)
        redirect_url: str = construct_redirect_uri(
            pending.claude_redirect_uri,
            code=auth_code,
            state=pending.claude_state,
        )

        logger.info(f"IdP callback successful for user {user_info.subject}")
        return redirect_url
