"""AUTX downstream token verification helpers for sellers."""

from __future__ import annotations

import time
from typing import Any

import httpx
from jose import jwt
from jose.exceptions import JWTError


class AutxTokenVerifier:
    def __init__(
        self,
        jwks_url: str = "http://localhost:8000/.well-known/jwks.json",
        issuer: str = "https://auth.autx.ai",
        cache_ttl_seconds: int = 3600,
    ) -> None:
        self.jwks_url = jwks_url
        self.issuer = issuer
        self.cache_ttl_seconds = cache_ttl_seconds
        self._cached_jwks: dict[str, Any] | None = None
        self._cached_at: float = 0

    async def _get_jwks(self) -> dict[str, Any]:
        now = time.time()
        if self._cached_jwks and (now - self._cached_at) < self.cache_ttl_seconds:
            return self._cached_jwks

        async with httpx.AsyncClient(timeout=10) as client:
            res = await client.get(self.jwks_url)
            res.raise_for_status()
            self._cached_jwks = res.json()
            self._cached_at = now
            return self._cached_jwks

    async def verify(self, token: str, audience: str) -> dict[str, Any]:
        jwks = await self._get_jwks()
        unverified = jwt.get_unverified_header(token)
        kid = unverified.get("kid")
        if not kid:
            raise JWTError("Missing kid")

        key = next((k for k in jwks.get("keys", []) if k.get("kid") == kid), None)
        if not key:
            raise JWTError("Unable to find key for kid")

        payload = jwt.decode(
            token,
            key,
            algorithms=["RS256"],
            audience=audience,
            issuer=self.issuer,
        )

        if payload.get("type") != "downstream":
            raise JWTError("Invalid token type")
        if not isinstance(payload.get("act"), dict):
            raise JWTError("Missing act claim")

        return payload


async def verify_autx_token(token: str, audience: str) -> dict[str, Any]:
    verifier = AutxTokenVerifier()
    return await verifier.verify(token=token, audience=audience)
