"""Pluggable auth principal models.

These describe *who* is invoking a skill. The runtime auth provider produces
an instance of the agent's declared ``auth_model`` and hands it to the
:class:`RunContext`.
"""
from __future__ import annotations

import inspect
from collections.abc import Mapping
from typing import Any, Generic, TypeVar

from pydantic import BaseModel, ConfigDict, Field

AuthT = TypeVar("AuthT", bound=BaseModel)


class NoAuth(BaseModel):
    """Public agent: no caller identity required."""

    model_config = ConfigDict(extra="forbid", frozen=True)


class APIKeyAuth(BaseModel):
    """Caller authenticated by a long-lived API key."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    api_key_id: str
    scopes: list[str] = Field(default_factory=list)


class JWTAuth(BaseModel):
    """Caller authenticated by a JWT (typically from a user-facing login)."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    sub: str
    org_id: str | None = None
    email: str | None = None
    scopes: list[str] = Field(default_factory=list)


class AuthError(Exception):
    """Raised by auth resolvers when a caller token is missing or invalid."""

    def __init__(self, message: str, *, status_code: int = 401) -> None:
        super().__init__(message)
        self.status_code = status_code
        self.message = message


class AuthResolver(Generic[AuthT]):
    """Resolve an inbound bearer token into an agent's typed auth principal.

    Agent authors set ``auth_model`` to the principal shape their skills want
    and optionally set ``auth_resolver`` to an instance of this class. The
    resolver may validate JWTs, call a customer API, hit an OIDC userinfo or
    introspection endpoint, or bridge any other identity system.
    """

    async def resolve(
        self,
        token: str | None,
        *,
        headers: Mapping[str, str],
        agent: Any,
    ) -> AuthT:
        raise NotImplementedError


class NoAuthResolver(AuthResolver[NoAuth]):
    async def resolve(
        self,
        token: str | None,
        *,
        headers: Mapping[str, str],
        agent: Any,
    ) -> NoAuth:
        return NoAuth()


class APIKeyAuthResolver(AuthResolver[APIKeyAuth]):
    """Default resolver for ``auth_model = APIKeyAuth``.

    Requests must provide the configured bearer value. Public agents should use
    ``NoAuth`` rather than an API-key model without configured verification.
    """

    def __init__(
        self,
        *,
        accepted_key: str | None = None,
        api_key_id: str = "configured",
        scopes: list[str] | None = None,
    ) -> None:
        self.accepted_key = accepted_key
        self.api_key_id = api_key_id
        self.scopes = scopes or []

    async def resolve(
        self,
        token: str | None,
        *,
        headers: Mapping[str, str],
        agent: Any,
    ) -> APIKeyAuth:
        if self.accepted_key is None:
            raise AuthError("API key auth is not configured", status_code=500)
        if token is None:
            raise AuthError("missing bearer token")
        if token != self.accepted_key:
            raise AuthError("invalid bearer token")
        return APIKeyAuth(api_key_id=self.api_key_id, scopes=list(self.scopes))


class StaticAuthResolver(AuthResolver[AuthT]):
    """Resolver for tests/local adapters that always returns one principal."""

    def __init__(self, principal: AuthT) -> None:
        self.principal = principal

    async def resolve(
        self,
        token: str | None,
        *,
        headers: Mapping[str, str],
        agent: Any,
    ) -> AuthT:
        return self.principal


class RemoteBearerAuthResolver(AuthResolver[AuthT]):
    """Resolve bearer tokens by calling an external API.

    Works with OIDC ``userinfo`` endpoints and most homegrown
    ``GET /me``/``POST /introspect`` APIs. The response JSON is mapped into
    ``auth_model``. For SAML-backed apps, put your SAML/session exchange
    behind a bearer-token endpoint and use this same resolver.
    """

    def __init__(
        self,
        url: str,
        *,
        auth_model: type[AuthT],
        method: str = "GET",
        timeout_seconds: float = 5.0,
        token_header: str = "authorization",
        token_prefix: str = "bearer",
        active_field: str | None = "active",
        subject_field: str = "sub",
        email_field: str = "email",
        org_field: str = "org_id",
        scopes_field: str = "scope",
        extra_fields: Mapping[str, str] | None = None,
    ) -> None:
        self.url = url
        self.auth_model = auth_model
        self.method = method.upper()
        self.timeout_seconds = timeout_seconds
        self.token_header = token_header
        self.token_prefix = token_prefix
        self.active_field = active_field
        self.subject_field = subject_field
        self.email_field = email_field
        self.org_field = org_field
        self.scopes_field = scopes_field
        self.extra_fields = dict(extra_fields or {})

    async def resolve(
        self,
        token: str | None,
        *,
        headers: Mapping[str, str],
        agent: Any,
    ) -> AuthT:
        if token is None:
            raise AuthError("missing bearer token")
        import httpx

        request_headers = {
            self.token_header: f"{self.token_prefix} {token}".strip()
        }
        async with httpx.AsyncClient(timeout=self.timeout_seconds) as client:
            if self.method == "POST":
                resp = await client.post(self.url, headers=request_headers, data={"token": token})
            else:
                resp = await client.get(self.url, headers=request_headers)
        if resp.status_code in (401, 403):
            raise AuthError("invalid bearer token", status_code=resp.status_code)
        if resp.status_code >= 400:
            raise AuthError(
                f"auth resolver upstream returned HTTP {resp.status_code}",
                status_code=502,
            )
        try:
            data = resp.json()
        except ValueError as exc:
            raise AuthError(
                "auth resolver returned invalid JSON",
                status_code=502,
            ) from exc
        if not isinstance(data, dict):
            raise AuthError("auth resolver returned non-object JSON", status_code=502)
        if self.active_field and data.get(self.active_field) is False:
            raise AuthError("inactive bearer token")
        try:
            return self.auth_model.model_validate(self._map_payload(data))
        except Exception as exc:  # noqa: BLE001
            raise AuthError(
                f"auth resolver payload did not match {self.auth_model.__name__}: {exc}",
                status_code=502,
            ) from exc

    def _map_payload(self, data: dict[str, Any]) -> dict[str, Any]:
        payload: dict[str, Any] = {}
        sub = data.get(self.subject_field) or data.get("sub") or data.get("id")
        if sub is not None:
            payload["sub"] = str(sub)
        email = data.get(self.email_field)
        if email is not None:
            payload["email"] = str(email)
        org = data.get(self.org_field) or data.get("org") or data.get("organization_id")
        if org is not None:
            payload["org_id"] = str(org)
        scopes = data.get(self.scopes_field)
        if isinstance(scopes, str):
            payload["scopes"] = [s for s in scopes.split() if s]
        elif isinstance(scopes, list):
            payload["scopes"] = [str(s) for s in scopes]
        for target, source in self.extra_fields.items():
            if source in data:
                payload[target] = data[source]
        # Preserve exact matching fields for custom auth models.
        for key in self.auth_model.model_fields:
            if key in data and key not in payload:
                payload[key] = data[key]
        return payload


OIDCUserInfoAuthResolver = RemoteBearerAuthResolver


async def resolve_auth(
    resolver: AuthResolver[AuthT],
    token: str | None,
    *,
    headers: Mapping[str, str],
    agent: Any,
) -> AuthT:
    """Call sync or async custom resolver implementations."""
    result = resolver.resolve(token, headers=headers, agent=agent)
    if inspect.isawaitable(result):
        return await result
    return result
