"""JWT authentication handler for API Gateway events."""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass
from typing import Mapping

import jwt
from aws_lambda_typing.events import APIGatewayProxyEventV1
from jwt import InvalidTokenError

from ..authz import Principal, Role
from ..context import set_principal
from ..exceptions import DomainValidationError, UnauthorizedError
from ..protocols import IJwtTokenHandler

log = logging.getLogger(__name__)


@dataclass(frozen=True, slots=True)
class JwtTokenHandler(IJwtTokenHandler):
    """Extract Authorization headers and hydrate principal in context."""

    public_key: str
    issuer: str = os.getenv("AUTH_JWT_ISSUER", "stock-agent").strip() or "stock-agent"
    audience: str = os.getenv("AUTH_JWT_AUDIENCE", "stock-agent-api").strip() or "stock-agent-api"

    def __call__(self, event: APIGatewayProxyEventV1) -> None:
        try:
            token = self._extract_bearer_token(event.get("headers") or {}) or ""
            if not token:
                log.warning("No JWT token provided, setting principal to guest")
                set_principal(Principal.guest())
            else:
                claims = jwt.decode(
                    token,
                    self.public_key,
                    algorithms=["RS256"],
                    issuer=self.issuer,
                    audience=self.audience,
                    options={"require": ["exp", "iat", "nbf", "iss", "aud"]},
                )
                set_principal(self._principal_from_claims(claims))
        except InvalidTokenError as exc:
            log.warning("Invalid JWT token")
            raise UnauthorizedError("Invalid token") from exc

    def _principal_from_claims(self, claims: Mapping[str, object]) -> Principal:
        try:
            name = claims["name"]
            role = claims["role"]
            tenant_id = claims["tenant_id"]
        except KeyError as exc:
            raise InvalidTokenError("Missing identity claims") from exc

        if not isinstance(name, str) or not isinstance(role, str) or not isinstance(tenant_id, str):
            raise InvalidTokenError("Missing or invalid identity claims")

        try:
            return Principal.of(name=name, role=Role(role), tenant_id=tenant_id)
        except (DomainValidationError, TypeError, ValueError) as exc:
            raise InvalidTokenError("Missing or invalid identity claims") from exc

    def _extract_bearer_token(self, headers: Mapping[str, str]) -> str | None:
        for key, value in headers.items():
            if key.lower() == "authorization":
                auth_value = value.strip()
                break
        else:
            return None

        parts = auth_value.split(" ", 1)
        if len(parts) != 2 or parts[0].lower() != "bearer":
            raise UnauthorizedError("Authorization header must be of type Bearer")
        return parts[1].strip()


__all__ = ["JwtTokenHandler"]

