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

from __future__ import annotations

import asyncio
import logging
import os
import random
from contextlib import contextmanager
from contextvars import ContextVar
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 urljoin
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.traitmix import Role
from collinear.schemas.traitmix import SimulationResult
from collinear.schemas.traitmix import TraitMixCombination
from collinear.schemas.traitmix import TraitMixConfig

if TYPE_CHECKING:
    from collections.abc import Iterator

HTTP_UNAUTHORIZED = 401
HTTP_UNPROCESSABLE = 422
MIN_MASK_VISIBLE = 4
DEFAULT_TRAITMIX_SEED = -1
MAX_ALLOWED_CONCURRENCY = 8
DEFAULT_CONCURRENCY = 1


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()


@dataclass(frozen=True)
class _AssistantOptions:
    max_tokens: int | None
    seed: int | None


_ASSISTANT_OPTS: ContextVar[_AssistantOptions | None] = ContextVar("ASSISTANT_OPTS", default=None)


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

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

    ASSISTANT_PROMPT_TEMPLATE = (
        "You are the ASSISTANT. You are a helpful, respectful, and succinct "
        "customer support assistant.\n\n"
        "Respond only to the customer's most recent message. Write only the "
        "assistant's next message as plain text. Do not include role names, quotes, "
        "or any markup. Avoid lists unless the customer explicitly asks for step-by-step "
        "instructions. Keep the reply under 150 words.\n\n"
        "If key details are missing, ask one brief, specific follow-up question. "
        "If you are unsure, say so and suggest a practical next step."
    )

    def __init__(
        self,
        assistant_model_url: str,
        assistant_model_api_key: str,
        assistant_model_name: str,
        *,
        collinear_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.
            collinear_api_key: Collinear API key sent as the ``API-Key`` header to
                the traitmix service.
            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 collinear_api_key:
            raise ValueError("COLLINEAR_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.collinear_api_key = collinear_api_key
        self.timeout = timeout
        self.max_retries = max_retries
        self.rate_limit_retries = rate_limit_retries
        self.logger = logging.getLogger("collinear")
        self.traitmix_temperature: float = 0.7
        self.traitmix_max_tokens: int = 256
        self.traitmix_seed: int = DEFAULT_TRAITMIX_SEED
        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: TraitMixConfig,
        k: int | None = None,
        num_exchanges: int = 2,
        batch_delay: float = 0.1,
        *,
        traitmix_temperature: float | None = None,
        traitmix_max_tokens: int | None = None,
        traitmix_seed: int | None = None,
        assistant_max_tokens: int | None = None,
        assistant_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: TraitMix 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).
            traitmix_temperature: Optional sampling temperature for traitmix generation
                (default 0.7).
            traitmix_max_tokens: Optional max tokens for traitmix generation (default 256).
            traitmix_seed: Optional deterministic seed for traitmix sampling. ``-1``
                (default) delegates randomness to the service.
            assistant_max_tokens: Optional max tokens for assistant responses
                (sent only if provided).
            assistant_seed: Optional deterministic seed for the assistant model
                (sent only if provided).
            mix_traits: If True, use pairwise mixing (exactly two traits per traitmix).
                Requires at least two traits with available levels.
            progress: Whether to display a tqdm-style bar tracking traitmix 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,
                traitmix_temperature=traitmix_temperature,
                traitmix_max_tokens=traitmix_max_tokens,
                traitmix_seed=traitmix_seed,
                assistant_max_tokens=assistant_max_tokens,
                assistant_seed=assistant_seed,
                mix_traits=mix_traits,
                progress=progress,
                max_concurrency=max_concurrency,
            )
        )

    async def run_async(
        self,
        config: TraitMixConfig,
        k: int | None = None,
        num_exchanges: int = 2,
        batch_delay: float = 0.1,
        *,
        traitmix_temperature: float | None = None,
        traitmix_max_tokens: int | None = None,
        traitmix_seed: int | None = None,
        assistant_max_tokens: int | None = None,
        assistant_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: TraitMix 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).
            traitmix_temperature: Optional sampling temperature for traitmix generation
                (default 0.7).
            mix_traits: If True, use pairwise mixing (exactly two traits per traitmix).
                Requires at least two traits with available levels.
            traitmix_max_tokens: Optional max tokens for traitmix generation (default 256).
            traitmix_seed: Optional deterministic seed for traitmix sampling. ``-1``
                (default) delegates randomness to the service.
            assistant_max_tokens: Optional max tokens for assistant responses
                (sent only if provided).
            assistant_seed: Optional deterministic seed for the assistant model
                (sent only if provided).
            progress: Whether to display a tqdm-style bar tracking traitmix 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._traitmix_settings(
                traitmix_temperature=traitmix_temperature,
                traitmix_max_tokens=traitmix_max_tokens,
                traitmix_seed=traitmix_seed,
            ),
            self._assistant_settings(
                assistant_max_tokens=assistant_max_tokens,
                assistant_seed=assistant_seed,
            ),
        ):
            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 _traitmix_settings(
        self,
        *,
        traitmix_temperature: float | None,
        traitmix_max_tokens: int | None,
        traitmix_seed: int | None,
    ) -> Iterator[None]:
        prev_temp = self.traitmix_temperature
        prev_max = self.traitmix_max_tokens
        prev_seed = self.traitmix_seed
        try:
            if traitmix_temperature is not None:
                self.traitmix_temperature = float(traitmix_temperature)
            if traitmix_max_tokens is not None:
                self.traitmix_max_tokens = int(traitmix_max_tokens)
            if traitmix_seed is not None:
                self.traitmix_seed = int(traitmix_seed)
            yield
        finally:
            self.traitmix_temperature = prev_temp
            self.traitmix_max_tokens = prev_max
            self.traitmix_seed = prev_seed

    @contextmanager
    def _assistant_settings(
        self,
        *,
        assistant_max_tokens: int | None,
        assistant_seed: int | None,
    ) -> Iterator[None]:
        token = _ASSISTANT_OPTS.set(
            _AssistantOptions(max_tokens=assistant_max_tokens, seed=assistant_seed)
        )
        try:
            yield
        finally:
            _ASSISTANT_OPTS.reset(token)

    class EmptyTraitMixResponseError(RuntimeError):
        """Raised when the TraitMix API returns empty content."""

    @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[TraitMixCombination], k: int | None
    ) -> list[TraitMixCombination]:
        total = len(combinations)
        self.logger.info("Total traitmix combinations: %d", total)
        if total == 0:
            self.logger.warning("No traitmix 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[TraitMixCombination],
        num_exchanges: int,
        batch_delay: float,
        max_concurrency: int,
    ) -> list[SimulationResult]:
        sem = asyncio.Semaphore(self._calculate_semaphore_limit(max_concurrency))

        async def run_one(i: int, combo: TraitMixCombination) -> SimulationResult | 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
                    )
                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)}")
                    return None
                else:
                    result = SimulationResult(
                        conv_prefix=conversation[:-1],
                        response=final_response,
                        traitmix=combo,
                    )
                    self.logger.info(f"Completed simulation {i + 1}/{len(samples)}")
                    return result

        results = await asyncio.gather(*(run_one(i, combo) for i, combo in enumerate(samples)))
        return [result for result in results if result is not None]

    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:
                opts = _ASSISTANT_OPTS.get()
                response = await self.client.chat.completions.create(
                    model=self.assistant_model_name,
                    messages=full_messages,
                    temperature=0.8,
                    max_tokens=(
                        opts.max_tokens
                        if (opts is not None and opts.max_tokens is not None)
                        else openai.NOT_GIVEN
                    ),
                    seed=(
                        opts.seed
                        if (opts is not None and opts.seed is not None)
                        else openai.NOT_GIVEN
                    ),
                )
            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: TraitMixCombination,
        conversation: list[ChatCompletionMessageParam],
        role: Role,
    ) -> str:
        """Generate a single turn in the conversation."""
        if role is Role.USER:
            self.logger.info("Generating USER turn")
            if len(combo.traits) == 1:
                trait = next(iter(combo.traits))
                intensity = next(iter(combo.traits.values()))
                response = await self._call_collinear_traitmix_api(
                    trait=trait,
                    intensity=intensity,
                    combo=combo,
                    conversation=conversation,
                )
            else:
                response = await self._call_collinear_traitmix_api_trait_dict(
                    trait_dict={k: str(v) for k, v in combo.traits.items()},
                    combo=combo,
                    conversation=conversation,
                )
        else:
            self.logger.info("Generating ASSISTANT turn")
            system_prompt = self.ASSISTANT_PROMPT_TEMPLATE
            max_empty_retries = 2
            attempts = 0
            response = ""
            while attempts <= max_empty_retries:
                candidate = await self._call_with_retry(conversation, system_prompt)
                attempts += 1
                if candidate and candidate.strip():
                    response = candidate
                    break
                self.logger.warning(
                    "Assistant returned empty response (attempt %s/%s)",
                    attempts,
                    max_empty_retries + 1,
                )
            if not response.strip():
                response = "I'm sorry, I don't have anything to add right now."
        self.logger.info(response)

        return response

    def _mask_key_preview(self) -> str:
        key = (self.collinear_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(
            "TraitMix 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 TraitMix 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 TraitMix 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_traitmix_endpoint(self, endpoint: str) -> str:
        base = os.getenv("COLLINEAR_TRAITMIX_URL", "https://steer.collinear.ai/")
        if not base:
            base = "https://steer.collinear.ai/"
        parsed = urlparse(base)
        normalized_endpoint = endpoint.lstrip("/")
        stripped = parsed._replace(params="", query="", fragment="")
        base_url = urlunparse(stripped)
        if not base_url.endswith("/"):
            base_url = f"{base_url}/"
        return urljoin(base_url, normalized_endpoint)

    def _traitmix_batch_url(self) -> str:
        return self._resolve_traitmix_endpoint("steer_batch")

    def _traitmix_headers(self) -> dict[str, str]:
        return {
            "Content-Type": "application/json",
            "API-Key": self.collinear_api_key,
        }

    def _build_traitmix_payload(
        self,
        *,
        trait_dict: dict[str, int | str],
        conversation: list[ChatCompletionMessageParam],
        combo: TraitMixCombination,
    ) -> dict[str, object]:
        user_characteristics = self._user_characteristics_payload(combo)
        messages = self._transform_conversation_for_traitmix(conversation)

        def _to_api_label(value: int | str) -> str:
            mapping = {0: "low", 1: "medium", 2: "high"}
            if isinstance(value, int):
                if value not in mapping:
                    raise ValueError(f"Trait level int must be 0,1,2; got {value!r}")
                return mapping[value]
            raw = str(value).strip()
            try:
                iv = int(raw)
            except Exception as err:
                lv = raw.lower()
                if lv in {"low", "medium", "high"}:
                    return lv
                message = f"Unknown trait level: {value!r}. Use 0/1/2 or low/medium/high."
                raise ValueError(message) from err
            else:
                if iv not in mapping:
                    raise ValueError(f"Trait level string-int must be 0,1,2; got {raw!r}")
                return mapping[iv]

        api_trait_dict: dict[str, str] = {k: _to_api_label(v) for k, v in trait_dict.items()}

        payload: dict[str, object] = {
            "trait_dict": api_trait_dict,
            "user_characteristics": user_characteristics,
            "messages": messages,
            "temperature": float(self.traitmix_temperature),
            "max_tokens": int(self.traitmix_max_tokens),
            "seed": int(self.traitmix_seed),
        }
        return payload

    def _transform_conversation_for_traitmix(
        self, conversation: list[ChatCompletionMessageParam]
    ) -> list[dict[str, object]]:
        transformed: list[dict[str, object]] = []
        for msg in conversation:
            raw = dict(msg)
            role = str(raw.get("role"))
            if role == Role.USER.value:
                swapped_role = Role.ASSISTANT.value
            elif role == Role.ASSISTANT.value:
                swapped_role = Role.USER.value
            else:
                swapped_role = role

            content = raw.get("content")
            transformed.append(
                {
                    "role": swapped_role,
                    "content": "" if content is None else str(content),
                }
            )
        return transformed

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

        if combo.age is not None:
            payload["age"] = combo.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

    def _collect_traits_from_payloads(self, payloads: list[dict[str, object]]) -> set[str]:
        traits: set[str] = set()
        for payload in payloads:
            trait_dict = payload.get("trait_dict") if isinstance(payload, dict) else None
            if isinstance(trait_dict, dict):
                traits.update(str(name) for name in trait_dict)
        return traits

    def _first_trait_dict(self, payloads: list[dict[str, object]]) -> dict[str, object] | None:
        if not payloads:
            return None
        first = payloads[0]
        if not isinstance(first, dict):
            return None
        trait_dict = first.get("trait_dict")
        return trait_dict if isinstance(trait_dict, dict) else None

    def _parse_batch_responses(
        self,
        resp: httpx.Response,
        expected_count: int,
    ) -> list[str]:
        raw, err = self._read_json_or_error(resp)
        if err is not None:
            raise RuntimeError(err)

        if not isinstance(raw, dict):
            raise TypeError("Unexpected response payload")

        responses = raw.get("responses")
        if not isinstance(responses, list):
            raise TypeError("Unexpected response payload")

        normalized: list[str] = []
        for entry in responses:
            if isinstance(entry, str):
                normalized.append(entry)
            elif entry is None:
                normalized.append("")
            else:
                normalized.append(str(entry))

        if len(normalized) != expected_count:
            message = (
                f"Batch response count mismatch (expected {expected_count}, got {len(normalized)})."
            )
            raise RuntimeError(message)

        return normalized

    def _handle_unprocessable_payloads(
        self,
        payloads: list[dict[str, object]],
        resp: httpx.Response,
    ) -> None:
        collected_traits = self._collect_traits_from_payloads(payloads)
        first_trait_dict = self._first_trait_dict(payloads)

        if len(payloads) == 1 and isinstance(first_trait_dict, dict):
            if len(first_trait_dict) == 1:
                trait = next(iter(first_trait_dict))
                self._handle_unprocessable_or_skip(str(trait), resp)
                return
            if collected_traits:
                self._handle_unprocessable_or_skip_mixed(collected_traits, resp)
                return

        if collected_traits:
            self._handle_unprocessable_or_skip_mixed(collected_traits, resp)

    @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_traitmix(
        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: TraitMix API call failed. "
                f"Details: {e!s}. Check COLLINEAR_API_KEY and COLLINEAR_TRAITMIX_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 TraitMix 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 TraitMix 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: TraitMix API call failed. "
                f"Details: {e!s}. Check COLLINEAR_API_KEY and COLLINEAR_TRAITMIX_URL.",
            )
        return raw, None

    async def _call_collinear_traitmix_api(
        self,
        *,
        trait: str,
        intensity: int | str,
        combo: TraitMixCombination,
        conversation: list[ChatCompletionMessageParam] | None = None,
    ) -> str:
        return await self._call_collinear_traitmix_api_trait_dict(
            trait_dict={trait: intensity},
            combo=combo,
            conversation=conversation,
        )

    async def _call_collinear_traitmix_api_trait_dict(
        self,
        *,
        trait_dict: dict[str, int | str],
        combo: TraitMixCombination,
        conversation: list[ChatCompletionMessageParam] | None = None,
    ) -> str:
        conv = conversation if conversation is not None else []
        payload: dict[str, object] = self._build_traitmix_payload(
            trait_dict=trait_dict,
            conversation=conv,
            combo=combo,
        )
        headers = self._traitmix_headers()
        responses = await self._call_batch_endpoint([payload], headers=headers)

        if not responses:
            raise SimulationRunner.EmptyTraitMixResponseError("TraitMix API returned no responses")

        response = responses[0].strip()
        if not response:
            raise SimulationRunner.EmptyTraitMixResponseError(
                "TraitMix API returned empty response",
            )
        return response

    async def _call_batch_endpoint(
        self,
        payloads: list[dict[str, object]],
        *,
        headers: dict[str, str],
    ) -> list[str]:
        if not payloads:
            return []
        url = self._traitmix_batch_url()
        resp, err = await self._request_traitmix(url, headers, payloads)
        if err is not None:
            raise RuntimeError(err)
        if resp is None:
            raise RuntimeError("Error: TraitMix batch API call failed.")

        self._log_if_unauthorized(resp)

        if resp.status_code == HTTP_UNPROCESSABLE:
            self._handle_unprocessable_payloads(payloads, resp)
            raise RuntimeError("TraitMix API returned 422 Unprocessable Entity")

        return self._parse_batch_responses(resp, len(payloads))

    async def _run_user_turn(
        self,
        combo: TraitMixCombination,
        conversation: list[ChatCompletionMessageParam],
    ) -> tuple[bool, str]:
        response = await self._generate_turn(combo, conversation, role=Role.USER)
        message: ChatCompletionUserMessageParam = {
            "role": Role.USER.value,
            "content": response,
        }
        conversation.append(message)
        should_stop = self._should_stop(response)
        if should_stop:
            conversation.pop()
        return should_stop, response

    async def _run_assistant_turn(
        self,
        combo: TraitMixCombination,
        conversation: list[ChatCompletionMessageParam],
    ) -> str:
        response = await self._generate_turn(combo, conversation, role=Role.ASSISTANT)
        message: ChatCompletionAssistantMessageParam = {
            "role": Role.ASSISTANT.value,
            "content": response,
        }
        conversation.append(message)
        return response

    def _should_stop(self, response: str) -> bool:
        return response == "###STOP###" or "###STOP###" in response

    async def _build_conversation(
        self, combo: TraitMixCombination, 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 = ""
        completed_user_turns = 0
        last_assistant_response = ""
        last_response = ""
        try:
            for turn in range(1, total_turns + 1):
                if turn % 2 == 1:
                    try:
                        stopped, last_response = await self._run_user_turn(combo, conversation)
                    finally:
                        completed_user_turns += 1
                        self._advance_progress(1)
                    if stopped:
                        final_response = last_assistant_response
                        break
                else:
                    last_assistant_response = await self._run_assistant_turn(combo, conversation)
                    last_response = last_assistant_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:
            if self._should_stop(final_response) or self._should_stop(last_response):
                final_response = last_assistant_response
            if not final_response and last_assistant_response:
                final_response = last_assistant_response
            return conversation, final_response

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

        Resolves the traits endpoint to the same host as the TraitMix API by reusing
        ``COLLINEAR_TRAITMIX_URL``. When unset, defaults to
        ``https://traitmix.collinear.com/traits``.

        Network errors or unexpected payloads result in an empty list.
        """
        try:
            traitmix_url = os.getenv("COLLINEAR_TRAITMIX_URL")
            if traitmix_url:
                traits_url = self._resolve_traitmix_endpoint("traits")
            else:
                traits_url = "https://traitmix.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.collinear_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 []
