import json
import logging
from typing import Optional
from httpx import Limits
from dataclasses import dataclass


from ppio_sandbox.core.api.client.client import AuthenticatedClient
from ppio_sandbox.core.connection_config import ConnectionConfig
from ppio_sandbox.core.api.metadata import default_headers
from ppio_sandbox.core.exceptions import (
    AuthenticationException,
    SandboxException,
    RateLimitException,
)
from ppio_sandbox.core.api.client.types import Response

logger = logging.getLogger(__name__)


@dataclass
class SandboxCreateResponse:
    sandbox_id: str
    sandbox_domain: Optional[str]
    envd_version: str
    envd_access_token: str


def handle_api_exception(e: Response):
    try:
        body = json.loads(e.content) if e.content else {}
    except json.JSONDecodeError:
        body = {}

    if e.status_code == 429:
        message = f"{e.status_code}: Rate limit exceeded, please try again later"
        if body.get("message"):
            message += f" - {body['message']}"

        return RateLimitException(message)

    if "message" in body:
        return SandboxException(f"{e.status_code}: {body['message']}")
    return SandboxException(f"{e.status_code}: {e.content}")


class ApiClient(AuthenticatedClient):
    """
    The client for interacting with the PPIO Agent Sandbox API.
    """

    def __init__(
        self,
        config: ConnectionConfig,
        require_api_key: bool = True,
        require_access_token: bool = False,
        limits: Optional[Limits] = None,
        *args,
        **kwargs,
    ):
        if require_api_key and require_access_token:
            raise AuthenticationException(
                "Only one of api_key or access_token can be required, not both",
            )

        if not require_api_key and not require_access_token:
            raise AuthenticationException(
                "Either api_key or access_token is required",
            )

        token = None
        if require_api_key:
            if config.api_key is None:
                raise AuthenticationException(
                    "API key is required, please visit the Key Management at https://ppio.com/settings/key-management to get your API key. "
                    "You can either set the environment variable `PPIO_API_KEY` "
                    'or you can pass it directly to the sandbox like Sandbox(api_key="sk_...")',
                )
            token = config.api_key

        if require_access_token:
            if config.access_token is None:
                raise AuthenticationException(
                    "Access token is required, please visit the Key Management tab at https://ppio.com/settings/key-management to get your access token (which is same as API key). "
                    "You can set the environment variable `PPIO_ACCESS_TOKEN` or pass the `access_token` in options.",
                )
            token = config.access_token

        auth_header_name = "X-API-KEY" if require_api_key else "Authorization"
        prefix = "" if require_api_key else "Bearer"

        headers = {
            **default_headers,
            **(config.headers or {}),
        }

        # Prevent passing these parameters twice
        kwargs.pop("headers", None)
        kwargs.pop("token", None)
        kwargs.pop("auth_header_name", None)
        kwargs.pop("prefix", None)

        super().__init__(
            base_url=config.api_url,
            httpx_args={
                "event_hooks": {
                    "request": [self._log_request],
                    "response": [self._log_response],
                },
                "proxy": config.proxy,
                "limits": limits,
            },
            headers=headers,
            token=token,
            auth_header_name=auth_header_name,
            prefix=prefix,
            *args,
            **kwargs,
        )

    def _log_request(self, request):
        logger.info(f"Request {request.method} {request.url}")

    def _log_response(self, response: Response):
        if response.status_code >= 400:
            logger.error(f"Response {response.status_code}")
        else:
            logger.info(f"Response {response.status_code}")


# We need to override the logging hooks for the async usage
class AsyncApiClient(ApiClient):
    async def _log_request(self, request):
        logger.info(f"Request {request.method} {request.url}")

    async def _log_response(self, response: Response):
        if response.status_code >= 400:
            logger.error(f"Response {response.status_code}")
        else:
            logger.info(f"Response {response.status_code}")
