"""Authentication handlers: BearerToken, OAuth2, GA4GH Passports."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Protocol

import httpx
import jwt

from ai4gh.core.exceptions import AuthError

if TYPE_CHECKING:
    from ai4gh.core.config import AuthConfig


class AuthHandler(Protocol):
    """Protocol for authentication handlers."""

    async def get_headers(self) -> dict[str, str]:
        """Return authentication headers for a request."""
        ...


class NoAuth:
    """No authentication."""

    async def get_headers(self) -> dict[str, str]:
        """Return empty headers."""
        return {}


class BearerTokenAuth:
    """Simple bearer token authentication."""

    def __init__(self, token: str) -> None:
        self._token = token

    async def get_headers(self) -> dict[str, str]:
        """Return Authorization header with bearer token."""
        return {"Authorization": f"Bearer {self._token}"}


class OAuth2Auth:
    """OAuth2 client credentials authentication."""

    def __init__(
        self,
        token_url: str,
        client_id: str,
        client_secret: str,
    ) -> None:
        self._token_url = token_url
        self._client_id = client_id
        self._client_secret = client_secret
        self._access_token: str | None = None

    async def _fetch_token(self) -> str:
        """Fetch an OAuth2 access token."""
        async with httpx.AsyncClient() as client:
            response = await client.post(
                self._token_url,
                data={
                    "grant_type": "client_credentials",
                    "client_id": self._client_id,
                    "client_secret": self._client_secret,
                },
            )
            if not response.is_success:
                raise AuthError(
                    f"OAuth2 token fetch failed: {response.status_code}",
                    status_code=response.status_code,
                )
            data = response.json()
            return data["access_token"]  # type: ignore[no-any-return]

    async def get_headers(self) -> dict[str, str]:
        """Return Authorization header with OAuth2 token."""
        if self._access_token is None:
            self._access_token = await self._fetch_token()
        return {"Authorization": f"Bearer {self._access_token}"}

    def invalidate(self) -> None:
        """Invalidate the cached token, forcing re-fetch on next use."""
        self._access_token = None


class PassportAuth:
    """GA4GH Passport authentication via OIDC + Visa extraction.

    Passports are JWT tokens containing Visa claims that authorize
    access to specific datasets via GA4GH-compliant services.
    """

    def __init__(
        self,
        passport_token: str,
        *,
        verify: bool = True,
        algorithms: list[str] | None = None,
    ) -> None:
        self._token = passport_token
        self._verify = verify
        self._algorithms = algorithms or ["RS256", "ES256"]
        self._decoded: dict[str, Any] | None = None

    async def get_headers(self) -> dict[str, str]:
        """Return Authorization header with Passport token."""
        return {"Authorization": f"Bearer {self._token}"}

    def decode(self, *, verify: bool | None = None) -> dict[str, Any]:
        """Decode the Passport JWT.

        Args:
            verify: Override the verification setting. Use False for inspection only.

        Returns:
            Decoded JWT payload.
        """
        should_verify = verify if verify is not None else self._verify
        options = {"verify_signature": should_verify}

        try:
            self._decoded = jwt.decode(
                self._token,
                options=options,  # type: ignore[arg-type]
                algorithms=self._algorithms,
            )
        except jwt.InvalidTokenError as e:
            raise AuthError(f"Invalid Passport token: {e}") from e

        return self._decoded

    def get_visas(self) -> list[dict[str, Any]]:
        """Extract GA4GH Visas from the decoded Passport.

        Returns:
            List of Visa claim objects.
        """
        if self._decoded is None:
            self.decode(verify=False)
        assert self._decoded is not None

        return self._decoded.get("ga4gh_passport_v1", [])  # type: ignore[no-any-return]

    def has_visa_type(self, visa_type: str) -> bool:
        """Check if the Passport contains a Visa of the given type.

        Standard Visa types:
        - AffiliationAndRole
        - AcceptedTermsAndPolicies
        - ResearcherStatus
        - ControlledAccessGrants
        - LinkedIdentities
        """
        visas = self.get_visas()
        return any(v.get("type") == visa_type for v in visas)


def auto_detect_auth(config: AuthConfig) -> AuthHandler:
    """Auto-detect the appropriate auth handler from config.

    Args:
        config: Authentication configuration.

    Returns:
        An appropriate auth handler instance.
    """
    if config.token:
        return BearerTokenAuth(config.token)
    if config.token_url and config.client_id and config.client_secret:
        return OAuth2Auth(
            token_url=config.token_url,
            client_id=config.client_id,
            client_secret=config.client_secret,
        )
    return NoAuth()
