"""JWT token validation for hosted mode."""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Optional

try:
    import jwt
    HAS_JWT = True
except ImportError:
    HAS_JWT = False

from ..errors import AuthenticationError, ConfigurationError
from ..logging import get_logger

logger = get_logger(__name__)


@dataclass
class TokenClaims:
    """Decoded JWT token claims.
    
    Attributes:
        user_id: User identifier
        ad_account_id: Optional ad account ID
        expires_at: Token expiration timestamp
        issued_at: Token issuance timestamp
        issuer: Token issuer
        raw_claims: Raw claims dictionary
    """
    user_id: str
    ad_account_id: Optional[str]
    expires_at: datetime
    issued_at: datetime
    issuer: str
    raw_claims: dict


class JWTValidator:
    """Validates MCP service tokens from auth broker.
    
    Verifies JWT signatures, expiration, and issuer claims.
    """
    
    def __init__(
        self,
        public_key: str,
        issuer: str,
        audience: str | None = None,
        algorithm: str = "RS256",
        leeway: int = 0,
    ):
        """Initialize JWT validator.
        
        Args:
            public_key: Public key for verifying signatures
            issuer: Expected token issuer
            algorithm: JWT algorithm (default: RS256)
            leeway: Leeway in seconds for expiration validation
            
        Raises:
            ConfigurationError: If PyJWT is not installed or config is invalid
        """
        if not HAS_JWT:
            raise ConfigurationError(
                "PyJWT is required for JWT validation. "
                "Install with: pip install pyjwt[crypto]",
                setting="JWT_VALIDATION"
            )
        
        self.public_key = public_key
        self.issuer = issuer
        self.audience = audience
        self.algorithm = algorithm
        self.leeway = leeway
    
    async def validate(self, token: str) -> TokenClaims:
        """Validate a JWT token and extract claims.
        
        Args:
            token: JWT token string
            
        Returns:
            TokenClaims with decoded claims
            
        Raises:
            AuthenticationError: If token is invalid or expired
        """
        try:
            # Decode and verify token
            payload = jwt.decode(
                token,
                self.public_key,
                algorithms=[self.algorithm],
                issuer=self.issuer,
                audience=self.audience,
                options={
                    "verify_signature": True,
                    "verify_exp": True,
                    "verify_iss": True,
                    "verify_aud": bool(self.audience),
                },
                leeway=self.leeway,
            )
            
            # Extract required claims
            user_id = payload.get("sub")
            if not user_id:
                raise AuthenticationError(
                    "Token missing required 'sub' claim",
                    reason="missing_subject"
                )
            
            # Extract optional claims
            ad_account_id = payload.get("ad_account_id")
            
            # Extract timestamps
            exp = payload.get("exp")
            iat = payload.get("iat")
            
            if not exp:
                raise AuthenticationError(
                    "Token missing required 'exp' claim",
                    reason="missing_expiration"
                )
            
            # Create claims object
            claims = TokenClaims(
                user_id=user_id,
                ad_account_id=ad_account_id,
                expires_at=datetime.fromtimestamp(exp, tz=timezone.utc),
                issued_at=datetime.fromtimestamp(iat, tz=timezone.utc) if iat else datetime.now(timezone.utc),
                issuer=payload.get("iss", self.issuer),
                raw_claims=payload,
            )
            
            logger.info(
                "jwt_validated user_id=%s ad_account_id=%s expires_at=%s",
                user_id,
                ad_account_id,
                claims.expires_at.isoformat(),
            )
            
            return claims
            
        except jwt.ExpiredSignatureError as exc:
            logger.warning("jwt_expired error=%s", exc)
            raise AuthenticationError(
                "Token has expired",
                reason="expired_token"
            ) from exc
        
        except jwt.InvalidIssuerError as exc:
            logger.warning("jwt_invalid_issuer error=%s", exc)
            raise AuthenticationError(
                "Token has invalid issuer",
                reason="invalid_issuer"
            ) from exc
        
        except jwt.InvalidAudienceError as exc:
            logger.warning("jwt_invalid_audience error=%s", exc)
            raise AuthenticationError(
                "Token has invalid audience",
                reason="invalid_audience",
            ) from exc
        
        except jwt.InvalidSignatureError as exc:
            logger.warning("jwt_invalid_signature error=%s", exc)
            raise AuthenticationError(
                "Token has invalid signature",
                reason="invalid_signature"
            ) from exc
        
        except jwt.DecodeError as exc:
            logger.warning("jwt_decode_error error=%s", exc)
            raise AuthenticationError(
                "Token could not be decoded",
                reason="decode_error"
            ) from exc
        
        except Exception as exc:
            logger.error("jwt_validation_error type=%s error=%s", type(exc).__name__, exc)
            raise AuthenticationError(
                "Token validation failed",
                reason="validation_error"
            ) from exc

