from __future__ import annotations

import asyncio
import json
from functools import partial
from typing import Any, Mapping, Optional

from .headers import _header_get
from .transport import _request_raw


def call_agent(
    url: str,
    *,
    payload: Optional[Any] = None,
    method: str = "POST",
    headers: Optional[Mapping[str, str]] = None,
    timeout: float = 30.0,
) -> Any:
    """
    Call a backend agent URL while forwarding tracker IDs.

    Adds these headers when available from context:
    - X-Request-ID
    - X-Session-ID
    - X-Trace-ID

    If payload is provided, it is JSON-encoded.
    """
    raw = _request_raw(
        url,
        payload=payload,
        method=method,
        headers=headers,
        timeout=timeout,
    )
    body = raw["body"]
    response_headers = raw["headers"]
    content_type = _header_get(response_headers, "Content-Type") or ""

    if "application/json" in content_type:
        return json.loads(body.decode("utf-8")) if body else None

    return body.decode("utf-8") if body else None


async def call_agent_async(
    url: str,
    *,
    payload: Optional[Any] = None,
    method: str = "POST",
    headers: Optional[Mapping[str, str]] = None,
    timeout: float = 30.0,
) -> Any:
    """Async wrapper around call_agent."""
    fn = partial(
        call_agent,
        url,
        payload=payload,
        method=method,
        headers=headers,
        timeout=timeout,
    )
    return await asyncio.to_thread(fn)
