"""JWT token management."""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any

import jwt
import structlog

from curadise_agent.errors import AuthenticationError

log = structlog.get_logger(__name__)


@dataclass
class TokenInfo:
    """Information extracted from a JWT token."""

    subject: str | None
    issuer: str | None
    audience: str | list[str] | None
    issued_at: datetime | None
    expires_at: datetime | None
    claims: dict[str, Any]

    @property
    def is_expired(self) -> bool:
        """Check if token is expired."""
        if self.expires_at is None:
            return False
        return datetime.now(UTC) >= self.expires_at

    @property
    def expires_in_seconds(self) -> float | None:
        """Get seconds until expiration."""
        if self.expires_at is None:
            return None
        delta = self.expires_at - datetime.now(UTC)
        return max(0.0, delta.total_seconds())

    @property
    def needs_refresh(self) -> bool:
        """Check if token should be refreshed (within 5 minutes of expiry)."""
        expires_in = self.expires_in_seconds
        if expires_in is None:
            return False
        return expires_in < 300  # 5 minutes


@dataclass
class TokenManager:
    """
    Manages JWT tokens for agent authentication.

    Handles token parsing, validation, refresh scheduling,
    and secure storage.
    """

    refresh_threshold_seconds: int = 300  # Refresh 5 minutes before expiry
    _access_token: str | None = field(default=None, init=False)
    _refresh_token: str | None = field(default=None, init=False)
    _token_info: TokenInfo | None = field(default=None, init=False)

    @property
    def access_token(self) -> str | None:
        """Get current access token."""
        return self._access_token

    @property
    def refresh_token(self) -> str | None:
        """Get current refresh token."""
        return self._refresh_token

    @property
    def token_info(self) -> TokenInfo | None:
        """Get parsed token information."""
        return self._token_info

    @property
    def is_authenticated(self) -> bool:
        """Check if we have a valid access token."""
        return self._access_token is not None and not self.is_expired

    @property
    def is_expired(self) -> bool:
        """Check if current token is expired."""
        if self._token_info is None:
            return True
        return self._token_info.is_expired

    @property
    def needs_refresh(self) -> bool:
        """Check if token needs refreshing."""
        if self._token_info is None:
            return True
        if self._refresh_token is None:
            return False
        return self._token_info.needs_refresh

    def set_tokens(
        self,
        access_token: str,
        refresh_token: str | None = None,
    ) -> None:
        """
        Set new tokens.

        Args:
            access_token: JWT access token
            refresh_token: Optional refresh token
        """
        self._access_token = access_token
        self._refresh_token = refresh_token
        self._token_info = self.parse_token(access_token)

        log.info(
            "tokens_set",
            expires_in=self._token_info.expires_in_seconds if self._token_info else None,
        )

    def clear_tokens(self) -> None:
        """Clear all tokens."""
        self._access_token = None
        self._refresh_token = None
        self._token_info = None
        log.info("tokens_cleared")

    @staticmethod
    def parse_token(token: str, verify: bool = False) -> TokenInfo:
        """
        Parse a JWT token without verification.

        Args:
            token: JWT token string
            verify: Whether to verify signature (requires key)

        Returns:
            TokenInfo with extracted claims

        Raises:
            AuthenticationError: If token is malformed
        """
        try:
            # Decode without verification to extract claims
            claims = jwt.decode(
                token,
                options={
                    "verify_signature": verify,
                    "verify_exp": False,
                    "verify_iat": False,
                },
            )

            # Extract standard claims
            subject = claims.get("sub")
            issuer = claims.get("iss")
            audience = claims.get("aud")

            # Parse timestamps
            issued_at = None
            if "iat" in claims:
                issued_at = datetime.fromtimestamp(claims["iat"], tz=UTC)

            expires_at = None
            if "exp" in claims:
                expires_at = datetime.fromtimestamp(claims["exp"], tz=UTC)

            return TokenInfo(
                subject=subject,
                issuer=issuer,
                audience=audience,
                issued_at=issued_at,
                expires_at=expires_at,
                claims=claims,
            )

        except jwt.DecodeError as e:
            raise AuthenticationError(f"Invalid token format: {e}", cause=e) from e

    def get_authorization_header(self) -> dict[str, str]:
        """
        Get authorization header for requests.

        Returns:
            Header dict with Bearer token

        Raises:
            AuthenticationError: If no valid token
        """
        if not self._access_token:
            raise AuthenticationError("No access token available")

        return {"Authorization": f"Bearer {self._access_token}"}

    def get_expiry_timestamp(self) -> float | None:
        """
        Get token expiry as Unix timestamp.

        Returns:
            Expiry timestamp or None if no token
        """
        if self._token_info is None or self._token_info.expires_at is None:
            return None
        return self._token_info.expires_at.timestamp()

    def schedule_refresh_at(self) -> float | None:
        """
        Calculate when to schedule token refresh.

        Returns:
            Unix timestamp for refresh, or None if no refresh needed
        """
        if self._refresh_token is None:
            return None

        expiry = self.get_expiry_timestamp()
        if expiry is None:
            return None

        # Schedule refresh before expiry
        refresh_at = expiry - self.refresh_threshold_seconds
        now = time.time()

        if refresh_at <= now:
            # Already past refresh time, refresh immediately
            return now

        return refresh_at
