from __future__ import annotations

from typing import Any

import httpx

from ._async_run import AsyncRun, AsyncRuns
from ._base_client import BaseClient
from ._constants import DEFAULT_MAX_RETRIES, DEFAULT_TIMEOUT
from ._exceptions import APIConnectionError, APITimeoutError
from ._response import RungateResponse
from ._run import Run, Runs
from .resources.chat import AsyncChat, Chat


class Rungate(BaseClient):
    """Synchronous Rungate client.

    Usage:
        client = Rungate(api_key="rg_sk_...")
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": "Hello"}],
        )
        print(response.rungate.cost_usd)

    Or as a context manager:
        with Rungate(api_key="rg_sk_...") as client:
            response = client.chat.completions.create(...)
    """

    def __init__(
        self,
        api_key: str | None = None,
        base_url: str | None = None,
        timeout: float = DEFAULT_TIMEOUT,
        max_retries: int = DEFAULT_MAX_RETRIES,
    ) -> None:
        super().__init__(api_key=api_key, base_url=base_url, timeout=timeout, max_retries=max_retries)
        self._http_client = httpx.Client(timeout=self.timeout)
        self.chat = Chat(self)
        self.runs = Runs(self)

    def run(
        self,
        *,
        name: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> Run:
        """Create a run context manager.

        Usage:
            with rg.run(name="research") as run:
                response = run.chat.completions.create(model="gpt-4o", messages=[...])
                print(run.budget.remaining_usd)
        """
        return self.runs.create(name=name, metadata=metadata)

    def _request(
        self,
        method: str,
        path: str,
        *,
        json: dict[str, Any] | None = None,
        headers: dict[str, str] | None = None,
    ) -> RungateResponse:
        """Make a synchronous HTTP request with connection retry logic.

        Retries are for connection failures only (ConnectError, TimeoutException).
        HTTP status errors are NOT retried — Rungate handles provider-level retries
        server-side.
        """
        url = self._build_url(path)
        request_headers = self._build_headers(headers)

        last_error: Exception | None = None
        for attempt in range(self.max_retries + 1):
            try:
                http_response = self._http_client.request(
                    method,
                    url,
                    json=json,
                    headers=request_headers,
                )
                response = RungateResponse(http_response)
                return self._handle_response(response)
            except httpx.TimeoutException as e:
                last_error = e
                if attempt >= self.max_retries:
                    raise APITimeoutError(f"Request timed out after {self.timeout}s: {e}") from e
            except httpx.ConnectError as e:
                last_error = e
                if attempt >= self.max_retries:
                    raise APIConnectionError(f"Failed to connect to {self.base_url}: {e}") from e
            except httpx.RequestError as e:
                # Catches all other httpx request errors (protocol errors, etc.)
                last_error = e
                if attempt >= self.max_retries:
                    raise APIConnectionError(f"Request failed: {e}") from e

        raise APIConnectionError(f"Request failed after {self.max_retries + 1} attempts: {last_error}")

    def _request_stream(
        self,
        method: str,
        path: str,
        *,
        json: dict[str, Any] | None = None,
        headers: dict[str, str] | None = None,
    ) -> httpx.Response:
        """Make a streaming request. Returns httpx.Response with body not yet read.

        The caller is responsible for closing the response.

        If the server returns a non-streaming error (JSON), this method reads the
        body, detects the error, and raises the appropriate exception.
        """
        url = self._build_url(path)
        request_headers = self._build_headers(headers)

        request = self._http_client.build_request(method, url, json=json, headers=request_headers)
        response = self._http_client.send(request, stream=True)

        # Check for errors or non-SSE responses
        content_type = response.headers.get("content-type", "")
        media_type = content_type.split(";")[0].strip().lower()

        if response.status_code >= 400:
            try:
                response.read()
                rr = RungateResponse(response)
                self._handle_response(rr)
            except Exception:
                response.close()
                raise
        elif media_type != "text/event-stream":
            try:
                response.read()
            except Exception:
                response.close()
                raise
            from ._exceptions import APIError

            raise APIError(
                message=f"Expected streaming response (text/event-stream), got {content_type or 'unknown'}",
                status_code=response.status_code,
                body=RungateResponse(response).body,
            )

        return response

    def close(self) -> None:
        """Close the underlying HTTP client and release connection pool."""
        self._http_client.close()

    def __enter__(self) -> Rungate:
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()


class AsyncRungate(BaseClient):
    """Asynchronous Rungate client.

    Usage:
        client = AsyncRungate(api_key="rg_sk_...")
        response = await client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": "Hello"}],
        )

    Or as an async context manager:
        async with AsyncRungate(api_key="rg_sk_...") as client:
            response = await client.chat.completions.create(...)
    """

    def __init__(
        self,
        api_key: str | None = None,
        base_url: str | None = None,
        timeout: float = DEFAULT_TIMEOUT,
        max_retries: int = DEFAULT_MAX_RETRIES,
    ) -> None:
        super().__init__(api_key=api_key, base_url=base_url, timeout=timeout, max_retries=max_retries)
        self._http_client = httpx.AsyncClient(timeout=self.timeout)
        self.chat = AsyncChat(self)
        self.runs = AsyncRuns(self)

    def run(
        self,
        *,
        name: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> AsyncRun:
        """Create an async run context manager."""
        return self.runs.create(name=name, metadata=metadata)

    async def _request(
        self,
        method: str,
        path: str,
        *,
        json: dict[str, Any] | None = None,
        headers: dict[str, str] | None = None,
    ) -> RungateResponse:
        """Make an async HTTP request with connection retry logic."""
        url = self._build_url(path)
        request_headers = self._build_headers(headers)

        last_error: Exception | None = None
        for attempt in range(self.max_retries + 1):
            try:
                http_response = await self._http_client.request(
                    method,
                    url,
                    json=json,
                    headers=request_headers,
                )
                response = RungateResponse(http_response)
                return self._handle_response(response)
            except httpx.TimeoutException as e:
                last_error = e
                if attempt >= self.max_retries:
                    raise APITimeoutError(f"Request timed out after {self.timeout}s: {e}") from e
            except httpx.ConnectError as e:
                last_error = e
                if attempt >= self.max_retries:
                    raise APIConnectionError(f"Failed to connect to {self.base_url}: {e}") from e
            except httpx.RequestError as e:
                last_error = e
                if attempt >= self.max_retries:
                    raise APIConnectionError(f"Request failed: {e}") from e

        raise APIConnectionError(f"Request failed after {self.max_retries + 1} attempts: {last_error}")

    async def _request_stream(
        self,
        method: str,
        path: str,
        *,
        json: dict[str, Any] | None = None,
        headers: dict[str, str] | None = None,
    ) -> httpx.Response:
        """Make an async streaming request. Returns httpx.Response with body not yet read."""
        url = self._build_url(path)
        request_headers = self._build_headers(headers)

        request = self._http_client.build_request(method, url, json=json, headers=request_headers)
        response = await self._http_client.send(request, stream=True)

        content_type = response.headers.get("content-type", "")
        media_type = content_type.split(";")[0].strip().lower()

        if response.status_code >= 400:
            try:
                await response.aread()
                rr = RungateResponse(response)
                self._handle_response(rr)
            except Exception:
                await response.aclose()
                raise
        elif media_type != "text/event-stream":
            try:
                await response.aread()
            except Exception:
                await response.aclose()
                raise
            from ._exceptions import APIError

            raise APIError(
                message=f"Expected streaming response (text/event-stream), got {content_type or 'unknown'}",
                status_code=response.status_code,
                body=RungateResponse(response).body,
            )

        return response

    async def close(self) -> None:
        """Close the underlying async HTTP client."""
        await self._http_client.aclose()

    async def __aenter__(self) -> AsyncRungate:
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.close()
