from typing import Optional
from jose import jwt, JWTError
from analogai.foundation.common.logging.app_logger import app_logger
from analogai.foundation.extensions.authentication.api_token.api_token_payload import ApiTokenPayload


class ApiTokenAuthenticationService:
    def __init__(self, claim_keys: Optional[list[str]] = None):
        self.claim_keys = claim_keys or ["agent_id"]

    def _redact_token(self, token: str) -> str:
        if not token:
            return "<empty>"
        if len(token) <= 10:
            return f"{token[0]}***{token[-1]}"
        return f"{token[:6]}...{token[-4:]}"

    def extract_payload(self, token: str) -> ApiTokenPayload:
        if not token:
            app_logger.debug("API token authentication failed: missing token")
            raise ValueError("Missing auth token")

        try:
            payload = jwt.decode(
                token,
                key="dummy-key",
                options={
                    "verify_signature": False,
                    "verify_aud": False,
                    "verify_iss": False,
                    "verify_exp": False,
                },
            )
        except JWTError as error:
            app_logger.debug(
                "API token parsing failed for token %s: %s",
                self._redact_token(token),
                str(error),
            )
            raise ValueError("Invalid token format") from error

        agent_id = self._extract_agent_id(payload)
        if not agent_id:
            app_logger.debug(
                "API token missing agent_id claim for token %s",
                self._redact_token(token),
            )
            raise ValueError("No 'agent_id' claim in token")

        return ApiTokenPayload(
            token=token,
            agent_id=str(agent_id),
            subject=payload.get("sub"),
            client_id=payload.get("client_id") or payload.get("azp"),
            issuer=payload.get("iss"),
            audience=payload.get("aud"),
        )

    def _extract_agent_id(self, payload: dict) -> Optional[str]:
        for key in self.claim_keys:
            value = payload.get(key)
            if value:
                return value
        return None
