from __future__ import annotations

import json
import math
import random
import time
from typing import (
    Any,
    Generic,
    Mapping,
    Type,
    TypeVar,
    Union,
    overload,
)

import anyio
import httpx

from ._constants import (
    DEFAULT_MAX_RETRIES,
    DEFAULT_TIMEOUT,
    RETRYABLE_STATUS_CODES,
)
from ._exceptions import (
    APIConnectionError,
    APITimeoutError,
    make_status_error,
)
from ._files import to_httpx_files
from ._models import BaseModel, FinalRequestOptions
from ._types import Headers, NotGiven, is_given
from ._version import __version__

_T = TypeVar("_T")


def _build_headers(
    *,
    api_key: str,
    custom_headers: Mapping[str, str] | None,
) -> dict[str, str]:
    headers: dict[str, str] = {
        "Accept": "application/json",
        "User-Agent": f"PonyFlash/Python {__version__}",
        "Authorization": f"Bearer {api_key}",
    }
    if custom_headers:
        headers.update(custom_headers)
    return headers


def _should_retry(response: httpx.Response) -> bool:
    return response.status_code in RETRYABLE_STATUS_CODES


def _retry_delay(attempt: int, retry_after: str | None) -> float:
    if retry_after:
        try:
            return float(retry_after)
        except ValueError:
            pass
    delay = min(0.5 * (2**attempt), 8.0)
    jitter = random.random() * 0.5
    return delay + jitter


def _parse_json(response: httpx.Response) -> Any:
    content_type = response.headers.get("content-type", "")
    if "application/json" in content_type:
        return response.json()
    text = response.text
    if text:
        try:
            return json.loads(text)
        except (json.JSONDecodeError, ValueError):
            return None
    return None


def _unwrap_envelope(data: Any) -> Any:
    """Extract ``data`` from the unified response envelope ``{ error, data }``."""
    if not isinstance(data, dict) or "error" not in data or "data" not in data:
        return data
    err = data.get("error")
    if isinstance(err, dict):
        code = err.get("code", 0)
        if code != 0:
            from ._exceptions import APIBusinessError
            raise APIBusinessError(
                code=code,
                message=err.get("msg") or err.get("message") or f"API error code {code}",
            )
    payload = data.get("data")
    return payload if payload is not None else data


def _process_response(response: httpx.Response, cast_to: Type[_T]) -> _T:
    data = _parse_json(response)
    if cast_to is httpx.Response:
        return response  # type: ignore[return-value]
    if data is None:
        data = {}
    data = _unwrap_envelope(data)
    if cast_to is None or cast_to is type(None):
        return data  # type: ignore[return-value]
    if isinstance(cast_to, type) and issubclass(cast_to, BaseModel):
        return cast_to.model_validate(data)  # type: ignore[return-value]
    return data  # type: ignore[return-value]


def _serialize_multipart(data: dict[str, Any]) -> dict[str, Any]:
    """Flatten a dict into simple key-value pairs suitable for multipart form data."""
    result: dict[str, Any] = {}
    for key, value in data.items():
        if value is None or not is_given(value):
            continue
        if isinstance(value, (str, int, float, bool)):
            result[key] = str(value) if not isinstance(value, str) else value
        elif isinstance(value, list):
            for item in value:
                arr_key = f"{key}[]"
                if isinstance(item, dict):
                    for sub_k, sub_v in item.items():
                        result[f"{arr_key}[{sub_k}]"] = str(sub_v) if not isinstance(sub_v, str) else sub_v
                else:
                    result.setdefault(arr_key, [])
                    if isinstance(result[arr_key], list):
                        result[arr_key].append(str(item) if not isinstance(item, str) else item)
                    else:
                        result[arr_key] = [result[arr_key], str(item) if not isinstance(item, str) else item]
        elif isinstance(value, dict):
            for sub_k, sub_v in value.items():
                result[f"{key}[{sub_k}]"] = str(sub_v) if not isinstance(sub_v, str) else sub_v
    return result


class SyncAPIClient:
    _base_url: str
    _api_key: str
    _max_retries: int
    _timeout: httpx.Timeout
    _client: httpx.Client

    def __init__(
        self,
        *,
        base_url: str,
        api_key: str,
        max_retries: int = DEFAULT_MAX_RETRIES,
        timeout: httpx.Timeout = DEFAULT_TIMEOUT,
        http_client: httpx.Client | None = None,
        custom_headers: Mapping[str, str] | None = None,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._api_key = api_key
        self._max_retries = max_retries
        self._timeout = timeout
        self._default_headers = _build_headers(api_key=api_key, custom_headers=custom_headers)
        if http_client is not None:
            self._client = http_client
        else:
            self._client = httpx.Client(timeout=timeout)

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

    def __enter__(self) -> SyncAPIClient:
        return self

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

    def _build_url(self, path: str) -> str:
        if path.startswith("http"):
            return path
        return f"{self._base_url}{path}" if path.startswith("/") else f"{self._base_url}/{path}"

    def _build_request(self, options: FinalRequestOptions) -> httpx.Request:
        headers = {**self._default_headers}
        if options.headers:
            headers.update(options.headers)

        url = self._build_url(options.url)
        timeout = options.timeout if options.timeout is not None else self._timeout

        kwargs: dict[str, Any] = {}

        content_type = headers.get("Content-Type", "")
        files = options.files

        if content_type.startswith("multipart/form-data"):
            if "boundary" not in content_type:
                headers.pop("Content-Type", None)
            if options.json_data:
                kwargs["data"] = _serialize_multipart(options.json_data)
            if files:
                kwargs["files"] = files
        elif options.json_data is not None:
            kwargs["json"] = options.json_data

        if options.params:
            kwargs["params"] = options.params

        return self._client.build_request(
            options.method,
            url,
            headers=headers,
            timeout=timeout,
            **kwargs,
        )

    def request(self, cast_to: Type[_T], options: FinalRequestOptions) -> _T:
        max_retries = options.max_retries if options.max_retries is not None else self._max_retries
        retries = 0

        while True:
            try:
                request = self._build_request(options)
                response = self._client.send(request)
            except httpx.TimeoutException as exc:
                if retries >= max_retries:
                    raise APITimeoutError(request=request) from exc
                retries += 1
                time.sleep(_retry_delay(retries, None))
                continue
            except httpx.ConnectError as exc:
                if retries >= max_retries:
                    raise APIConnectionError(request=request) from exc
                retries += 1
                time.sleep(_retry_delay(retries, None))
                continue

            if response.is_success:
                return _process_response(response, cast_to)

            if _should_retry(response) and retries < max_retries:
                retries += 1
                retry_after = response.headers.get("Retry-After")
                time.sleep(_retry_delay(retries, retry_after))
                continue

            body = _parse_json(response)
            raise make_status_error(response, body)

    def get(self, path: str, *, cast_to: Type[_T], options: FinalRequestOptions | None = None) -> _T:
        opts = FinalRequestOptions(method="GET", url=path)
        if options:
            opts = opts.model_copy(update=options.model_dump(exclude_none=True))
        return self.request(cast_to, opts)

    def post(
        self,
        path: str,
        *,
        cast_to: Type[_T],
        body: Any = None,
        files: Any = None,
        options: FinalRequestOptions | None = None,
    ) -> _T:
        headers: dict[str, str] | None = None
        if files:
            headers = {"Content-Type": "multipart/form-data"}
            files = to_httpx_files(files)
        opts = FinalRequestOptions(method="POST", url=path, json_data=body, files=files, headers=headers)
        if options:
            opts = opts.model_copy(update={k: v for k, v in options.model_dump(exclude_none=True).items() if v is not None})
        return self.request(cast_to, opts)

    def delete(self, path: str, *, cast_to: Type[_T], options: FinalRequestOptions | None = None) -> _T:
        opts = FinalRequestOptions(method="DELETE", url=path)
        if options:
            opts = opts.model_copy(update=options.model_dump(exclude_none=True))
        return self.request(cast_to, opts)

    def get_api_list(
        self,
        path: str,
        *,
        page: Type[_T],
        model: Type[Any],
        body: Any = None,
        options: FinalRequestOptions | None = None,
    ) -> _T:
        opts = FinalRequestOptions(method="GET", url=path, json_data=body)
        if options:
            opts = opts.model_copy(update={k: v for k, v in options.model_dump(exclude_none=True).items() if v is not None})
        result = self.request(page, opts)
        if hasattr(result, "_set_private_attributes"):
            result._set_private_attributes(client=self, model=model, options=opts)  # type: ignore[union-attr]
        return result


class AsyncAPIClient:
    _base_url: str
    _api_key: str
    _max_retries: int
    _timeout: httpx.Timeout
    _client: httpx.AsyncClient

    def __init__(
        self,
        *,
        base_url: str,
        api_key: str,
        max_retries: int = DEFAULT_MAX_RETRIES,
        timeout: httpx.Timeout = DEFAULT_TIMEOUT,
        http_client: httpx.AsyncClient | None = None,
        custom_headers: Mapping[str, str] | None = None,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._api_key = api_key
        self._max_retries = max_retries
        self._timeout = timeout
        self._default_headers = _build_headers(api_key=api_key, custom_headers=custom_headers)
        if http_client is not None:
            self._client = http_client
        else:
            self._client = httpx.AsyncClient(timeout=timeout)

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

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

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

    def _build_url(self, path: str) -> str:
        if path.startswith("http"):
            return path
        return f"{self._base_url}{path}" if path.startswith("/") else f"{self._base_url}/{path}"

    def _build_request(self, options: FinalRequestOptions) -> httpx.Request:
        headers = {**self._default_headers}
        if options.headers:
            headers.update(options.headers)

        url = self._build_url(options.url)
        timeout = options.timeout if options.timeout is not None else self._timeout

        kwargs: dict[str, Any] = {}

        content_type = headers.get("Content-Type", "")
        files = options.files

        if content_type.startswith("multipart/form-data"):
            if "boundary" not in content_type:
                headers.pop("Content-Type", None)
            if options.json_data:
                kwargs["data"] = _serialize_multipart(options.json_data)
            if files:
                kwargs["files"] = files
        elif options.json_data is not None:
            kwargs["json"] = options.json_data

        if options.params:
            kwargs["params"] = options.params

        return self._client.build_request(
            options.method,
            url,
            headers=headers,
            timeout=timeout,
            **kwargs,
        )

    async def request(self, cast_to: Type[_T], options: FinalRequestOptions) -> _T:
        max_retries = options.max_retries if options.max_retries is not None else self._max_retries
        retries = 0

        while True:
            try:
                request = self._build_request(options)
                response = await self._client.send(request)
            except httpx.TimeoutException as exc:
                if retries >= max_retries:
                    raise APITimeoutError(request=request) from exc
                retries += 1
                await anyio.sleep(_retry_delay(retries, None))
                continue
            except httpx.ConnectError as exc:
                if retries >= max_retries:
                    raise APIConnectionError(request=request) from exc
                retries += 1
                await anyio.sleep(_retry_delay(retries, None))
                continue

            if response.is_success:
                return _process_response(response, cast_to)

            if _should_retry(response) and retries < max_retries:
                retries += 1
                retry_after = response.headers.get("Retry-After")
                await anyio.sleep(_retry_delay(retries, retry_after))
                continue

            body = _parse_json(response)
            raise make_status_error(response, body)

    async def get(self, path: str, *, cast_to: Type[_T], options: FinalRequestOptions | None = None) -> _T:
        opts = FinalRequestOptions(method="GET", url=path)
        if options:
            opts = opts.model_copy(update=options.model_dump(exclude_none=True))
        return await self.request(cast_to, opts)

    async def post(
        self,
        path: str,
        *,
        cast_to: Type[_T],
        body: Any = None,
        files: Any = None,
        options: FinalRequestOptions | None = None,
    ) -> _T:
        headers: dict[str, str] | None = None
        if files:
            headers = {"Content-Type": "multipart/form-data"}
            files = to_httpx_files(files)
        opts = FinalRequestOptions(method="POST", url=path, json_data=body, files=files, headers=headers)
        if options:
            opts = opts.model_copy(update={k: v for k, v in options.model_dump(exclude_none=True).items() if v is not None})
        return await self.request(cast_to, opts)

    async def delete(self, path: str, *, cast_to: Type[_T], options: FinalRequestOptions | None = None) -> _T:
        opts = FinalRequestOptions(method="DELETE", url=path)
        if options:
            opts = opts.model_copy(update=options.model_dump(exclude_none=True))
        return await self.request(cast_to, opts)

    async def get_api_list(
        self,
        path: str,
        *,
        page: Type[_T],
        model: Type[Any],
        body: Any = None,
        options: FinalRequestOptions | None = None,
    ) -> _T:
        opts = FinalRequestOptions(method="GET", url=path, json_data=body)
        if options:
            opts = opts.model_copy(update={k: v for k, v in options.model_dump(exclude_none=True).items() if v is not None})
        result = await self.request(page, opts)
        if hasattr(result, "_set_private_attributes"):
            result._set_private_attributes(client=self, model=model, options=opts)  # type: ignore[union-attr]
        return result
