"""Sync and async GraphQL clients for the gaiia API."""

from __future__ import annotations

import os
import time
from types import TracebackType
from typing import Any, AsyncIterator, Iterator

import httpx

from gaiia.errors import (
    GaiiaAPIError,
    GaiiaMutationError,
    GaiiaRateLimitedError,
    GaiiaTransportError,
    GraphQLError,
)
from gaiia.pagination import aiter_nodes, iter_nodes
from gaiia.ratelimit import RateLimitInfo

DEFAULT_BASE_URL = "https://api.gaiia.com/api/v1"
DEFAULT_TIMEOUT = 30.0
API_KEY_HEADER = "X-Gaiia-Api-Key"
TIMEZONE_HEADER = "x-timezone"
USER_AGENT = "gaiia-py/0.1.0"


def _build_headers(
    api_key: str, timezone: str | None, extra: dict[str, str] | None
) -> dict[str, str]:
    headers = {
        "Content-Type": "application/json",
        "Accept": "application/json",
        "User-Agent": USER_AGENT,
        API_KEY_HEADER: api_key,
    }
    if timezone:
        headers[TIMEZONE_HEADER] = timezone
    if extra:
        headers.update(extra)
    return headers


def _classify_errors(raw_errors: list[dict[str, Any]]) -> tuple[bool, list[GraphQLError]]:
    """Return ``(rate_limited, parsed_errors)`` for top-level GraphQL errors."""
    errors = [GraphQLError.from_dict(e) for e in raw_errors]
    rate_limited = any(e.code == "RATE_LIMITED" for e in errors)
    return rate_limited, errors


def _raise_for_payload(payload: dict[str, Any]) -> None:
    raw_errors = payload.get("errors")
    if not raw_errors:
        return
    rate_limited, errors = _classify_errors(raw_errors)
    if rate_limited:
        ext = errors[0].extensions
        retry_at = None
        if ext.get("retryAt"):
            from datetime import datetime

            try:
                retry_at = datetime.fromisoformat(str(ext["retryAt"]).replace("Z", "+00:00"))
            except ValueError:
                retry_at = None
        raise GaiiaRateLimitedError(
            errors,
            cost=ext.get("cost"),
            limit=ext.get("limit"),
            used=ext.get("used"),
            remaining=ext.get("remaining"),
            retry_at=retry_at,
        )
    raise GaiiaAPIError(errors, data=payload.get("data"))


class _Base:
    def __init__(
        self,
        api_key: str | None = None,
        *,
        base_url: str = DEFAULT_BASE_URL,
        timezone: str | None = None,
        timeout: float = DEFAULT_TIMEOUT,
        default_headers: dict[str, str] | None = None,
        max_retries: int = 0,
    ) -> None:
        resolved = api_key if api_key is not None else os.environ.get("GAIIA_KEY")
        if not resolved:
            raise ValueError(
                "api_key is required (pass api_key= or set the GAIIA_KEY env var)"
            )
        self._api_key = resolved
        self._base_url = base_url
        self._timezone = timezone
        self._timeout = timeout
        self._default_headers = default_headers
        self._max_retries = max_retries
        self.last_rate_limit: RateLimitInfo | None = None

    def _headers(self, extra: dict[str, str] | None = None) -> dict[str, str]:
        merged = dict(self._default_headers or {})
        if extra:
            merged.update(extra)
        return _build_headers(self._api_key, self._timezone, merged)


class GaiiaClient(_Base):
    """Synchronous gaiia GraphQL client."""

    def __init__(
        self,
        api_key: str | None = None,
        *,
        base_url: str = DEFAULT_BASE_URL,
        timezone: str | None = None,
        timeout: float = DEFAULT_TIMEOUT,
        default_headers: dict[str, str] | None = None,
        max_retries: int = 0,
        transport: httpx.BaseTransport | None = None,
    ) -> None:
        super().__init__(
            api_key,
            base_url=base_url,
            timezone=timezone,
            timeout=timeout,
            default_headers=default_headers,
            max_retries=max_retries,
        )
        self._http = httpx.Client(timeout=timeout, transport=transport)

    def __enter__(self) -> GaiiaClient:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        self.close()

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

    def query(
        self,
        query: str,
        variables: dict[str, Any] | None = None,
        *,
        operation_name: str | None = None,
        timezone: str | None = None,
        extra_headers: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        """Run a GraphQL query/mutation and return the ``data`` payload.

        Raises :class:`GaiiaAPIError` on top-level errors,
        :class:`GaiiaRateLimitedError` on RATE_LIMITED, and
        :class:`GaiiaTransportError` on network failures.
        """
        body: dict[str, Any] = {"query": query}
        if variables is not None:
            body["variables"] = variables
        if operation_name is not None:
            body["operationName"] = operation_name

        per_call_headers = dict(extra_headers or {})
        if timezone is not None:
            per_call_headers[TIMEZONE_HEADER] = timezone
        headers = self._headers(per_call_headers)

        attempt = 0
        while True:
            attempt += 1
            try:
                resp = self._http.post(self._base_url, json=body, headers=headers)
            except httpx.HTTPError as e:
                raise GaiiaTransportError(str(e)) from e

            self.last_rate_limit = RateLimitInfo.from_headers(resp.headers)
            try:
                payload = resp.json()
            except ValueError as e:
                raise GaiiaTransportError(
                    f"non-JSON response from gaiia (HTTP {resp.status_code})"
                ) from e

            try:
                _raise_for_payload(payload)
            except GaiiaRateLimitedError as e:
                if attempt > self._max_retries:
                    raise
                _sleep_until(e.retry_at)
                continue

            if resp.status_code >= 500:
                raise GaiiaTransportError(
                    f"gaiia returned HTTP {resp.status_code}: {resp.text[:200]}"
                )
            return payload.get("data") or {}

    def mutate(
        self,
        query: str,
        mutation_name: str,
        variables: dict[str, Any] | None = None,
        *,
        timezone: str | None = None,
        raise_on_payload_errors: bool = True,
    ) -> dict[str, Any]:
        """Run a mutation and inspect its ``errors`` payload field.

        Returns the mutation's result object (e.g. the contents of
        ``data.createTicket``). If ``raise_on_payload_errors`` is true and the
        result contains a non-empty ``errors`` list, :class:`GaiiaMutationError`
        is raised.
        """
        data = self.query(query, variables, timezone=timezone)
        result = data.get(mutation_name)
        if result is None:
            raise GaiiaAPIError(
                [GraphQLError(message=f"mutation {mutation_name!r} not in response", code=None)],
                data=data,
            )
        errors = result.get("errors") or []
        if errors and raise_on_payload_errors:
            raise GaiiaMutationError(mutation_name, errors, payload=result)
        return result

    def iter_nodes(
        self,
        query: str,
        connection_path: str,
        variables: dict[str, Any] | None = None,
        *,
        cursor_var: str = "after",
        timezone: str | None = None,
    ) -> Iterator[Any]:
        """Iterate every node across all pages of a Connection.

        ``query`` must accept the cursor variable (default ``$after``) and use
        it as the ``after:`` argument on the paginated field.
        """
        base_vars = dict(variables or {})

        def fetch_page(cursor: str | None) -> dict[str, Any]:
            v = dict(base_vars)
            v[cursor_var] = cursor
            return self.query(query, v, timezone=timezone)

        yield from iter_nodes(fetch_page, connection_path)


class AsyncGaiiaClient(_Base):
    """Asynchronous gaiia GraphQL client."""

    def __init__(
        self,
        api_key: str | None = None,
        *,
        base_url: str = DEFAULT_BASE_URL,
        timezone: str | None = None,
        timeout: float = DEFAULT_TIMEOUT,
        default_headers: dict[str, str] | None = None,
        max_retries: int = 0,
        transport: httpx.AsyncBaseTransport | None = None,
    ) -> None:
        super().__init__(
            api_key,
            base_url=base_url,
            timezone=timezone,
            timeout=timeout,
            default_headers=default_headers,
            max_retries=max_retries,
        )
        self._http = httpx.AsyncClient(timeout=timeout, transport=transport)

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

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        await self.aclose()

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

    async def query(
        self,
        query: str,
        variables: dict[str, Any] | None = None,
        *,
        operation_name: str | None = None,
        timezone: str | None = None,
        extra_headers: dict[str, str] | None = None,
    ) -> dict[str, Any]:
        body: dict[str, Any] = {"query": query}
        if variables is not None:
            body["variables"] = variables
        if operation_name is not None:
            body["operationName"] = operation_name

        per_call_headers = dict(extra_headers or {})
        if timezone is not None:
            per_call_headers[TIMEZONE_HEADER] = timezone
        headers = self._headers(per_call_headers)

        attempt = 0
        while True:
            attempt += 1
            try:
                resp = await self._http.post(self._base_url, json=body, headers=headers)
            except httpx.HTTPError as e:
                raise GaiiaTransportError(str(e)) from e

            self.last_rate_limit = RateLimitInfo.from_headers(resp.headers)
            try:
                payload = resp.json()
            except ValueError as e:
                raise GaiiaTransportError(
                    f"non-JSON response from gaiia (HTTP {resp.status_code})"
                ) from e

            try:
                _raise_for_payload(payload)
            except GaiiaRateLimitedError as e:
                if attempt > self._max_retries:
                    raise
                _sleep_until(e.retry_at)
                continue

            if resp.status_code >= 500:
                raise GaiiaTransportError(
                    f"gaiia returned HTTP {resp.status_code}: {resp.text[:200]}"
                )
            return payload.get("data") or {}

    async def mutate(
        self,
        query: str,
        mutation_name: str,
        variables: dict[str, Any] | None = None,
        *,
        timezone: str | None = None,
        raise_on_payload_errors: bool = True,
    ) -> dict[str, Any]:
        data = await self.query(query, variables, timezone=timezone)
        result = data.get(mutation_name)
        if result is None:
            raise GaiiaAPIError(
                [GraphQLError(message=f"mutation {mutation_name!r} not in response", code=None)],
                data=data,
            )
        errors = result.get("errors") or []
        if errors and raise_on_payload_errors:
            raise GaiiaMutationError(mutation_name, errors, payload=result)
        return result

    async def iter_nodes(
        self,
        query: str,
        connection_path: str,
        variables: dict[str, Any] | None = None,
        *,
        cursor_var: str = "after",
        timezone: str | None = None,
    ) -> AsyncIterator[Any]:
        base_vars = dict(variables or {})

        async def fetch_page(cursor: str | None) -> dict[str, Any]:
            v = dict(base_vars)
            v[cursor_var] = cursor
            return await self.query(query, v, timezone=timezone)

        async for node in aiter_nodes(fetch_page, connection_path):
            yield node


def _sleep_until(retry_at: Any) -> None:
    if not retry_at:
        time.sleep(1.0)
        return
    delta = retry_at.timestamp() - time.time()
    if delta > 0:
        time.sleep(min(delta, 60.0))
    else:
        time.sleep(1.0)
