"""Simulation runner for generating conversations."""

from __future__ import annotations

import asyncio
import logging
import os
import random
from contextlib import contextmanager
from dataclasses import dataclass
from importlib import import_module
from typing import TYPE_CHECKING
from typing import Protocol
from typing import cast
from urllib.parse import urlparse
from urllib.parse import urlunparse

import httpx
import openai
from openai.types.chat import ChatCompletionAssistantMessageParam
from openai.types.chat import ChatCompletionMessageParam
from openai.types.chat import ChatCompletionSystemMessageParam
from openai.types.chat import ChatCompletionUserMessageParam

from collinear.schemas.steer import Role
from collinear.schemas.steer import SimulationResult
from collinear.schemas.steer import SteerCombination
from collinear.schemas.steer import SteerConfig

if TYPE_CHECKING:
    from collections.abc import Iterator

HTTP_UNAUTHORIZED = 401
HTTP_UNPROCESSABLE = 422
MIN_MASK_VISIBLE = 4
DEFAULT_STEER_SEED = -1
MAX_ALLOWED_CONCURRENCY = 8
DEFAULT_CONCURRENCY = 1
BATCH_FLUSH_DELAY_SECONDS = 0.01


class _ProgressBar(Protocol):
    total: int
    n: int

    def update(self, value: int) -> None:
        """Advance the bar by ``value`` units."""

    def refresh(self) -> None:
        """Redraw the bar if the backend supports it."""

    def close(self) -> None:
        """Finalize the bar display."""


class _ProgressFactory(Protocol):
    def __call__(
        self,
        *,
        total: int,
        desc: str | None = None,
        unit: str | None = None,
    ) -> _ProgressBar:
        """Return a new progress bar instance."""


def _load_progress_factory() -> _ProgressFactory | None:
    try:
        module = import_module("tqdm.auto")
    except Exception:
        return None
    creator_obj: object | None = getattr(module, "tqdm", None)
    if creator_obj is None or not callable(creator_obj):
        return None
    return cast("_ProgressFactory", creator_obj)


_PROGRESS_FACTORY: _ProgressFactory | None = _load_progress_factory()


class SimulationRunner:
    """Orchestrates simulation conversations between steers and models.

    Supports split endpoints:
    - USER turns are generated by the Collinear steer API.
    - ASSISTANT turns are generated via an OpenAI-compatible endpoint (customer model).
    """

    ASSISTANT_PROMPT_TEMPLATE = (
        "You are a helpful customer service assistant. Your responses are "
        "helpful, respectful, and succinct. You are talking to a customer.\n\n"
        "You are the ASSISTANT. Respond to the last user message.\n"
        "Write only the next assistant message.\n"
        "Do not include role names or quotes. Avoid ordered, unordered, "
        "numbered, or bullet pointed lists."
    )

    def __init__(
        self,
        assistant_model_url: str,
        assistant_model_api_key: str,
        assistant_model_name: str,
        *,
        steer_api_key: str,
        timeout: float = 30.0,
        max_retries: int = 3,
        rate_limit_retries: int = 6,
    ) -> None:
        """Initialize the simulation runner.

        Args:
            assistant_model_url: Base URL for OpenAI-compatible assistant endpoint.
            assistant_model_api_key: API key for the assistant model endpoint.
            assistant_model_name: Model name for the assistant.
            steer_api_key: Steer service API key (sent as ``API-Key`` to Steer).
            timeout: Request timeout in seconds.
            max_retries: Max retries for assistant calls.
            rate_limit_retries: Max retries upon rate limits for assistant calls.

        """
        if not assistant_model_name:
            raise ValueError("model_name is required")
        if not steer_api_key:
            raise ValueError("steer_api_key is required")
        self.assistant_model_url = assistant_model_url
        self.assistant_model_api_key = assistant_model_api_key
        self.assistant_model_name = assistant_model_name
        self.steer_api_key = steer_api_key
        self.timeout = timeout
        self.max_retries = max_retries
        self.rate_limit_retries = rate_limit_retries
        self.logger = logging.getLogger("collinear")
        self.steer_temperature: float = 0.7
        self.steer_max_tokens: int = 256
        self.steer_seed: int = DEFAULT_STEER_SEED
        self._steer_batch_size: int = DEFAULT_CONCURRENCY
        self._batcher: SimulationRunner._SteerBatcher | None = None
        self._progress_bar: _ProgressBar | None = None
        self._user_turns_completed: int = 0
        self.client = openai.AsyncOpenAI(
            base_url=assistant_model_url,
            api_key=assistant_model_api_key,
            max_retries=max_retries,
            timeout=timeout,
        )

    def run(
        self,
        config: SteerConfig,
        k: int | None = None,
        num_exchanges: int = 2,
        batch_delay: float = 0.1,
        *,
        steer_temperature: float | None = None,
        steer_max_tokens: int | None = None,
        steer_seed: int | None = None,
        mix_traits: bool = False,
        progress: bool = True,
        max_concurrency: int = DEFAULT_CONCURRENCY,
    ) -> list[SimulationResult]:
        """Run simulations with the given configuration.

        Args:
            config: Steer configuration.
            k: Optional number of simulations to run. If ``None``, run all
                available combinations in deterministic order. If provided and
                smaller than the total number of combinations, a random subset
                of size ``k`` is selected.
            num_exchanges: Number of user-assistant exchanges (e.g., 2 = 2 user
                turns + 2 assistant turns).
            batch_delay: Delay between simulations to avoid rate limits (seconds).
            steer_temperature: Optional sampling temperature for steer generation
                (default 0.7).
            steer_max_tokens: Optional max tokens for steer generation (default 256).
            steer_seed: Optional deterministic seed for steer sampling. ``-1``
                (default) delegates randomness to the service.
            mix_traits: If True, use pairwise mixing (exactly two traits per steer).
                Requires at least two traits with available levels.
            progress: Whether to display a tqdm-style bar tracking steer API calls.
                Defaults to ``True``.
            max_concurrency: The maximum number of concurrent conversations to have (default 1).
                Values greater than 1 trigger batch requests (up to 8) via ``/steer_batch``.

        Returns:
            List of simulation results with conv_prefix and response.

        """
        return asyncio.run(
            self.run_async(
                config,
                k,
                num_exchanges,
                batch_delay,
                steer_temperature=steer_temperature,
                steer_max_tokens=steer_max_tokens,
                steer_seed=steer_seed,
                mix_traits=mix_traits,
                progress=progress,
                max_concurrency=max_concurrency,
            )
        )

    async def run_async(
        self,
        config: SteerConfig,
        k: int | None = None,
        num_exchanges: int = 2,
        batch_delay: float = 0.1,
        *,
        steer_temperature: float | None = None,
        steer_max_tokens: int | None = None,
        steer_seed: int | None = None,
        mix_traits: bool = False,
        progress: bool = True,
        max_concurrency: int = DEFAULT_CONCURRENCY,
    ) -> list[SimulationResult]:
        """Async Run simulations with the given configuration.

        Args:
            config: Steer configuration.
            k: Optional number of simulations to run. If ``None``, run all
                available combinations in deterministic order. If provided and
                smaller than the total number of combinations, a random subset
                of size ``k`` is selected.
            num_exchanges: Number of user-assistant exchanges (e.g., 2 = 2 user
                turns + 2 assistant turns).
            batch_delay: Delay between simulations to avoid rate limits (seconds).
            steer_temperature: Optional sampling temperature for steer generation
                (default 0.7).
            mix_traits: If True, use pairwise mixing (exactly two traits per steer).
                Requires at least two traits with available levels.
            steer_max_tokens: Optional max tokens for steer generation (default 256).
            steer_seed: Optional deterministic seed for steer sampling. ``-1``
                (default) delegates randomness to the service.
            progress: Whether to display a tqdm-style bar tracking steer API calls.
                Defaults to ``True``.
            max_concurrency: The maximum number of concurrent conversations to have (default 1).
                Values greater than 1 trigger batch requests (up to 8) via ``/steer_batch``.

        Returns:
            List of simulation results with conv_prefix and response.

        """
        with (
            self._steer_settings(
                steer_temperature=steer_temperature,
                steer_max_tokens=steer_max_tokens,
                steer_seed=steer_seed,
            ),
            self._steer_batch_mode(max_concurrency=max_concurrency),
        ):
            combinations = config.combinations(mix_traits=mix_traits)
            samples = self._select_samples(combinations, k)
            if not samples:
                return []
            total_queries = len(samples) * num_exchanges
            with self._progress_tracking(enabled=progress, total=total_queries):
                return await self._execute_samples(
                    samples, num_exchanges, batch_delay, max_concurrency
                )

    @contextmanager
    def _steer_settings(
        self,
        *,
        steer_temperature: float | None,
        steer_max_tokens: int | None,
        steer_seed: int | None,
    ) -> Iterator[None]:
        prev_temp = self.steer_temperature
        prev_max = self.steer_max_tokens
        prev_seed = self.steer_seed
        try:
            if steer_temperature is not None:
                self.steer_temperature = float(steer_temperature)
            if steer_max_tokens is not None:
                self.steer_max_tokens = int(steer_max_tokens)
            if steer_seed is not None:
                self.steer_seed = int(steer_seed)
            else:
                self.steer_seed = DEFAULT_STEER_SEED
            yield
        finally:
            self.steer_temperature = prev_temp
            self.steer_max_tokens = prev_max
            self.steer_seed = prev_seed

    @contextmanager
    def _steer_batch_mode(self, *, max_concurrency: int) -> Iterator[None]:
        prev_batch_size = self._steer_batch_size
        prev_batcher = self._batcher
        batch_size = self._calculate_semaphore_limit(max_concurrency)
        self._steer_batch_size = batch_size
        self._batcher = None
        if batch_size > 1:
            self._batcher = SimulationRunner._SteerBatcher(self, batch_size)
        try:
            yield
        finally:
            if self._batcher is not None:
                self._batcher.close()
            self._batcher = prev_batcher
            self._steer_batch_size = prev_batch_size

    @dataclass
    class _BatchItem:
        payload: dict[str, object]
        headers: dict[str, str]
        trait_dict: dict[str, int]
        future: asyncio.Future[str]

    class _SteerBatcher:
        def __init__(self, runner: SimulationRunner, batch_size: int) -> None:
            self._runner = runner
            self._batch_size = batch_size
            self._pending: list[SimulationRunner._BatchItem] = []
            self._lock = asyncio.Lock()
            self._flush_task: asyncio.Task[None] | None = None
            self._closed = False

        async def submit(
            self,
            payload: dict[str, object],
            headers: dict[str, str],
            trait_dict: dict[str, int],
        ) -> str:
            if self._closed:
                return await self._runner.call_single_steer(
                    payload,
                    headers=headers,
                    trait_dict=trait_dict,
                )

            loop = asyncio.get_running_loop()
            fut: asyncio.Future[str] = loop.create_future()
            item = SimulationRunner._BatchItem(payload, headers, trait_dict, fut)

            async with self._lock:
                self._pending.append(item)
                if self._flush_task is None or self._flush_task.done():
                    self._flush_task = asyncio.create_task(self._flush())

            return await fut

        async def _flush(self) -> None:
            try:
                await asyncio.sleep(BATCH_FLUSH_DELAY_SECONDS)
                while True:
                    batch = await self._pop_batch()
                    if not batch:
                        return
                    await self._process_batch(batch)
            except asyncio.CancelledError:
                raise
            except Exception:
                self._runner.logger.exception("Batch flush failed")
            finally:
                async with self._lock:
                    if not self._pending:
                        self._flush_task = None

        def close(self) -> None:
            self._closed = True
            flush = self._flush_task
            if flush is not None:
                flush.cancel()

        async def _pop_batch(self) -> list[SimulationRunner._BatchItem]:
            async with self._lock:
                if not self._pending:
                    self._flush_task = None
                    return []
                batch = self._pending[: self._batch_size]
                del self._pending[: len(batch)]
                return batch

        async def _process_batch(self, batch: list[SimulationRunner._BatchItem]) -> None:
            payloads = [item.payload for item in batch]
            headers = batch[0].headers

            success, result = await self._runner.call_batch_steer(payloads, headers=headers)
            if success and isinstance(result, list):
                for item, content in zip(batch, result, strict=False):
                    if not item.future.done():
                        item.future.set_result(content)
                return

            reason = result if isinstance(result, str) else "Batch request failed"
            self._runner.logger.debug(
                "Falling back to single steer requests after batch failure: %s",
                reason,
            )
            for item in batch:
                try:
                    response = await self._runner.call_single_steer(
                        item.payload,
                        headers=item.headers,
                        trait_dict=item.trait_dict,
                    )
                except Exception as exc:
                    if not item.future.done():
                        item.future.set_exception(exc)
                    continue
                if not item.future.done():
                    item.future.set_result(response)

    @contextmanager
    def _progress_tracking(self, *, enabled: bool, total: int) -> Iterator[None]:
        if not enabled or total <= 0:
            yield
            return
        factory = _PROGRESS_FACTORY
        if factory is None:
            self.logger.debug("tqdm not available; progress disabled.")
            yield
            return
        try:
            bar = factory(total=total, desc="User/Assistant turns", unit="query")
        except Exception:
            self.logger.exception("Failed to initialize progress bar; continuing without it.")
            yield
            return
        self._progress_bar = bar
        try:
            yield
        finally:
            self._progress_bar = None
            try:
                bar.close()
            except Exception:
                self.logger.debug("Failed to close progress bar", exc_info=True)

    def _advance_progress(self, step: int) -> None:
        if step <= 0:
            return
        bar = self._progress_bar
        if bar is None:
            return
        try:
            bar.update(step)
        except Exception:
            self.logger.debug("Progress update failed", exc_info=True)

    def _adjust_progress_total(self, decrement: int) -> None:
        if decrement <= 0:
            return
        bar = self._progress_bar
        if bar is None:
            return
        try:
            current_total = bar.total
            new_total = max(bar.n, current_total - decrement)
            if new_total != current_total:
                bar.total = new_total
                bar.refresh()
        except Exception:
            self.logger.debug("Progress total adjustment failed", exc_info=True)

    def _select_samples(
        self, combinations: list[SteerCombination], k: int | None
    ) -> list[SteerCombination]:
        total = len(combinations)
        self.logger.info("Total steer combinations: %d", total)
        if total == 0:
            self.logger.warning("No steer combinations generated; nothing to run.")
            return []
        if k is None or k >= total:
            if k is None:
                self.logger.info("Running all %d combinations (k=None).", total)
            else:
                self.logger.info(
                    "k=%d >= total=%d; running all %d combinations.",
                    k,
                    total,
                    total,
                )
            return combinations
        self.logger.info("Sampling k=%d of %d combinations at random.", k, total)
        return random.sample(combinations, k)

    def _calculate_semaphore_limit(self, max_concurrency: int) -> int:
        """Calculate the semaphore limit respecting bounds.

        Args:
            max_concurrency: Requested concurrency limit

        Returns:
            Actual concurrency limit (bounded by [1, MAX_ALLOWED_CONCURRENCY])

        """
        return min(MAX_ALLOWED_CONCURRENCY, max(1, max_concurrency))

    def calculate_semaphore_limit(self, max_concurrency: int) -> int:
        """Public accessor for the computed concurrency limit."""
        return self._calculate_semaphore_limit(max_concurrency)

    async def _execute_samples(
        self,
        samples: list[SteerCombination],
        num_exchanges: int,
        batch_delay: float,
        max_concurrency: int,
    ) -> list[SimulationResult]:
        simulations: list[SimulationResult] = []
        sem = asyncio.Semaphore(self._calculate_semaphore_limit(max_concurrency))

        async def run_one(i: int, combo: SteerCombination) -> None:
            if i > 0 and batch_delay > 0:
                await asyncio.sleep(batch_delay)
            async with sem:
                try:
                    self.logger.info("=" * 40)
                    conversation, final_response = await self._build_conversation(
                        combo, num_exchanges
                    )
                    simulations.append(
                        SimulationResult(
                            conv_prefix=conversation[:-1],
                            response=final_response,
                            steer=combo,
                        )
                    )
                    self.logger.info(f"Completed simulation {i + 1}/{len(samples)}")
                except SimulationRunner.BuildConversationError as e:
                    remaining = max(0, num_exchanges - e.completed_user_turns)
                    if remaining:
                        self._adjust_progress_total(remaining)
                    if e.invalid_trait:
                        self.logger.warning(
                            "Skipping simulation %d/%d due to invalid trait '%s'.",
                            i + 1,
                            len(samples),
                            e.trait or "<unknown>",
                        )
                    else:
                        self.logger.exception(f"Failed simulation {i + 1}/{len(samples)}")

        await asyncio.gather(*(run_one(i, combo) for i, combo in enumerate(samples)))
        return simulations

    async def _call_with_retry(
        self,
        messages: list[ChatCompletionMessageParam],
        system_prompt: str,
    ) -> str:
        """Make API call with retry logic."""
        sys_msg: ChatCompletionSystemMessageParam = {"role": "system", "content": system_prompt}
        full_messages: list[ChatCompletionMessageParam] = [sys_msg, *messages]
        attempt = 0
        while True:
            try:
                response = await self.client.chat.completions.create(
                    model=self.assistant_model_name,
                    messages=full_messages,
                    temperature=0.8,
                    max_tokens=200,
                )
            except openai.RateLimitError as e:
                attempt += 1
                self.logger.warning(f"Rate limit hit, attempt {attempt}: {e}")
                if attempt >= self.rate_limit_retries:
                    raise
                delay = min(60.0, max(1.0, (2.0 ** (attempt - 1)) + random.random()))
                await asyncio.sleep(delay)
            except Exception as e:
                self.logger.exception("Error getting response")
                return f"Error: {e!s}"
            else:
                content = response.choices[0].message.content
                return content or ""

    async def _generate_turn(
        self,
        combo: SteerCombination,
        conversation: list[ChatCompletionMessageParam],
        role: Role,
    ) -> str:
        """Generate a single turn in the conversation."""
        if role is Role.USER:
            self.logger.info("Generating USER turn")
            use_batch = self._steer_batch_size > 1
            if len(combo.traits) == 1:
                trait = next(iter(combo.traits))
                intensity = next(iter(combo.traits.values()))
                response = await self._call_collinear_steer_api(
                    _conversation=conversation,
                    trait=trait,
                    intensity=intensity,
                    combo=combo,
                    use_batch=use_batch,
                )
            else:
                response = await self._call_collinear_steer_api_trait_dict(
                    _conversation=conversation,
                    trait_dict={k: int(v) for k, v in combo.traits.items()},
                    combo=combo,
                    use_batch=use_batch,
                )
        else:
            self.logger.info("Generating ASSISTANT turn")
            system_prompt = self.ASSISTANT_PROMPT_TEMPLATE
            response = await self._call_with_retry(conversation, system_prompt)
        self.logger.info(response)

        return response

    def _mask_key_preview(self) -> str:
        key = (self.steer_api_key or "").strip()
        return key if len(key) <= MIN_MASK_VISIBLE else key[:2] + "***" + key[-2:]

    def _log_unauthorized(self, resp: httpx.Response) -> None:
        self.logger.error(
            "Steer API unauthorized (401). API-Key preview=%s. Body=%s",
            self._mask_key_preview(),
            resp.text,
        )

    def _log_if_unauthorized(self, resp: httpx.Response) -> None:
        if resp.status_code == HTTP_UNAUTHORIZED:
            self._log_unauthorized(resp)

    class InvalidTraitError(Exception):
        """Raised when the Steer API signals an unknown/unsupported trait."""

        trait: str

        def __init__(self, trait: str) -> None:
            """Initialize with the given trait."""
            super().__init__(trait)
            self.trait = trait

    class BuildConversationError(Exception):
        """Raised when building a conversation fails during simulation.

        Carries metadata about how many user turns were completed before failure,
        and whether the failure was due to an invalid trait.
        """

        def __init__(
            self,
            completed_user_turns: int,
            *,
            invalid_trait: bool = False,
            trait: str | None = None,
        ) -> None:
            """Initialize the exception with failure metadata.

            Args:
                completed_user_turns: Number of user turns completed before failure.
                    Used to adjust progress tracking by reducing the remaining
                    user turns from the progress bar total.
                invalid_trait: Whether this failure was due to an invalid trait
                    (e.g., trait not recognized by the Steer API). Defaults to False.
                trait: The trait name that caused the failure, if applicable.
                    Only meaningful when invalid_trait is True. Defaults to None.

            """
            super().__init__("Conversation build failed")
            self.completed_user_turns = completed_user_turns
            self.invalid_trait = invalid_trait
            self.trait = trait

    def _resolve_steer_endpoint(self, endpoint: str) -> str:
        base = os.getenv("COLLINEAR_STEER_URL", "https://steer.collinear.ai/steer")
        if not base:
            base = "https://steer.collinear.ai/steer"
        parsed = urlparse(base)
        path = parsed.path or ""
        if path.endswith("/"):
            prefix = path
        else:
            sep_index = path.rfind("/")
            prefix = path[: sep_index + 1] if sep_index >= 0 else "/"
        if not prefix.startswith("/"):
            prefix = f"/{prefix}"
        normalized = endpoint.lstrip("/")
        new_path = f"{prefix}{normalized}" if prefix.endswith("/") else f"{prefix}/{normalized}"
        return urlunparse(parsed._replace(path=new_path, params="", query="", fragment=""))

    def _steer_url(self) -> str:
        return self._resolve_steer_endpoint("steer")

    def _steer_batch_url(self) -> str:
        return self._resolve_steer_endpoint("steer_batch")

    def _build_steer_payload(
        self,
        *,
        trait_dict: dict[str, int],
        conversation: list[ChatCompletionMessageParam],
        combo: SteerCombination,
    ) -> dict[str, object]:
        user_characteristics = self._user_characteristics_payload(combo)
        messages = [
            {"role": str(msg["role"]), "content": str(msg["content"])} for msg in conversation
        ]

        payload: dict[str, object] = {
            "trait_dict": {k: int(v) for k, v in trait_dict.items()},
            "user_characteristics": user_characteristics,
            "messages": messages,
            "temperature": float(self.steer_temperature),
            "max_tokens": int(self.steer_max_tokens),
            "seed": int(self.steer_seed),
        }
        return payload

    def _user_characteristics_payload(self, combo: SteerCombination) -> dict[str, object]:
        payload: dict[str, object] = {}

        age = self._coerce_age(combo.age)
        if age is not None:
            payload["age"] = age

        optional_fields: dict[str, str | None] = {
            "gender": combo.gender,
            "occupation": combo.occupation,
            "location": combo.location,
            "language": combo.language,
            "intent": combo.intent,
            "task": combo.task,
        }
        for field, raw in optional_fields.items():
            value = self._normalize_optional_str(raw)
            if value is not None:
                payload[field] = value

        return payload

    @staticmethod
    def _coerce_age(value: int | str | None) -> int | None:
        if value is None:
            return None
        if isinstance(value, int):
            return value
        if isinstance(value, str):
            stripped = value.strip()
            if not stripped:
                return None
            value = stripped
        try:
            return int(value)
        except (TypeError, ValueError) as exc:
            raise ValueError("age must be an integer for Steer API requests") from exc

    @staticmethod
    def _normalize_optional_str(value: str | None) -> str | None:
        if value is None:
            return None
        stripped = value.strip()
        return stripped or None

    async def _request_steer(
        self, url: str, headers: dict[str, str], payload: object
    ) -> tuple[httpx.Response | None, str | None]:
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                self.logger.debug("POST %s (API-Key present=%s)", url, "API-Key" in headers)
                resp = await client.post(url, headers=headers, json=payload)
                return resp, None
        except Exception as e:
            self.logger.exception("User service error")
            return (
                None,
                "Error: Steer API call failed. "
                f"Details: {e!s}. Check STEER_API_KEY and COLLINEAR_STEER_URL.",
            )

    def _should_skip_trait_for_422(self, trait: str) -> tuple[bool, list[str]]:
        available = self.list_traits()
        if not available:
            return True, []
        return (trait not in set(available)), available

    def _handle_unprocessable_or_skip(self, trait: str, _resp: httpx.Response) -> None:
        should_skip, available = self._should_skip_trait_for_422(trait)
        if should_skip:
            avail_str = ", ".join(sorted(available)) if available else "<unavailable>"
            self.logger.warning(
                "Trait '%s' not recognized by Steer API (422). Available traits: [%s]. "
                "Skipping this combination.",
                trait,
                avail_str,
            )
            raise SimulationRunner.InvalidTraitError(trait)

    def _handle_unprocessable_or_skip_mixed(self, traits: set[str], _resp: httpx.Response) -> None:
        available = set(self.list_traits())
        if not available:
            raise SimulationRunner.InvalidTraitError(",".join(sorted(traits)))
        missing = [t for t in traits if t not in available]
        if missing:
            avail_str = ", ".join(sorted(available))
            self.logger.warning(
                "Traits '%s' not recognized by Steer API (422). Available traits: [%s]. "
                "Skipping this combination.",
                ", ".join(missing),
                avail_str,
            )
            raise SimulationRunner.InvalidTraitError(",".join(sorted(traits)))

    def _read_json_or_error(self, resp: httpx.Response) -> tuple[object | None, str | None]:
        try:
            resp.raise_for_status()
            raw: object = resp.json()
        except Exception as e:
            self.logger.exception("User service error")
            return (
                None,
                "Error: Steer API call failed. "
                f"Details: {e!s}. Check STEER_API_KEY and COLLINEAR_STEER_URL.",
            )
        return raw, None

    async def _call_collinear_steer_api(
        self,
        *,
        trait: str,
        intensity: int,
        combo: SteerCombination,
        conversation: list[ChatCompletionMessageParam] | None = None,
        _conversation: list[ChatCompletionMessageParam] | None = None,
        use_batch: bool | None = None,
    ) -> str:
        return await self._call_collinear_steer_api_trait_dict(
            trait_dict={trait: int(intensity)},
            combo=combo,
            conversation=conversation,
            _conversation=_conversation,
            use_batch=use_batch,
        )

    async def _call_collinear_steer_api_trait_dict(
        self,
        *,
        trait_dict: dict[str, int],
        combo: SteerCombination,
        conversation: list[ChatCompletionMessageParam] | None = None,
        _conversation: list[ChatCompletionMessageParam] | None = None,
        use_batch: bool | None = None,
    ) -> str:
        conv = conversation if conversation is not None else (_conversation or [])
        payload: dict[str, object] = self._build_steer_payload(
            trait_dict=trait_dict,
            conversation=conv,
            combo=combo,
        )
        headers: dict[str, str] = {
            "Content-Type": "application/json",
            "API-Key": self.steer_api_key,
        }

        effective_batch = use_batch if use_batch is not None else (self._steer_batch_size > 1)
        if effective_batch and self._batcher is not None:
            return await self._batcher.submit(payload, headers, trait_dict)

        return await self._call_single_endpoint(payload, headers, trait_dict)

    async def _call_single_endpoint(
        self,
        payload: dict[str, object],
        headers: dict[str, str],
        trait_dict: dict[str, int],
    ) -> str:
        url = self._steer_url()
        resp, err = await self._request_steer(url, headers, payload)
        if err is not None:
            return err
        if resp is None:
            return "Error: Steer API call failed."

        self._log_if_unauthorized(resp)

        if resp.status_code == HTTP_UNPROCESSABLE:
            if len(trait_dict) == 1:
                trait = next(iter(trait_dict))
                self._handle_unprocessable_or_skip(trait, resp)
            else:
                self._handle_unprocessable_or_skip_mixed(set(trait_dict.keys()), resp)

        raw, err2 = self._read_json_or_error(resp)
        if err2 is not None:
            return err2

        data: dict[str, object] | None = (
            cast("dict[str, object]", raw) if isinstance(raw, dict) else None
        )

        value: object | None = data.get("response") if data is not None else None

        if isinstance(value, str):
            return value
        return "" if value is None else str(value)

    async def call_single_steer(
        self,
        payload: dict[str, object],
        *,
        headers: dict[str, str],
        trait_dict: dict[str, int],
    ) -> str:
        """Expose single-request pathway for batching helpers."""
        return await self._call_single_endpoint(payload, headers, trait_dict)

    async def call_batch_steer(
        self,
        payloads: list[dict[str, object]],
        *,
        headers: dict[str, str],
    ) -> tuple[bool, list[str] | str]:
        """Expose batch-request pathway for batching helpers."""
        return await self._call_batch_endpoint(payloads, headers=headers)

    async def _call_batch_endpoint(
        self,
        payloads: list[dict[str, object]],
        *,
        headers: dict[str, str],
    ) -> tuple[bool, list[str] | str]:
        if not payloads:
            return True, []
        url = self._steer_batch_url()
        resp, err = await self._request_steer(url, headers, payloads)
        error: str | None = None
        if err is not None:
            error = err
        elif resp is None:
            error = "Error: Steer batch API call failed."
        else:
            self._log_if_unauthorized(resp)
            if resp.status_code == HTTP_UNPROCESSABLE:
                error = "422 Unprocessable Entity"
            else:
                raw, err2 = self._read_json_or_error(resp)
                if err2 is not None:
                    error = err2
                else:
                    responses = self._extract_batch_responses(raw)
                    if responses is None:
                        self.logger.warning(
                            "Unexpected payload from steer_batch endpoint: %r",
                            raw,
                        )
                        error = "Unexpected response payload"
                    elif len(responses) != len(payloads):
                        self.logger.warning(
                            "Batch response count mismatch (expected %d, got %d).",
                            len(payloads),
                            len(responses),
                        )
                        error = "Batch response count mismatch"
                    else:
                        return True, responses

        return False, error or "Error: Steer batch API call failed."

    def _extract_batch_responses(self, raw: object) -> list[str] | None:
        entries = self._batch_entries(raw)
        if entries is None:
            return None
        return [self._normalize_batch_entry(entry) for entry in entries]

    def _batch_entries(self, raw: object) -> list[object] | None:
        if isinstance(raw, list):
            return raw
        if isinstance(raw, dict):
            for key in ("responses", "response", "data"):
                if key in raw:
                    value = raw[key]
                    if isinstance(value, list):
                        return value
                    return [value]
            return None
        return [raw]

    def _normalize_batch_entry(self, entry: object) -> str:
        if isinstance(entry, dict):
            for key in ("response", "text", "message"):
                if key in entry:
                    return self._stringify_batch_value(entry[key])
        return self._stringify_batch_value(entry)

    @staticmethod
    def _stringify_batch_value(value: object) -> str:
        if isinstance(value, str):
            return value
        if value is None:
            return ""
        return str(value)

    async def _build_conversation(
        self, combo: SteerCombination, num_exchanges: int
    ) -> tuple[list[ChatCompletionMessageParam], str]:
        """Build a conversation with specified number of exchanges.

        Each exchange consists of one user turn followed by one assistant turn.
        The final assistant turn uses the actual model being tested.
        """
        conversation: list[ChatCompletionMessageParam] = []
        total_turns = num_exchanges * 2
        final_response: str = ""
        completed_user_turns = 0
        try:
            for turn in range(1, total_turns + 1):
                is_user_turn = turn % 2 == 1
                role = Role.USER if is_user_turn else Role.ASSISTANT
                if is_user_turn:
                    try:
                        response = await self._generate_turn(combo, conversation, role=role)
                    finally:
                        completed_user_turns += 1
                        self._advance_progress(1)
                    u: ChatCompletionUserMessageParam = {
                        "role": Role.USER.value,
                        "content": response,
                    }
                    conversation.append(u)
                else:
                    response = await self._generate_turn(combo, conversation, role=role)
                    a: ChatCompletionAssistantMessageParam = {
                        "role": Role.ASSISTANT.value,
                        "content": response,
                    }
                    conversation.append(a)
                    if turn == total_turns:
                        final_response = response
        except SimulationRunner.InvalidTraitError as e:
            raise SimulationRunner.BuildConversationError(
                completed_user_turns,
                invalid_trait=True,
                trait=e.trait,
            ) from e
        except Exception as e:
            raise SimulationRunner.BuildConversationError(completed_user_turns) from e
        else:
            return conversation, final_response

    def list_traits(self) -> list[str]:
        """Return available traits from the Steer service.

        Resolves the traits endpoint to the same host as the Steer URL. If
        ``COLLINEAR_STEER_URL`` is set, its scheme/host are used and the path
        is replaced with ``/traits``. Otherwise defaults to
        ``https://steer.collinear.com/traits``.

        Network errors or unexpected payloads result in an empty list.
        """
        try:
            steer_url = os.getenv(
                "COLLINEAR_STEER_URL",
                "https://steer.collinear.ai/steer",
            )
            if steer_url:
                parsed = urlparse(steer_url)
                base = parsed._replace(path="/traits", params="", query="", fragment="")
                traits_url = urlunparse(base)
            else:
                traits_url = "https://steer.collinear.ai/traits"

            with httpx.Client(timeout=self.timeout) as client:
                self.logger.debug("GET %s", traits_url)
                resp = client.get(
                    traits_url,
                    headers={
                        "API-Key": self.steer_api_key,
                    },
                )
                resp.raise_for_status()
                raw: object = resp.json()
        except Exception:
            self.logger.exception("Failed to fetch traits")
            return []

        if isinstance(raw, dict):
            traits = raw.get("traits")
            if isinstance(traits, list):
                return [str(t) for t in traits if isinstance(t, (str, bytes))]
        if isinstance(raw, list):
            return [str(t) for t in raw if isinstance(t, (str, bytes))]
        return []
