"""
Shared HTTP layer: inject user-defined headers into every request.
The user provides a callback that mutates the headers dict (e.g. from STAX_KEY, STAX_MODULE_ID).
"""

from __future__ import annotations

from collections.abc import Callable
from typing import Any, MutableMapping

import httpx


def _default_inject(headers: MutableMapping[str, str]) -> None:
    """No-op; user replaces with their own injector."""
    pass


class HeaderInjectorTransport(httpx.HTTPTransport):
    """Sync transport that calls inject_headers on the request headers before sending."""

    def __init__(
        self,
        inject_headers: Callable[[MutableMapping[str, str]], None] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self._inject = inject_headers or _default_inject

    def handle_request(self, request: httpx.Request) -> httpx.Response:
        # httpx headers are immutable; get a dict, let user inject, then new request with merged headers
        headers = dict(request.headers)
        self._inject(headers)
        new_request = request.copy(headers=headers)
        return super().handle_request(new_request)


class AsyncHeaderInjectorTransport(httpx.AsyncHTTPTransport):
    """Async transport that calls inject_headers on the request headers before sending."""

    def __init__(
        self,
        inject_headers: Callable[[MutableMapping[str, str]], None] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self._inject = inject_headers or _default_inject

    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
        headers = dict(request.headers)
        self._inject(headers)
        new_request = request.copy(headers=headers)
        return await super().handle_async_request(new_request)


def sync_client(
    base_url: str,
    inject_headers: Callable[[MutableMapping[str, str]], None] | None = None,
    **kwargs: Any,
) -> httpx.Client:
    """Build a sync httpx.Client that calls inject_headers on every request."""
    transport = HeaderInjectorTransport(inject_headers=inject_headers)
    return httpx.Client(base_url=base_url, transport=transport, **kwargs)


def async_client(
    base_url: str,
    inject_headers: Callable[[MutableMapping[str, str]], None] | None = None,
    **kwargs: Any,
) -> httpx.AsyncClient:
    """Build an async httpx.AsyncClient that calls inject_headers on every request."""
    transport = AsyncHeaderInjectorTransport(inject_headers=inject_headers)
    return httpx.AsyncClient(base_url=base_url, transport=transport, **kwargs)
