from __future__ import annotations

import asyncio
import functools
import inspect
import os
import time
import warnings
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from importlib.metadata import PackageNotFoundError, version
from typing import Any, ParamSpec, TypeVar
from uuid import UUID

import httpx

from .recorder import TraceRecorder, safe_json_value
from .types import (
    CIRunStatus,
    CIRunThresholds,
    CampaignReliability,
    CampaignRunResult,
    FailureRecord,
    RunStatus,
    TraceBootstrapResponse,
    TraceMetadata,
    TracePayload,
    TraceTimelineStep,
)

SDK_PACKAGE_NAME = "sepurux"
DEMO_API_KEY = "demo_sepurux_key"
DEFAULT_BASE_URL = "http://localhost:8000"
P = ParamSpec("P")
R = TypeVar("R")


class ManagedTraceRecorder(TraceRecorder):
    def __init__(
        self,
        task_name: str,
        task_input: dict[str, Any] | None = None,
        *,
        source: str = "sdk",
        trace_version: str = "0.1",
    ) -> None:
        super().__init__(
            task_name=task_name,
            task_input=task_input,
            source=source,
            trace_version=trace_version,
        )
        self.trace_id: str | None = None
        self.batch_id: str | None = None
        self.run_id: str | None = None  # first run_id in the batch (backward compat)
        self.upload_error: Exception | None = None


class SepuruxClient:
    def __init__(
        self,
        base_url: str,
        api_key: str | None = None,
        project_id: str | None = None,
        timeout: float = 30,
        sdk_header: str | None = None,
        demo_mode: bool = False,
    ) -> None:
        normalized = base_url.strip().rstrip("/")
        if not normalized:
            raise ValueError("base_url must be a non-empty URL")

        self.base_url = normalized
        self.api_key = api_key
        self.project_id = project_id
        self.timeout = timeout
        self.sdk_header = (sdk_header or _default_sdk_header()).strip() or _default_sdk_header()
        self.demo_mode = demo_mode or api_key == DEMO_API_KEY

        self._client = httpx.Client(
            base_url=self.base_url,
            timeout=timeout,
            headers={"Accept": "application/json", "X-Sepurux-SDK": self.sdk_header},
        )
        if self.demo_mode:
            warnings.warn("This is a limited demo environment.", stacklevel=2)

    @classmethod
    def from_env(
        cls,
        *,
        base_url_env: str = "SEPURUX_API_BASE_URL",
        api_key_env: str = "SEPURUX_API_KEY",
        project_id_env: str = "SEPURUX_PROJECT_ID",
        default_base_url: str = DEFAULT_BASE_URL,
        timeout: float = 30,
        sdk_header: str | None = None,
        demo_mode: bool = False,
    ) -> SepuruxClient:
        base_url = os.getenv(base_url_env, default_base_url)
        return cls(
            base_url=base_url,
            api_key=os.getenv(api_key_env),
            project_id=os.getenv(project_id_env),
            timeout=timeout,
            sdk_header=sdk_header,
            demo_mode=demo_mode,
        )

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

    def upload_trace(self, trace: dict[str, Any] | TracePayload) -> str:
        if not isinstance(trace, dict):
            raise TypeError("trace must be a dictionary")

        response = self._client.post(
            "/v1/traces",
            headers=self._request_headers(),
            json=trace,
        )
        response.raise_for_status()

        payload = _parse_json_response(response)
        trace_id = payload.get("trace_id")
        if not isinstance(trace_id, str) or not trace_id:
            raise ValueError("Trace upload succeeded but trace_id was missing")
        return trace_id

    def list_traces(self, *, limit: int = 50, offset: int = 0) -> list[TraceMetadata]:
        response = self._client.get(
            "/v1/traces",
            headers=self._request_headers(),
            params={"limit": limit, "offset": offset},
        )
        response.raise_for_status()
        payload = _parse_json_list_response(response)
        return [item for item in payload if isinstance(item, dict)]

    def get_trace(self, trace_id: str, *, include_payload: bool = False) -> TraceMetadata:
        response = self._client.get(
            f"/v1/traces/{trace_id}",
            headers=self._request_headers(),
            params={"include_payload": "true"} if include_payload else None,
        )
        response.raise_for_status()
        payload = _parse_json_response(response)
        return payload

    def get_trace_timeline(self, trace_id: str) -> list[TraceTimelineStep]:
        response = self._client.get(
            f"/v1/traces/{trace_id}/timeline",
            headers=self._request_headers(),
        )
        response.raise_for_status()
        payload = _parse_json_list_response(response)
        return [item for item in payload if isinstance(item, dict)]

    def bootstrap_trace(self, trace_id: str) -> TraceBootstrapResponse:
        response = self._client.post(
            f"/v1/traces/{trace_id}/bootstrap",
            headers=self._request_headers(),
        )
        response.raise_for_status()
        payload = _parse_json_response(response)
        if not isinstance(payload.get("tools_discovered"), list):
            payload["tools_discovered"] = []
        if not isinstance(payload.get("campaigns"), list):
            payload["campaigns"] = []
        return payload

    def create_campaign(
        self,
        *,
        name: str,
        mutation_set: Any,
        eval_set: Any,
        mutation_pack_id: str | None = None,
    ) -> str:
        payload: dict[str, Any] = {
            "name": name,
            "mutation_set": mutation_set,
            "eval_set": eval_set,
        }
        if mutation_pack_id:
            payload["mutation_pack_id"] = mutation_pack_id

        response = self._client.post(
            "/v1/campaigns",
            headers=self._request_headers(),
            json=payload,
        )
        response.raise_for_status()

        data = _parse_json_response(response)
        campaign_id = data.get("campaign_id")
        if not isinstance(campaign_id, str) or not campaign_id:
            raise ValueError("Campaign creation succeeded but campaign_id was missing")
        return campaign_id

    def start_run(
        self,
        *,
        trace_id: str,
        campaign_id: str,
        mutation_pack_id: str | None = None,
        thresholds: dict[str, Any] | CIRunThresholds | None = None,
        preflight: bool = False,
    ) -> str:
        """Start a CI batch run. Returns the batch_id."""
        if preflight:
            self.preflight_run_inputs(trace_id=trace_id, campaign_id=campaign_id)

        ci_payload: dict[str, Any] = {
            "trace_id": trace_id,
            "campaign_id": campaign_id,
            "thresholds": thresholds or {"min_pass_rate": 0.9, "max_unsafe": 0},
        }
        if mutation_pack_id:
            ci_payload["mutation_pack_id"] = mutation_pack_id
        github_repo = (os.getenv("GITHUB_REPO") or "").strip()
        github_pr_number_raw = (os.getenv("PR_NUMBER") or "").strip()
        if github_repo:
            ci_payload["repo"] = github_repo
        if github_pr_number_raw:
            try:
                ci_payload["pull_request_number"] = int(github_pr_number_raw)
            except ValueError:
                pass

        response = self._client.post(
            "/v1/ci/runs",
            headers=self._request_headers(),
            json=ci_payload,
        )
        response.raise_for_status()
        payload = _parse_json_response(response)
        batch_id = payload.get("batch_id")
        if not isinstance(batch_id, str) or not batch_id:
            raise ValueError("Run creation succeeded but batch_id was missing")
        return batch_id

    def get_ci_run(self, run_id: str) -> CIRunStatus:
        response = self._client.get(
            f"/v1/ci/runs/{run_id}",
            headers=self._request_headers(),
        )
        response.raise_for_status()
        payload = _parse_json_response(response)

        top_failing_tools = payload.get("top_failing_tools")
        mutation_categories_raw = payload.get("mutation_categories")
        mutation_categories = (
            [item for item in mutation_categories_raw if isinstance(item, str)]
            if isinstance(mutation_categories_raw, list)
            else []
        )
        mutation_coverage_raw = payload.get("mutation_coverage_score")
        mutation_coverage = (
            int(mutation_coverage_raw) if isinstance(mutation_coverage_raw, int | float) else None
        )
        return {
            "run_id": str(payload.get("run_id") or run_id),
            "status": str(payload.get("status") or "unknown"),
            "decision": str(payload.get("decision") or "pending"),
            "pass_rate": (
                float(payload.get("pass_rate"))
                if isinstance(payload.get("pass_rate"), int | float)
                else 0.0
            ),
            "unsafe_attempts": (
                int(payload.get("unsafe_attempts"))
                if isinstance(payload.get("unsafe_attempts"), int | float)
                else 0
            ),
            "failures": (
                int(payload.get("failures"))
                if isinstance(payload.get("failures"), int | float)
                else 0
            ),
            "top_failing_tools": (
                [item for item in top_failing_tools if isinstance(item, dict)]
                if isinstance(top_failing_tools, list)
                else []
            ),
            "mutation_categories": mutation_categories,
            "mutation_coverage_score": mutation_coverage,
            "dashboard_url": str(payload.get("dashboard_url") or ""),
        }

    def get_ci_batch(self, batch_id: str) -> CIRunStatus:
        response = self._client.get(
            f"/v1/ci/batch/{batch_id}",
            headers=self._request_headers(),
        )
        response.raise_for_status()
        payload = _parse_json_response(response)

        top_failing_tools = payload.get("top_failing_tools")
        mutation_categories_raw = payload.get("mutation_categories")
        mutation_categories = (
            [item for item in mutation_categories_raw if isinstance(item, str)]
            if isinstance(mutation_categories_raw, list)
            else []
        )
        mutation_coverage_raw = payload.get("mutation_coverage_score")
        mutation_coverage = (
            int(mutation_coverage_raw) if isinstance(mutation_coverage_raw, int | float) else None
        )
        run_ids_raw = payload.get("run_ids")
        run_ids = (
            [str(r) for r in run_ids_raw if r] if isinstance(run_ids_raw, list) else []
        )
        return {
            "batch_id": str(payload.get("batch_id") or batch_id),
            "run_id": run_ids[0] if run_ids else batch_id,
            "run_ids": run_ids,
            "trace_count": int(payload.get("trace_count") or 1),
            "scenario_coverage_pct": float(payload.get("scenario_coverage_pct") or 0.0),
            "status": str(payload.get("status") or "unknown"),
            "decision": str(payload.get("decision") or "pending"),
            "pass_rate": (
                float(payload.get("pass_rate"))
                if isinstance(payload.get("pass_rate"), int | float)
                else 0.0
            ),
            "unsafe_attempts": (
                int(payload.get("unsafe_attempts"))
                if isinstance(payload.get("unsafe_attempts"), int | float)
                else 0
            ),
            "failures": (
                int(payload.get("failures"))
                if isinstance(payload.get("failures"), int | float)
                else 0
            ),
            "top_failing_tools": (
                [item for item in top_failing_tools if isinstance(item, dict)]
                if isinstance(top_failing_tools, list)
                else []
            ),
            "mutation_categories": mutation_categories,
            "mutation_coverage_score": mutation_coverage,
            "dashboard_url": str(payload.get("dashboard_url") or ""),
        }

    def preflight_run_inputs(self, *, trace_id: str, campaign_id: str) -> None:
        self._assert_resource_exists(
            path=f"/v1/traces/{trace_id}",
            label="Trace",
            identifier=trace_id,
        )
        self._assert_resource_exists(
            path=f"/v1/campaigns/{campaign_id}",
            label="Campaign",
            identifier=campaign_id,
        )

    def _assert_resource_exists(self, *, path: str, label: str, identifier: str) -> None:
        response = self._client.get(path, headers=self._request_headers())
        if response.status_code == 404:
            detail = _extract_error_detail(response)
            message = f"{label} not found: {identifier}"
            if detail and detail.lower() not in {"not found", f"{label.lower()} not found"}:
                message = f"{message} ({detail})"
            raise ValueError(message)
        response.raise_for_status()

    def register_eval_traces(
        self,
        campaign_id: str,
        trace_ids: list[str],
        *,
        remove_trace_ids: list[str] | None = None,
        min_scenarios: int | None = None,
    ) -> None:
        """Register traces in a campaign's eval-set for multi-scenario CI runs.

        After calling this, POST /v1/ci/runs will automatically run all
        registered traces as one batch — no need to pass trace_id per call.
        """
        payload: dict[str, Any] = {
            "add_trace_ids": trace_ids,
            "remove_trace_ids": remove_trace_ids or [],
        }
        if min_scenarios is not None:
            payload["min_scenarios"] = min_scenarios
        response = self._client.patch(
            f"/v1/campaigns/{campaign_id}/eval-set",
            headers=self._request_headers(),
            json=payload,
        )
        response.raise_for_status()

    def run_pack(self, *, campaign_id: str, mutation_pack: str) -> str:
        response = self._client.post(
            f"/v1/campaigns/{campaign_id}/run-pack",
            headers=self._request_headers(),
            json={"pack": mutation_pack},
        )
        response.raise_for_status()
        payload = _parse_json_response(response)
        run_id = payload.get("run_id")
        if not isinstance(run_id, str) or not run_id:
            raise ValueError("Run-pack succeeded but run_id was missing")
        return run_id

    def get_demo_info(self) -> dict[str, Any]:
        response = self._client.get("/demo/info", headers={"Accept": "application/json"})
        response.raise_for_status()
        return _parse_json_response(response)

    @contextmanager
    def trace(
        self,
        task_name: str,
        task_input: dict[str, Any] | None = None,
        *,
        source: str = "sdk",
        trace_version: str = "0.1",
        campaign_id: str | None = None,
        mutation_pack: str | None = None,
        raise_on_error: bool = False,
        error_tool: str | None = None,
    ) -> Iterator[ManagedTraceRecorder]:
        if mutation_pack and not campaign_id:
            raise ValueError("campaign_id is required when mutation_pack is provided")

        recorder = ManagedTraceRecorder(
            task_name=task_name,
            task_input=task_input,
            source=source,
            trace_version=trace_version,
        )
        if task_input is not None:
            recorder.prompt(safe_json_value(task_input), name=task_name)

        try:
            yield recorder
        except Exception as exc:
            recorder.error(tool=error_tool, message=str(exc))
            try:
                self._finalize_trace(
                    recorder,
                    campaign_id=campaign_id,
                    mutation_pack=mutation_pack,
                )
            except Exception as finalize_exc:
                recorder.upload_error = finalize_exc
            raise
        else:
            try:
                self._finalize_trace(
                    recorder,
                    campaign_id=campaign_id,
                    mutation_pack=mutation_pack,
                )
            except Exception as finalize_exc:
                recorder.upload_error = finalize_exc
                if raise_on_error:
                    raise

    def traceable(
        self,
        *,
        name: str | None = None,
        campaign_id: str | None = None,
        mutation_pack: str | None = None,
        source: str = "sdk",
        raise_on_error: bool = False,
    ) -> Callable[[Callable[P, R]], Callable[P, R]]:
        if mutation_pack and not campaign_id:
            raise ValueError("campaign_id is required when mutation_pack is provided")

        def decorator(func: Callable[P, R]) -> Callable[P, R]:
            if inspect.iscoroutinefunction(func):
                @functools.wraps(func)
                async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
                    task_name = name or func.__name__
                    task_input = {
                        "function": func.__name__,
                        "args": safe_json_value(list(args)),
                        "kwargs": safe_json_value(dict(kwargs)),
                    }
                    recorder = ManagedTraceRecorder(
                        task_name=task_name,
                        task_input=task_input,
                        source=source,
                    )
                    recorder.prompt(safe_json_value(task_input), name=task_name)
                    recorder.model_step(
                        "function.call",
                        {"qualname": getattr(func, "__qualname__", func.__name__)},
                    )
                    try:
                        result = await func(*args, **kwargs)  # type: ignore[misc]
                    except Exception as exc:
                        recorder.error(tool=func.__name__, message=str(exc))
                        try:
                            await asyncio.to_thread(
                                self._finalize_trace,
                                recorder,
                                campaign_id=campaign_id,
                                mutation_pack=mutation_pack,
                            )
                        except Exception as finalize_exc:
                            recorder.upload_error = finalize_exc
                        raise
                    else:
                        recorder.model_step(
                            "function.return",
                            {"ok": True},
                            output=safe_json_value(result),
                        )
                        recorder.agent_output(safe_json_value(result), name=task_name)
                        try:
                            await asyncio.to_thread(
                                self._finalize_trace,
                                recorder,
                                campaign_id=campaign_id,
                                mutation_pack=mutation_pack,
                            )
                        except Exception as finalize_exc:
                            recorder.upload_error = finalize_exc
                            if raise_on_error:
                                raise
                        return result  # type: ignore[return-value]

                return async_wrapper  # type: ignore[return-value]

            @functools.wraps(func)
            def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
                task_name = name or func.__name__
                task_input = {
                    "function": func.__name__,
                    "args": safe_json_value(list(args)),
                    "kwargs": safe_json_value(dict(kwargs)),
                }

                with self.trace(
                    task_name,
                    task_input=task_input,
                    source=source,
                    campaign_id=campaign_id,
                    mutation_pack=mutation_pack,
                    raise_on_error=raise_on_error,
                    error_tool=func.__name__,
                ) as recorder:
                    recorder.model_step(
                        "function.call",
                        {"qualname": getattr(func, "__qualname__", func.__name__)},
                    )
                    result = func(*args, **kwargs)
                    recorder.model_step(
                        "function.return",
                        {"ok": True},
                        output=safe_json_value(result),
                    )
                    recorder.agent_output(safe_json_value(result), name=task_name)
                    return result

            return wrapper

        return decorator

    def run_campaign(
        self,
        campaign_name_or_id: str,
        *,
        trace: dict[str, Any] | None = None,
        mutation_pack: str | None = None,
        wait_for_completion: bool = True,
        timeout_seconds: float = 60,
        poll_interval: float = 1.0,
        preflight: bool = False,
    ) -> CampaignRunResult:
        campaign_id = self._resolve_campaign_id(campaign_name_or_id)
        trace_id: str | None = None
        batch_id: str | None = None
        run_id: str | None = None
        if mutation_pack:
            run_id = self.run_pack(campaign_id=campaign_id, mutation_pack=mutation_pack)
        else:
            trace_payload = trace or _build_demo_trace(campaign_name_or_id)
            trace_id = self.upload_trace(trace_payload)
            batch_id = self.start_run(trace_id=trace_id, campaign_id=campaign_id, preflight=preflight)
            run_id = batch_id  # backward compat alias

        if not wait_for_completion:
            return {
                "batch_id": batch_id,
                "run_id": run_id or "",
                "run_ids": [run_id] if run_id else [],
                "campaign_id": campaign_id,
                "trace_id": trace_id,
                "status": "queued",
                "summary": None,
                "reliability_score": None,
                "reliability_status": None,
                "failures": [],
            }

        if batch_id is not None:
            batch_status = self._wait_for_batch_completion(
                batch_id,
                timeout_seconds=timeout_seconds,
                poll_interval=poll_interval,
            )
            first_run_id = batch_status.get("run_id") or batch_id
            reliability: CampaignReliability | None = None
            failures: list[FailureRecord] = []
            if batch_status["status"] not in {"queued", "running", "pending"}:
                reliability = self.get_campaign_reliability(campaign_id, run_id=first_run_id)
                failures = self.get_campaign_failures(campaign_id, run_id=first_run_id)
            return {
                "batch_id": batch_id,
                "run_id": first_run_id,
                "run_ids": batch_status.get("run_ids") or [first_run_id],
                "campaign_id": campaign_id,
                "trace_id": trace_id,
                "status": batch_status["status"],
                "summary": None,
                "reliability_score": (
                    float(reliability["reliability_score"])
                    if reliability is not None
                    and isinstance(reliability.get("reliability_score"), int | float)
                    else None
                ),
                "reliability_status": (
                    str(reliability["status"])
                    if reliability is not None and isinstance(reliability.get("status"), str)
                    else None
                ),
                "failures": failures,
            }

        # run_pack path — still uses single run polling
        assert run_id is not None
        run_status = self._wait_for_run_completion(
            run_id,
            timeout_seconds=timeout_seconds,
            poll_interval=poll_interval,
        )
        reliability = None
        failures = []
        if run_status["status"] not in {"queued", "running", "pending"}:
            reliability = self.get_campaign_reliability(campaign_id, run_id=run_id)
            failures = self.get_campaign_failures(campaign_id, run_id=run_id)
        return {
            "batch_id": None,
            "run_id": run_id,
            "run_ids": [run_id],
            "campaign_id": campaign_id,
            "trace_id": trace_id,
            "status": run_status["status"],
            "summary": run_status.get("summary"),
            "reliability_score": (
                float(reliability["reliability_score"])
                if reliability is not None
                and isinstance(reliability.get("reliability_score"), int | float)
                else None
            ),
            "reliability_status": (
                str(reliability["status"])
                if reliability is not None and isinstance(reliability.get("status"), str)
                else None
            ),
            "failures": failures,
        }

    def get_run(self, run_id: str) -> RunStatus:
        response = self._client.get(
            f"/v1/runs/{run_id}",
            headers=self._request_headers(),
        )
        response.raise_for_status()

        payload = _parse_json_response(response)
        summary = payload.get("summary")
        return {
            "run_id": str(payload.get("run_id") or run_id),
            "status": str(payload.get("status") or "unknown"),
            "summary": summary if isinstance(summary, dict) else None,
        }

    def get_campaign_failures(
        self,
        campaign_id: str,
        *,
        run_id: str | None = None,
        limit: int = 200,
        offset: int = 0,
    ) -> list[FailureRecord]:
        params: dict[str, Any] = {"limit": limit, "offset": offset}
        if run_id:
            params["run_id"] = run_id
        response = self._client.get(
            f"/v1/campaigns/{campaign_id}/failures",
            headers=self._request_headers(),
            params=params,
        )
        response.raise_for_status()
        payload = _parse_json_list_response(response)
        return [item for item in payload if isinstance(item, dict)]

    def get_campaign_reliability(
        self,
        campaign_id: str,
        *,
        run_id: str | None = None,
    ) -> CampaignReliability:
        params: dict[str, Any] = {}
        if run_id:
            params["run_id"] = run_id
        response = self._client.get(
            f"/v1/campaigns/{campaign_id}/reliability",
            headers=self._request_headers(),
            params=params,
        )
        response.raise_for_status()
        payload = _parse_json_response(response)
        return {
            "reliability_score": float(payload.get("reliability_score", 0.0)),
            "status": str(payload.get("status") or "healthy"),
            "classified_failures": int(payload.get("classified_failures", 0)),
            "failure_breakdown": (
                payload.get("failure_breakdown")
                if isinstance(payload.get("failure_breakdown"), list)
                else []
            ),
        }

    def __enter__(self) -> SepuruxClient:
        return self

    def __exit__(self, exc_type: object, exc: object, tb: object) -> None:
        del exc_type, exc, tb
        self.close()

    def _request_headers(self) -> dict[str, str]:
        headers: dict[str, str] = {}
        if self.api_key:
            headers["X-API-Key"] = self.api_key
        if self.project_id:
            headers["X-Project-Id"] = self.project_id
        headers["X-Sepurux-SDK"] = self.sdk_header
        return headers

    def _resolve_campaign_id(self, campaign_name_or_id: str) -> str:
        try:
            return str(UUID(campaign_name_or_id))
        except ValueError:
            pass

        response = self._client.get("/v1/campaigns", headers=self._request_headers())
        response.raise_for_status()
        payload = response.json()
        if not isinstance(payload, list):
            raise ValueError("Expected campaign list response")
        for item in payload:
            if not isinstance(item, dict):
                continue
            name = item.get("name")
            campaign_id = item.get("campaign_id")
            if name == campaign_name_or_id and isinstance(campaign_id, str):
                return campaign_id
        raise ValueError(f"Campaign not found: {campaign_name_or_id}")

    def _finalize_trace(
        self,
        recorder: ManagedTraceRecorder,
        *,
        campaign_id: str | None = None,
        mutation_pack: str | None = None,
    ) -> None:
        recorder.trace_id = self.upload_trace(recorder.to_trace())
        if campaign_id:
            resolved_campaign_id = self._resolve_campaign_id(campaign_id)
            if mutation_pack:
                recorder.run_id = self.run_pack(
                    campaign_id=resolved_campaign_id,
                    mutation_pack=mutation_pack,
                )
            else:
                batch_id = self.start_run(
                    trace_id=recorder.trace_id,
                    campaign_id=resolved_campaign_id,
                )
                recorder.batch_id = batch_id
                recorder.run_id = batch_id  # backward compat

    def _wait_for_run_completion(
        self,
        run_id: str,
        *,
        timeout_seconds: float,
        poll_interval: float,
    ) -> RunStatus:
        deadline = time.monotonic() + timeout_seconds
        run_status = self.get_run(run_id)
        while (
            run_status["status"] in {"queued", "running", "pending"} and time.monotonic() < deadline
        ):
            time.sleep(poll_interval)
            run_status = self.get_run(run_id)
        return run_status

    def _wait_for_batch_completion(
        self,
        batch_id: str,
        *,
        timeout_seconds: float,
        poll_interval: float,
    ) -> CIRunStatus:
        deadline = time.monotonic() + timeout_seconds
        batch_status = self.get_ci_batch(batch_id)
        while (
            batch_status["status"] in {"queued", "running", "pending"}
            and time.monotonic() < deadline
        ):
            time.sleep(poll_interval)
            batch_status = self.get_ci_batch(batch_id)
        return batch_status


def _parse_json_response(response: httpx.Response) -> dict[str, Any]:
    data = response.json()
    if not isinstance(data, dict):
        raise ValueError("Expected JSON object response")
    return data


def _parse_json_list_response(response: httpx.Response) -> list[Any]:
    data = response.json()
    if not isinstance(data, list):
        raise ValueError("Expected JSON array response")
    return data


def _extract_error_detail(response: httpx.Response) -> str | None:
    try:
        data = response.json()
    except ValueError:
        return None
    if not isinstance(data, dict):
        return None
    detail = data.get("detail")
    return detail if isinstance(detail, str) else None


def _default_sdk_header() -> str:
    try:
        sdk_version = version(SDK_PACKAGE_NAME)
    except PackageNotFoundError:
        sdk_version = "dev"
    return f"py/{sdk_version}"


def _build_demo_trace(campaign_name_or_id: str) -> dict[str, Any]:
    return {
        "trace_version": "0.1",
        "source": "sdk-demo",
        "task": {
            "name": "demo_campaign_run",
            "input": {
                "campaign": campaign_name_or_id,
            },
        },
        "events": [
            {
                "type": "model",
                "name": "demo.start",
                "input": {"campaign": campaign_name_or_id},
            }
        ],
    }
