"""
Stax TA API client (sync and async). Generated from specs/ta.json.
Clients call the provided inject_headers callback on every request (e.g. for STAX_KEY, STAX_MODULE_ID).
"""

from __future__ import annotations

from collections.abc import Callable
from typing import MutableMapping

from stax_sdk._http import async_client as _async_client
from stax_sdk._http import sync_client as _sync_client

DEFAULT_TA_BASE_URL = "https://api.stax.ai"


class Client:
    """
    Sync Stax TA API client. Uses inject_headers on every request.
    """

    def __init__(
        self,
        base_url: str = DEFAULT_TA_BASE_URL,
        inject_headers: Callable[[MutableMapping[str, str]], None] | None = None,
        **kwargs: object,
    ) -> None:
        self._client = _sync_client(
            base_url=base_url.rstrip("/"),
            inject_headers=inject_headers,
            **kwargs,
        )

    @property
    def client(self):
        """Underlying httpx.Client for making requests."""
        return self._client

    def __enter__(self) -> Client:
        self._client.__enter__()
        return self

    def __exit__(self, *args: object) -> None:
        self._client.__exit__(*args)


class AsyncClient:
    """
    Async Stax TA API client. Uses inject_headers on every request.
    """

    def __init__(
        self,
        base_url: str = DEFAULT_TA_BASE_URL,
        inject_headers: Callable[[MutableMapping[str, str]], None] | None = None,
        **kwargs: object,
    ) -> None:
        self._client = _async_client(
            base_url=base_url.rstrip("/"),
            inject_headers=inject_headers,
            **kwargs,
        )

    @property
    def client(self):
        """Underlying httpx.AsyncClient for making requests."""
        return self._client

    async def __aenter__(self) -> AsyncClient:
        await self._client.__aenter__()
        return self

    async def __aexit__(self, *args: object) -> None:
        await self._client.__aexit__(*args)


__all__ = ["Client", "AsyncClient", "DEFAULT_TA_BASE_URL"]
