# retoor <retoor@molodetz.nl>

from __future__ import annotations

from collections.abc import AsyncIterator

import httpx
from curl_cffi import CurlHttpVersion
from curl_cffi.requests import AsyncSession
from curl_cffi.requests.exceptions import RequestException, Timeout

IMPERSONATE_TARGET: str = "chrome146"

DEFAULT_TIMEOUT_SECONDS: float = 30.0

STRIP_REQUEST_HEADERS: frozenset[str] = frozenset(
    {
        "host",
        "connection",
        "proxy-connection",
        "content-length",
        "transfer-encoding",
        "user-agent",
        "accept-encoding",
    }
)

STRIP_RESPONSE_HEADERS: frozenset[str] = frozenset(
    {
        "content-encoding",
        "content-length",
        "transfer-encoding",
        "connection",
    }
)

HTTP_VERSION_LABELS: dict[int, bytes] = {
    int(CurlHttpVersion.V1_0): b"HTTP/1.0",
    int(CurlHttpVersion.V1_1): b"HTTP/1.1",
    int(CurlHttpVersion.V2_0): b"HTTP/2",
    int(CurlHttpVersion.V2TLS): b"HTTP/2",
    int(CurlHttpVersion.V2_PRIOR_KNOWLEDGE): b"HTTP/2",
    int(CurlHttpVersion.V3): b"HTTP/3",
    int(CurlHttpVersion.V3ONLY): b"HTTP/3",
}


def http_version_for(url: httpx.URL):
    if url.scheme == "http":
        return CurlHttpVersion.V1_1
    return None


def resolve_timeout(request: httpx.Request) -> float:
    extension = request.extensions.get("timeout") or {}
    for key in ("read", "connect", "pool"):
        value = extension.get(key)
        if isinstance(value, (int, float)):
            return float(value)
    return DEFAULT_TIMEOUT_SECONDS


class CurlResponseStream(httpx.AsyncByteStream):
    def __init__(self, response: object) -> None:
        self._response = response

    async def __aiter__(self) -> AsyncIterator[bytes]:
        async for chunk in self._response.aiter_content():
            yield chunk

    async def aclose(self) -> None:
        await self._response.aclose()


class CurlTransport(httpx.AsyncBaseTransport):
    def __init__(
        self,
        *,
        impersonate: str = IMPERSONATE_TARGET,
        verify: bool = True,
        proxy: str | None = None,
    ) -> None:
        self._session = AsyncSession()
        self._impersonate = impersonate
        self._verify = verify
        self._proxy = proxy

    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
        headers = {
            key: value
            for key, value in request.headers.items()
            if key.lower() not in STRIP_REQUEST_HEADERS
        }
        body = await request.aread()
        extra = {}
        version = http_version_for(request.url)
        if version is not None:
            extra["http_version"] = version
        try:
            response = await self._session.request(
                request.method,
                str(request.url),
                headers=headers,
                data=body or None,
                impersonate=self._impersonate,
                verify=self._verify,
                proxy=self._proxy,
                stream=True,
                allow_redirects=False,
                timeout=resolve_timeout(request),
                **extra,
            )
        except Timeout as exc:
            raise httpx.ConnectTimeout(str(exc), request=request) from exc
        except RequestException as exc:
            raise httpx.ConnectError(str(exc), request=request) from exc

        response_headers = [
            (key, value)
            for key, value in response.headers.items()
            if key.lower() not in STRIP_RESPONSE_HEADERS
        ]
        http_version = HTTP_VERSION_LABELS.get(int(response.http_version), b"HTTP/2")
        return httpx.Response(
            status_code=response.status_code,
            headers=response_headers,
            stream=CurlResponseStream(response),
            extensions={"http_version": http_version},
            request=request,
        )

    async def aclose(self) -> None:
        await self._session.close()