from __future__ import annotations

import json
from typing import Any, AsyncIterator, Dict, Optional

import httpx
from httpx_sse import aconnect_sse

from .common import build_default_headers, extract_request_id
from .async_token import AsyncTokenProvider
from ..config import SDKConfig
from ..exceptions import APIError, NetworkError


class AsyncBaseClient:
    def __init__(self, config: SDKConfig) -> None:
        config.validate()
        self._config = config

        self._http = httpx.AsyncClient(
            base_url=config.base_url.rstrip("/"),
            timeout=config.timeout,
        )
        self._token_provider = AsyncTokenProvider(config, self._http)

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

    async def __aenter__(self) -> "AsyncBaseClient":
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        await self.aclose()

    @staticmethod
    def _extract_request_id(resp: httpx.Response) -> Optional[str]:
        return extract_request_id(resp)

    async def _build_headers(
        self,
        *,
        headers: Optional[Dict[str, str]],
        auth: bool,
        accept: str,
        content_type: Optional[str] = None,
    ) -> Dict[str, str]:
        final_headers = build_default_headers(headers=headers, accept=accept, content_type=content_type)

        if auth:
            token = await self._token_provider.get_token()
            final_headers["Authorization"] = f"Bearer {token}"

        return final_headers

    async def _request(
        self,
        method: str,
        path: str,
        *,
        params: Optional[Dict[str, Any]] = None,
        json_body: Any = None,
        headers: Optional[Dict[str, str]] = None,
        auth: bool = True,
    ) -> Dict[str, Any]:
        final_headers = await self._build_headers(
            headers=headers,
            auth=auth,
            accept="application/json",
            content_type="application/json",
        )

        try:
            resp = await self._http.request(
                method,
                path,
                params=params,
                json=json_body,
                headers=final_headers,
            )
        except httpx.RequestError as e:
            raise NetworkError(f"request failed: {e}") from e

        request_id = self._extract_request_id(resp)

        if resp.status_code >= 400:
            raise APIError(resp.status_code, resp.text, request_id=request_id, body=resp.text)

        data: Dict[str, Any] = resp.json()
        return data

    async def _stream_sse(
        self,
        method: str,
        path: str,
        *,
        params: Optional[Dict[str, Any]] = None,
        json_body: Any = None,
        headers: Optional[Dict[str, str]] = None,
        auth: bool = True,
    ) -> AsyncIterator[Dict[str, Any]]:
        final_headers = await self._build_headers(
            headers=headers,
            auth=auth,
            accept="text/event-stream",
            content_type="application/json",
        )
        final_headers.setdefault("Cache-Control", "no-cache")

        try:
            async with aconnect_sse(
                self._http,
                method,
                path,
                params=params,
                json=json_body,
                headers=final_headers,
            ) as event_source:
                async for sse in event_source.aiter_sse():
                    if sse.data == "[DONE]":
                        break

                    try:
                        payload = json.loads(sse.data) if sse.data else None
                    except json.JSONDecodeError:
                        payload = sse.data

                    yield {"event": sse.event or "message", "data": payload}

        except httpx.RequestError as e:
            raise NetworkError(f"request failed: {e}") from e
