from typing import Any
from uuid import UUID, uuid4

from archipy.configs.base_config import BaseConfig
from archipy.configs.config_template import AuthConfig
from archipy.helpers.utils.datetime_utils import DatetimeUtils
from archipy.models.errors import InvalidTokenError, TokenExpiredError


class JWTUtils:
    """Utility class for JWT token operations with enhanced security and datetime handling."""

    @classmethod
    def create_token(
        cls,
        data: dict[str, Any],
        expires_in: int,
        additional_claims: dict[str, Any] | None = None,
        auth_config: AuthConfig | None = None,
    ) -> str:
        """Creates a JWT token with enhanced security features.

        Args:
            data (dict[str, Any]): Base claims data to include in the token.
            expires_in (int): Token expiration time in seconds.
            additional_claims (dict[str, Any] | None): Optional additional claims to include in the token.
            auth_config (AuthConfig | None): Optional auth configuration override. If not provided, uses the global config.

        Returns:
            str: The encoded JWT token.

        Raises:
            ValueError: If data is empty or expiration is invalid
        """
        import jwt

        configs = BaseConfig.global_config().AUTH if auth_config is None else auth_config
        current_time = DatetimeUtils.get_datetime_utc_now()

        if not data:
            raise ValueError("Token data cannot be empty")
        if expires_in <= 0:
            raise ValueError("Token expiration must be positive")

        to_encode = data.copy()
        expire = DatetimeUtils.get_datetime_after_given_datetime_or_now(seconds=expires_in, datetime_given=current_time)

        # Add standard claims
        to_encode.update(
            {
                # Registered claims (RFC 7519)
                "iss": configs.JWT_ISSUER,
                "aud": configs.JWT_AUDIENCE,
                "exp": expire,
                "iat": current_time,
                "nbf": current_time,
            },
        )

        # Add JWT ID if enabled
        if configs.ENABLE_JTI_CLAIM:
            to_encode["jti"] = str(uuid4())

        # Add additional claims
        if additional_claims:
            to_encode.update(additional_claims)

        return jwt.encode(to_encode, configs.SECRET_KEY.get_secret_value(), algorithm=configs.HASH_ALGORITHM)

    @classmethod
    def create_access_token(
        cls,
        user_uuid: UUID,
        additional_claims: dict[str, Any] | None = None,
        auth_config: AuthConfig | None = None,
    ) -> str:
        """Creates an access token for a user.

        Args:
            user_uuid (UUID): The user's UUID to include in the token.
            additional_claims (dict[str, Any] | None): Optional additional claims to include in the token.
            auth_config (AuthConfig | None): Optional auth configuration override. If not provided, uses the global config.

        Returns:
            str: The encoded access token.
        """
        configs = BaseConfig.global_config().AUTH if auth_config is None else auth_config

        return cls.create_token(
            data={
                "sub": str(user_uuid),
                "type": "access",
                "token_version": configs.TOKEN_VERSION,
            },
            expires_in=configs.ACCESS_TOKEN_EXPIRES_IN,
            additional_claims=additional_claims,
            auth_config=configs,
        )

    @classmethod
    def create_refresh_token(
        cls,
        user_uuid: UUID,
        additional_claims: dict[str, Any] | None = None,
        auth_config: AuthConfig | None = None,
    ) -> str:
        """Creates a refresh token for a user.

        Args:
            user_uuid (UUID): The user's UUID to include in the token.
            additional_claims (dict[str, Any] | None): Optional additional claims to include in the token.
            auth_config (AuthConfig | None): Optional auth configuration override. If not provided, uses the global config.

        Returns:
            str: The encoded refresh token.
        """
        configs = BaseConfig.global_config().AUTH if auth_config is None else auth_config

        return cls.create_token(
            data={
                "sub": str(user_uuid),
                "type": "refresh",
                "token_version": configs.TOKEN_VERSION,
            },
            expires_in=configs.REFRESH_TOKEN_EXPIRES_IN,
            additional_claims=additional_claims,
            auth_config=configs,
        )

    @classmethod
    def decode_token(
        cls,
        token: str,
        verify_type: str | None = None,
        auth_config: AuthConfig | None = None,
    ) -> dict[str, Any]:
        """Decodes and verifies a JWT token with enhanced security checks.

        Args:
            token (str): The JWT token to decode.
            verify_type (str | None): Optional token type to verify (e.g., "access" or "refresh").
            auth_config (AuthConfig | None): Optional auth configuration override. If not provided, uses the global config.

        Returns:
            dict[str, Any]: The decoded token payload.

        Raises:
            TokenExpiredError: If the token has expired.
            InvalidTokenError: If the token is invalid (e.g., invalid signature, audience, issuer, or type).
        """
        import jwt
        from jwt.exceptions import (
            ExpiredSignatureError,
            InvalidAudienceError,
            InvalidIssuerError,
            InvalidSignatureError,
            InvalidTokenError as JWTInvalidTokenError,
        )

        configs = BaseConfig.global_config().AUTH if auth_config is None else auth_config
        required_claims = ["exp", "iat", "nbf", "aud", "iss", "sub", "type", "token_version"]
        if configs.ENABLE_JTI_CLAIM:
            required_claims.append("jti")

        try:
            payload = jwt.decode(
                token,
                configs.SECRET_KEY.get_secret_value(),
                algorithms=[configs.HASH_ALGORITHM],
                options={
                    "verify_signature": True,
                    "verify_exp": True,
                    "verify_nbf": True,
                    "verify_iat": True,
                    "verify_aud": True,
                    "verify_iss": True,
                    "require": required_claims,
                },
                audience=configs.JWT_AUDIENCE,
                issuer=configs.JWT_ISSUER,
            )

            # Verify token type
            if verify_type and payload.get("type") != verify_type:
                raise InvalidTokenError(f"Invalid token type. Expected {verify_type}")

            # Verify token version
            if payload.get("token_version") != configs.TOKEN_VERSION:
                raise InvalidTokenError("Token version is outdated")

            return payload

        except ExpiredSignatureError as exception:
            raise TokenExpiredError("Token has expired") from exception
        except InvalidSignatureError as exception:
            raise InvalidTokenError("Token signature is invalid") from exception
        except InvalidAudienceError as exception:
            raise InvalidTokenError("Token has invalid audience") from exception
        except InvalidIssuerError as exception:
            raise InvalidTokenError("Token has invalid issuer") from exception
        except JWTInvalidTokenError as exception:
            raise InvalidTokenError(f"Invalid token: {exception!s}") from exception

    @classmethod
    def verify_access_token(cls, token: str, auth_config: AuthConfig | None = None) -> dict[str, Any]:
        """Verifies an access token.

        Args:
            token (str): The access token to verify.
            auth_config (AuthConfig | None): Optional auth configuration override. If not provided, uses the global config.

        Returns:
            dict[str, Any]: The decoded access token payload.

        Raises:
            InvalidTokenException: If the token is invalid or not an access token.
            TokenExpiredException: If the token has expired.
        """
        configs = BaseConfig.global_config().AUTH if auth_config is None else auth_config
        return cls.decode_token(token, verify_type="access", auth_config=configs)

    @classmethod
    def verify_refresh_token(cls, token: str, auth_config: AuthConfig | None = None) -> dict[str, Any]:
        """Verifies a refresh token.

        Args:
            token (str): The refresh token to verify.
            auth_config (AuthConfig | None): Optional auth configuration override. If not provided, uses the global config.

        Returns:
            dict[str, Any]: The decoded refresh token payload.

        Raises:
            InvalidTokenException: If the token is invalid or not a refresh token.
            TokenExpiredException: If the token has expired.
        """
        configs = BaseConfig.global_config().AUTH if auth_config is None else auth_config
        return cls.decode_token(token, verify_type="refresh", auth_config=configs)

    @staticmethod
    def extract_user_uuid(payload: dict[str, Any]) -> UUID:
        """Extracts the user UUID from the token payload.

        Args:
            payload (dict[str, Any]): The decoded token payload.

        Returns:
            UUID: The user's UUID.

        Raises:
            InvalidTokenException: If the user identifier is invalid or missing.
        """
        try:
            return UUID(payload["sub"])
        except (KeyError, ValueError) as exception:
            raise InvalidTokenError("Invalid or missing user identifier in token") from exception

    @classmethod
    def get_token_expiry(cls, token: str, auth_config: AuthConfig | None = None) -> int:
        """Gets the token expiry timestamp.

        Args:
            token (str): The JWT token.
            auth_config (AuthConfig | None): Optional auth configuration override. If not provided, uses the global config.

        Returns:
            int: The token expiry timestamp in seconds.

        Raises:
            InvalidTokenException: If the token is invalid.
        """
        payload = cls.decode_token(token, auth_config=auth_config)
        return int(payload["exp"])
