import time
import logging
from typing import Any

import httpx

from .exceptions import (
    LiteFoldAPIError,
    LiteFoldAuthError,
    LiteFoldNotFoundError,
    LiteFoldRateLimitError,
)

logger = logging.getLogger("litefold.http")

_DEFAULT_TIMEOUT = 60.0
_DEFAULT_MAX_RETRIES = 3
_DEFAULT_RETRY_BACKOFF = 1.0

# Status codes that are safe to retry
_RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503, 504})


def _raise_for_status(response: httpx.Response) -> None:
    if response.is_success:
        return
    try:
        detail = response.json().get("detail", response.text)
    except Exception:
        detail = response.text

    status = response.status_code
    if status in (401, 403):
        raise LiteFoldAuthError(status, detail)
    if status == 404:
        raise LiteFoldNotFoundError(status, detail)
    if status == 429:
        raise LiteFoldRateLimitError(status, detail)
    raise LiteFoldAPIError(status, detail)


_RETRYABLE_TRANSPORT_ERRORS = (
    httpx.ConnectError,
    httpx.ReadError,
    httpx.ReadTimeout,
    httpx.WriteTimeout,
    httpx.RemoteProtocolError,
    httpx.CloseError,
    httpx.PoolTimeout,
)


def _should_retry(exc: Exception) -> bool:
    if isinstance(exc, _RETRYABLE_TRANSPORT_ERRORS):
        return True
    if isinstance(exc, LiteFoldAPIError) and exc.status_code in _RETRYABLE_STATUS_CODES:
        return True
    return False


class HttpClient:
    def __init__(
        self,
        api_key: str,
        base_url: str,
        timeout: float = _DEFAULT_TIMEOUT,
        max_retries: int = _DEFAULT_MAX_RETRIES,
        retry_backoff: float = _DEFAULT_RETRY_BACKOFF,
        pool_max_connections: int = 200,
        pool_max_keepalive: int = 40,
    ):
        self._max_retries = max_retries
        self._retry_backoff = retry_backoff

        # Pool sized for 100+ agents sharing this client.
        # max_connections: total sockets.  max_keepalive_connections: idle kept warm.
        self._client = httpx.Client(
            base_url=base_url.rstrip("/"),
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=timeout,
            limits=httpx.Limits(
                max_connections=pool_max_connections,
                max_keepalive_connections=pool_max_keepalive,
                keepalive_expiry=30,
            ),
        )

    def _request_with_retry(
        self, method: str, path: str, *, base_url: str | None = None, **kwargs
    ) -> httpx.Response:
        url = f"{base_url.rstrip('/')}{path}" if base_url else path
        last_exc: Exception | None = None
        for attempt in range(self._max_retries + 1):
            response = None
            try:
                response = self._client.request(method, url, **kwargs)
                _raise_for_status(response)
                return response
            except Exception as exc:
                last_exc = exc
                if not _should_retry(exc) or attempt >= self._max_retries:
                    raise
                delay = self._retry_backoff * (2 ** attempt)
                # Respect Retry-After header for 429s
                if isinstance(exc, LiteFoldRateLimitError) and response is not None:
                    try:
                        delay = max(delay, float(response.headers.get("retry-after", delay)))
                    except (ValueError, TypeError):
                        pass
                logger.warning(
                    f"Retrying {method} {path} (attempt {attempt + 1}/{self._max_retries}) "
                    f"after {delay:.1f}s: {exc}"
                )
                time.sleep(delay)
        raise last_exc  # type: ignore[misc]

    def get(
        self, path: str, params: dict[str, Any] | None = None, base_url: str | None = None,
    ) -> Any:
        response = self._request_with_retry("GET", path, base_url=base_url, params=params)
        return response.json()

    def get_raw(
        self, path: str, params: dict[str, Any] | None = None, base_url: str | None = None,
    ) -> httpx.Response:
        return self._request_with_retry("GET", path, base_url=base_url, params=params)

    def post(
        self, path: str, json: dict[str, Any] | None = None, base_url: str | None = None,
    ) -> Any:
        response = self._request_with_retry("POST", path, base_url=base_url, json=json)
        return response.json()

    def post_multipart(
        self,
        path: str,
        data: dict[str, Any] | None = None,
        files: list[tuple[str, tuple[str, bytes, str]]] | None = None,
        base_url: str | None = None,
    ) -> Any:
        response = self._request_with_retry("POST", path, base_url=base_url, data=data, files=files)
        return response.json()

    def put(
        self, path: str, json: dict[str, Any] | None = None, base_url: str | None = None,
    ) -> Any:
        response = self._request_with_retry("PUT", path, base_url=base_url, json=json)
        return response.json()

    def patch(
        self, path: str, json: dict[str, Any] | None = None, base_url: str | None = None,
    ) -> Any:
        response = self._request_with_retry("PATCH", path, base_url=base_url, json=json)
        return response.json()

    def close(self) -> None:
        self._client.close()

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.close()
