"""HTTP/2-first client with an aiohttp fallback for dependency-light installs."""

from __future__ import annotations

import asyncio
import importlib.util
import logging
import os
from typing import Any, AsyncIterator, Optional

import aiohttp

try:
    import httpx
except ImportError:  # Manual plugin copies may not install new requirements.
    httpx = None


def _verbose_http_logging_enabled() -> bool:
    value = os.environ.get("O1KEY_VERBOSE_LOG", "")
    return value.strip().lower() not in ("", "0", "false", "no", "off")


# httpx emits one INFO line for every polling request. Those access logs drown
# out the batch lifecycle and add no information on successful requests. Keep
# warnings/errors, and allow the existing verbose switch to restore raw logs.
if not _verbose_http_logging_enabled():
    logging.getLogger("httpx").setLevel(logging.WARNING)
    logging.getLogger("httpcore").setLevel(logging.WARNING)


HTTPX_AVAILABLE = httpx is not None
HTTP_CLIENT_ERRORS = (aiohttp.ClientError,)
HTTP_STREAM_ERRORS = (aiohttp.ClientPayloadError,)
if HTTPX_AVAILABLE:
    HTTP_CLIENT_ERRORS += (httpx.HTTPError,)
    HTTP_STREAM_ERRORS += (httpx.StreamError,)


class ResponseBodyIntegrityError(OSError):
    """The response body ended early or disagreed with Content-Length."""


class ResponseTaskIdMismatchError(OSError):
    """A task query returned a different task ID than the one requested."""


def _response_header(response: Any, name: str) -> str:
    headers = getattr(response, "headers", None) or {}
    value = headers.get(name)
    if value is None:
        value = headers.get(name.lower())
    if value is None:
        for key, candidate in headers.items():
            if str(key).lower() == name.lower():
                value = candidate
                break
    return "" if value is None else str(value).strip()


def response_body_diagnostics(response: Any, received_bytes: int) -> dict[str, Any]:
    """Describe whether a fully-read response agrees with its declared length."""
    declared_text = _response_header(response, "Content-Length")
    content_encoding = _response_header(response, "Content-Encoding") or "identity"
    transfer_encoding = _response_header(response, "Transfer-Encoding") or "<none>"
    declared_bytes = None
    length_check = "not-declared"

    if declared_text:
        try:
            declared_bytes = int(declared_text)
        except ValueError:
            length_check = "invalid-header"
        else:
            if content_encoding.lower() not in ("", "identity"):
                # aiohttp/httpx expose decoded bytes while Content-Length can
                # describe the compressed wire representation.
                length_check = "skipped-compressed"
            elif declared_bytes == received_bytes:
                length_check = "match"
            else:
                length_check = "mismatch"

    return {
        "http_version": str(getattr(response, "http_version", "") or "HTTP"),
        "status": int(getattr(response, "status", 0) or 0),
        "content_length": declared_text or "<none>",
        "declared_bytes": declared_bytes,
        "received_bytes": int(received_bytes),
        "content_encoding": content_encoding,
        "transfer_encoding": transfer_encoding,
        "length_check": length_check,
    }


def format_response_body_diagnostics(diagnostics: dict[str, Any]) -> str:
    return (
        f"{diagnostics['http_version']} {diagnostics['status']}"
        f" | Content-Length={diagnostics['content_length']}"
        f" | received={diagnostics['received_bytes']}B"
        f" | Content-Encoding={diagnostics['content_encoding']}"
        f" | Transfer-Encoding={diagnostics['transfer_encoding']}"
        f" | length_check={diagnostics['length_check']}"
    )


async def read_response_body_with_diagnostics(
    response: Any,
    *,
    chunk_size: int = 64 * 1024,
) -> tuple[bytes, dict[str, Any]]:
    """Read a body while retaining the received byte count if the stream fails."""
    buffer = bytearray()
    try:
        content = getattr(response, "content", None)
        iter_chunked = getattr(content, "iter_chunked", None)
        if callable(iter_chunked):
            async for chunk in iter_chunked(chunk_size):
                buffer.extend(chunk)
        else:
            read = getattr(response, "read", None)
            if callable(read):
                buffer.extend(await read())
            else:
                text = getattr(response, "text", None)
                if not callable(text):
                    raise TypeError("response does not expose a readable body")
                buffer.extend((await text()).encode("utf-8"))
    except asyncio.CancelledError:
        raise
    except Exception as exc:
        diagnostics = response_body_diagnostics(response, len(buffer))
        raise ResponseBodyIntegrityError(
            "响应体读取提前中断 | "
            f"{format_response_body_diagnostics(diagnostics)}"
            f" | cause={type(exc).__name__}: {exc}"
        ) from exc

    body = bytes(buffer)
    diagnostics = response_body_diagnostics(response, len(body))
    if diagnostics["length_check"] == "mismatch":
        raise ResponseBodyIntegrityError(
            "响应体长度与服务端声明不一致，可能在传输中提前中断 | "
            f"{format_response_body_diagnostics(diagnostics)}"
        )
    return body, diagnostics


def response_task_id(payload: Any) -> Optional[str]:
    """Return an explicit task_id/taskId without mistaking result item IDs for it."""
    if not isinstance(payload, dict):
        return None
    sources = [payload]
    for key in ("data", "result", "task"):
        nested = payload.get(key)
        if isinstance(nested, dict):
            sources.append(nested)
    for source in sources:
        for key in ("task_id", "taskId"):
            value = source.get(key)
            if value is not None and str(value).strip():
                return str(value).strip()
    return None


def validate_response_task_id(payload: Any, expected_task_id: str) -> Optional[str]:
    actual_task_id = response_task_id(payload)
    if actual_task_id is not None and actual_task_id != str(expected_task_id):
        raise ResponseTaskIdMismatchError(
            f"任务查询响应 ID 不匹配 | requested_task_id={expected_task_id}"
            f" | response_task_id={actual_task_id}"
        )
    return actual_task_id


def http2_runtime_available() -> bool:
    return HTTPX_AVAILABLE and importlib.util.find_spec("h2") is not None


def create_timeout(
    total: float,
    *,
    connect: float,
    read: float,
    write: float,
    pool: float,
):
    """Return a timeout object understood by the selected HTTP backend."""
    if HTTPX_AVAILABLE:
        return httpx.Timeout(total, connect=connect, read=read, write=write, pool=pool)
    return aiohttp.ClientTimeout(
        total=total,
        connect=connect,
        sock_connect=connect,
        sock_read=read,
    )


class _ResponseContent:
    def __init__(self, response: httpx.Response):
        self._response = response

    async def iter_chunked(self, chunk_size: int) -> AsyncIterator[bytes]:
        async for chunk in self._response.aiter_bytes(chunk_size):
            yield chunk


class HttpResponse:
    def __init__(self, response: httpx.Response):
        self._response = response
        self.content = _ResponseContent(response)

    @property
    def status(self) -> int:
        return self._response.status_code

    @property
    def headers(self):
        return self._response.headers

    @property
    def raw_headers(self):
        return tuple(self._response.headers.raw)

    @property
    def http_version(self) -> str:
        return self._response.http_version

    async def text(self) -> str:
        await self._response.aread()
        return self._response.text


class _RequestContext:
    def __init__(self, context):
        self._context = context

    async def __aenter__(self) -> HttpResponse:
        response = await self._context.__aenter__()
        return HttpResponse(response)

    async def __aexit__(self, exc_type, exc, tb):
        return await self._context.__aexit__(exc_type, exc, tb)


class O1keyAsyncHttpClient:
    """HTTP/2-first client; falls back to httpx HTTP/1.1, then aiohttp."""

    def __init__(
        self,
        *,
        max_connections: int = 32,
        max_keepalive_connections: int = 16,
        keepalive_expiry: float = 30.0,
        http2: bool = True,
    ):
        self.http2_enabled = bool(http2 and http2_runtime_available())
        self.backend = "httpx" if HTTPX_AVAILABLE else "aiohttp"
        self._max_connections = max_connections
        self._keepalive_expiry = keepalive_expiry
        self._client = None
        if HTTPX_AVAILABLE:
            self._client = httpx.AsyncClient(
                http2=self.http2_enabled,
                verify=True,
                limits=httpx.Limits(
                    max_connections=max_connections,
                    max_keepalive_connections=max_keepalive_connections,
                    keepalive_expiry=keepalive_expiry,
                ),
                follow_redirects=False,
            )

    async def __aenter__(self):
        if self.backend == "aiohttp":
            connector = aiohttp.TCPConnector(
                limit=self._max_connections,
                limit_per_host=self._max_connections,
                keepalive_timeout=self._keepalive_expiry,
            )
            self._client = aiohttp.ClientSession(connector=connector, trust_env=True)
        await self._client.__aenter__()
        return self

    async def __aexit__(self, exc_type, exc, tb):
        return await self._client.__aexit__(exc_type, exc, tb)

    def post(
        self,
        url: str,
        *,
        headers: Optional[dict[str, str]] = None,
        data: Any = None,
        files: Any = None,
        timeout: Any = None,
    ) -> _RequestContext:
        if self._client is None:
            raise RuntimeError("HTTP client must be entered before use")
        if self.backend == "httpx":
            return _RequestContext(self._client.stream(
                "POST",
                url,
                headers=headers,
                content=data if files is None else None,
                files=files,
                timeout=timeout,
            ))

        payload = data
        if files:
            payload = aiohttp.FormData()
            for field_name, (filename, source, content_type) in files.items():
                payload.add_field(
                    field_name,
                    source,
                    filename=filename,
                    content_type=content_type,
                )
        return self._client.post(
            url,
            headers=headers,
            data=payload,
            timeout=timeout,
        )

    def get(
        self,
        url: str,
        *,
        headers: Optional[dict[str, str]] = None,
        allow_redirects: bool = False,
        timeout: Any = None,
    ) -> _RequestContext:
        if self._client is None:
            raise RuntimeError("HTTP client must be entered before use")
        if self.backend == "httpx":
            return _RequestContext(self._client.stream(
                "GET",
                url,
                headers=headers,
                follow_redirects=allow_redirects,
                timeout=timeout,
            ))
        return self._client.get(
            url,
            headers=headers,
            allow_redirects=allow_redirects,
            timeout=timeout,
        )


def create_http_client(**kwargs) -> O1keyAsyncHttpClient:
    return O1keyAsyncHttpClient(**kwargs)


__all__ = [
    "HTTPX_AVAILABLE",
    "HTTP_CLIENT_ERRORS",
    "HTTP_STREAM_ERRORS",
    "O1keyAsyncHttpClient",
    "ResponseBodyIntegrityError",
    "ResponseTaskIdMismatchError",
    "create_http_client",
    "create_timeout",
    "format_response_body_diagnostics",
    "http2_runtime_available",
    "read_response_body_with_diagnostics",
    "response_body_diagnostics",
    "response_task_id",
    "validate_response_task_id",
]