import base64
import os
from typing import Dict, Optional, Tuple, cast

import httpx

from litellm._logging import verbose_proxy_logger
from litellm.llms.custom_httpx.http_handler import (
    get_async_httpx_client,
    httpxSpecialProvider,
)
from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth


class Oauth2Handler:
    """
    Handles OAuth2 token validation.
    """

    @staticmethod
    def _is_introspection_endpoint(
        token_info_endpoint: str,
        oauth_client_id: Optional[str],
        oauth_client_secret: Optional[str],
    ) -> bool:
        """
        Determine if this is an introspection endpoint (requires POST) or token info endpoint (uses GET).

        Args:
            token_info_endpoint: The OAuth2 endpoint URL
            oauth_client_id: OAuth2 client ID
            oauth_client_secret: OAuth2 client secret

        Returns:
            bool: True if this is an introspection endpoint
        """
        return (
            "introspect" in token_info_endpoint.lower()
            and oauth_client_id is not None
            and oauth_client_secret is not None
        )

    @staticmethod
    def _prepare_introspection_request(
        token: str,
        oauth_client_id: Optional[str],
        oauth_client_secret: Optional[str],
    ) -> Tuple[Dict[str, str], Dict[str, str]]:
        """
        Prepare headers and data for OAuth2 introspection endpoint (RFC 7662).

        Args:
            token: The OAuth2 token to validate
            oauth_client_id: OAuth2 client ID
            oauth_client_secret: OAuth2 client secret

        Returns:
            Tuple of (headers, data) for the introspection request
        """
        headers = {"Content-Type": "application/x-www-form-urlencoded"}
        data = {"token": token}

        # Add client authentication if credentials are provided
        if oauth_client_id and oauth_client_secret:
            # Use HTTP Basic authentication for client credentials
            credentials = base64.b64encode(
                f"{oauth_client_id}:{oauth_client_secret}".encode()
            ).decode()
            headers["Authorization"] = f"Basic {credentials}"
        elif oauth_client_id:
            # For public clients, include client_id in the request body
            data["client_id"] = oauth_client_id

        return headers, data

    @staticmethod
    def _prepare_token_info_request(token: str) -> Dict[str, str]:
        """
        Prepare headers for generic token info endpoint.

        Args:
            token: The OAuth2 token to validate

        Returns:
            Dict of headers for the token info request
        """
        return {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}

    @staticmethod
    def _extract_user_info(
        response_data: Dict,
        user_id_field_name: str,
        user_role_field_name: str,
        user_team_id_field_name: str,
    ) -> Tuple[Optional[str], Optional[str], Optional[str]]:
        """
        Extract user information from OAuth2 response.

        Args:
            response_data: The response data from OAuth2 endpoint
            user_id_field_name: Field name for user ID
            user_role_field_name: Field name for user role
            user_team_id_field_name: Field name for team ID

        Returns:
            Tuple of (user_id, user_role, user_team_id)
        """
        user_id = response_data.get(user_id_field_name)
        user_team_id = response_data.get(user_team_id_field_name)
        user_role = response_data.get(user_role_field_name)

        return user_id, user_role, user_team_id

    @staticmethod
    async def check_oauth2_token(token: str) -> UserAPIKeyAuth:
        """
        Makes a request to the token introspection endpoint to validate the OAuth2 token.

        This function implements OAuth2 token introspection according to RFC 7662.
        It supports both generic token info endpoints (GET) and OAuth2 introspection endpoints (POST).

        Args:
            token (str): The OAuth2 token to validate.

        Returns:
            UserAPIKeyAuth: If the token is valid, containing user information.

        Raises:
            ValueError: If the token is invalid, the request fails, or the token info endpoint is not set.
        """
        from litellm.proxy.proxy_server import premium_user

        if premium_user is not True:
            raise ValueError(
                "Oauth2 token validation is only available for premium users"
                + CommonProxyErrors.not_premium_user.value
            )

        verbose_proxy_logger.debug(
            "Oauth2 token validation for token=[set=%s]", token is not None
        )

        # Get the token info endpoint from environment variable
        token_info_endpoint = os.getenv("OAUTH_TOKEN_INFO_ENDPOINT")
        user_id_field_name = os.environ.get("OAUTH_USER_ID_FIELD_NAME", "sub")
        user_role_field_name = os.environ.get("OAUTH_USER_ROLE_FIELD_NAME", "role")
        user_team_id_field_name = os.environ.get(
            "OAUTH_USER_TEAM_ID_FIELD_NAME", "team_id"
        )

        # OAuth2 client credentials for introspection endpoint authentication
        oauth_client_id = os.environ.get("OAUTH_CLIENT_ID")
        oauth_client_secret = os.environ.get("OAUTH_CLIENT_SECRET")

        if not token_info_endpoint:
            raise ValueError(
                "OAUTH_TOKEN_INFO_ENDPOINT environment variable is not set"
            )

        client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)

        # Determine if this is an introspection endpoint (requires POST) or token info endpoint (uses GET)
        is_introspection_endpoint = Oauth2Handler._is_introspection_endpoint(
            token_info_endpoint=token_info_endpoint,
            oauth_client_id=oauth_client_id,
            oauth_client_secret=oauth_client_secret,
        )

        try:
            if is_introspection_endpoint:
                # OAuth2 Token Introspection (RFC 7662) - requires POST with form data
                verbose_proxy_logger.debug("Using OAuth2 introspection endpoint (POST)")

                headers, data = Oauth2Handler._prepare_introspection_request(
                    token=token,
                    oauth_client_id=oauth_client_id,
                    oauth_client_secret=oauth_client_secret,
                )

                response = await client.post(
                    token_info_endpoint, headers=headers, data=data
                )
            else:
                # Generic token info endpoint - uses GET with Bearer token
                verbose_proxy_logger.debug("Using generic token info endpoint (GET)")
                headers = Oauth2Handler._prepare_token_info_request(token=token)
                response = await client.get(token_info_endpoint, headers=headers)

            # if it's a bad token we expect it to raise an HTTPStatusError
            response.raise_for_status()

            # If we get here, the request was successful
            data = response.json()

            verbose_proxy_logger.debug(
                "Oauth2 token validation for token=%s, response from endpoint=%s",
                token,
                data,
            )

            # For introspection endpoints, check if token is active
            if is_introspection_endpoint and not data.get("active", True):
                raise ValueError("Token is not active")

            # Extract user information from response
            user_id, user_role, user_team_id = Oauth2Handler._extract_user_info(
                response_data=data,
                user_id_field_name=user_id_field_name,
                user_role_field_name=user_role_field_name,
                user_team_id_field_name=user_team_id_field_name,
            )

            return UserAPIKeyAuth(
                api_key=token,
                team_id=user_team_id,
                user_id=user_id,
                user_role=cast(LitellmUserRoles, user_role),
            )
        except httpx.HTTPStatusError as e:
            # This will catch any 4xx or 5xx errors
            raise ValueError(f"Oauth 2.0 Token validation failed: {e}")
        except Exception as e:
            # This will catch any other errors (like network issues)
            raise ValueError(f"An error occurred during token validation: {e}")
