"""JWT identity extraction and validation.

Extracts AgentIdentity from Authorization: Bearer tokens.
Supports two modes:
  - Verified: validates JWT signature against a configured JWKS endpoint
  - Claimed: decodes JWT claims without signature validation (when no JWKS configured)

For v1, a single JWKS endpoint is supported. Multi-IdP and mTLS
are deferred to v2.
"""

from __future__ import annotations

import logging
from typing import Any

import jwt
from jwt import PyJWKClient

from attestmcp.models import AgentIdentity, IdentitySource, TrustLevel

logger = logging.getLogger(__name__)


class IdentityExtractor:
    """Extracts AgentIdentity from HTTP request headers.

    If a JWKS URL is configured, JWT signatures are validated and
    identities are marked as verified. Without JWKS, tokens are
    decoded without verification and marked as claimed.
    """

    def __init__(
        self,
        jwks_url: str | None = None,
        audience: str | None = None,
        issuer: str | None = None,
    ) -> None:
        self._jwks_client: PyJWKClient | None = None
        self._audience = audience
        self._issuer = issuer

        if jwks_url:
            try:
                self._jwks_client = PyJWKClient(jwks_url, cache_keys=True)
                logger.info("JWKS client initialized: %s", jwks_url)
            except Exception:
                logger.exception("Failed to initialize JWKS client")
                self._jwks_client = None

    def extract(
        self,
        authorization: str | None,
        session_id: str | None = None,
    ) -> AgentIdentity:
        """Extract identity from an Authorization header value.

        Args:
            authorization: The full Authorization header value
                (e.g., "Bearer eyJ..."). None if no header present.
            session_id: MCP session ID from Mcp-Session-Id header.

        Returns:
            AgentIdentity with appropriate trust level.
        """
        if not authorization:
            return AgentIdentity.anonymous(session_id=session_id)

        # Extract the token from "Bearer <token>"
        parts = authorization.split(" ", 1)
        if len(parts) != 2 or parts[0].lower() != "bearer":
            logger.warning("Invalid Authorization header format")
            return AgentIdentity.anonymous(session_id=session_id)

        token = parts[1]

        # Try verified extraction first (if JWKS configured)
        if self._jwks_client is not None:
            return self._extract_verified(token, session_id)

        # Fall back to claimed extraction (decode without verification)
        return self._extract_claimed(token, session_id)

    def _extract_verified(
        self, token: str, session_id: str | None
    ) -> AgentIdentity:
        """Decode and verify JWT signature against JWKS."""
        assert self._jwks_client is not None

        try:
            signing_key = self._jwks_client.get_signing_key_from_jwt(token)

            decode_options: dict[str, Any] = {}
            if self._audience is None:
                decode_options["verify_aud"] = False

            claims = jwt.decode(
                token,
                signing_key.key,
                algorithms=["RS256", "ES256"],
                audience=self._audience,
                issuer=self._issuer,
                options=decode_options,  # type: ignore[arg-type]
            )

            return self._claims_to_identity(
                claims,
                session_id=session_id,
                verified=True,
                trust_level=TrustLevel.VERIFIED,
            )

        except jwt.ExpiredSignatureError:
            logger.warning("JWT has expired")
            return AgentIdentity(
                source=IdentitySource.OAUTH_JWT,
                session_id=session_id,
                verified=False,
                trust_level=TrustLevel.CLAIMED,
            )
        except jwt.InvalidTokenError:
            logger.warning("JWT validation failed", exc_info=True)
            return AgentIdentity(
                source=IdentitySource.OAUTH_JWT,
                session_id=session_id,
                verified=False,
                trust_level=TrustLevel.CLAIMED,
            )

    def _extract_claimed(
        self, token: str, session_id: str | None
    ) -> AgentIdentity:
        """Decode JWT without signature verification.

        Used when no JWKS endpoint is configured. Identity is
        extracted but marked as claimed, not verified.
        """
        try:
            claims = jwt.decode(
                token,
                options={
                    "verify_signature": False,
                    "verify_exp": False,
                    "verify_aud": False,
                },
            )

            return self._claims_to_identity(
                claims,
                session_id=session_id,
                verified=False,
                trust_level=TrustLevel.CLAIMED,
            )

        except jwt.DecodeError:
            logger.warning("Failed to decode JWT")
            return AgentIdentity.anonymous(session_id=session_id)

    @staticmethod
    def _claims_to_identity(
        claims: dict[str, Any],
        *,
        session_id: str | None,
        verified: bool,
        trust_level: TrustLevel,
    ) -> AgentIdentity:
        """Map JWT claims to AgentIdentity fields.

        Standard claim mappings:
          - sub / email → user
          - azp / client_id → agent_name
          - org / organization → org
          - role / roles → role
        """
        user = claims.get("email") or claims.get("sub")
        agent_name = claims.get("azp") or claims.get("client_id")
        org = claims.get("org") or claims.get("organization")

        # Role can be a string or first item of a list
        role_claim = claims.get("role") or claims.get("roles")
        if isinstance(role_claim, list) and role_claim:
            role = str(role_claim[0])
        elif isinstance(role_claim, str):
            role = role_claim
        else:
            role = None

        return AgentIdentity(
            source=IdentitySource.OAUTH_JWT,
            user=str(user) if user else None,
            agent_name=str(agent_name) if agent_name else None,
            org=str(org) if org else None,
            role=role,
            session_id=session_id,
            verified=verified,
            trust_level=trust_level,
        )
