"""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 typing import TYPE_CHECKING
from urllib.parse import urljoin

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

from collinear.clients.assistant_client import AssistantClient
from collinear.clients.steering_client import SteeringClient
from collinear.clients.steering_client import normalize_base_url
from collinear.common.exceptions import BuildConversationError
from collinear.common.exceptions import InvalidTraitError
from collinear.schemas.traitmix import Role
from collinear.schemas.traitmix import SimulationResult
from collinear.schemas.traitmix import TraitMixCombination
from collinear.schemas.traitmix import TraitMixConfig
from collinear.simulate import progress
from collinear.simulate.conversation import DEFAULT_ASSISTANT_PROMPT
from collinear.simulate.conversation import should_stop
from collinear.simulate.payloads import build_traitmix_payload
from collinear.simulate.payloads import user_characteristics_from_combo

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


@dataclass(frozen=True)
class AssistantOptions:
    """Per-call overrides for assistant completions."""

    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 = DEFAULT_ASSISTANT_PROMPT
    InvalidTraitError = InvalidTraitError
    BuildConversationError = BuildConversationError

    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_handle: progress.Progress | None = None
        # Initialize helper clients (internal)
        self.assistant_client = AssistantClient(
            base_url=assistant_model_url,
            api_key=assistant_model_api_key,
            model=assistant_model_name,
            timeout=timeout,
            max_retries=max_retries,
            rate_limit_retries=rate_limit_retries,
            logger=self.logger,
        )
        # Backward compatibility: expose underlying OpenAI client as before.
        self.client = self.assistant_client.client
        self.steering_client = SteeringClient(api_key=collinear_api_key, 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.track_progress(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]:
        """Temporarily override traitmix generation parameters."""
        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]:
        """Temporarily override assistant completion parameters."""
        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 track_progress(self, *, enabled: bool, total: int) -> Iterator[None]:
        """Manage the lifecycle of the progress indicator."""
        with progress.progress_manager(enabled=enabled, total=total, logger=self.logger) as handle:
            self.progress_handle = handle
            try:
                yield
            finally:
                self.progress_handle = None

    def advance_progress(self, step: int) -> None:
        """Advance the progress bar when user turns complete."""
        if step <= 0:
            return
        active_progress = self.progress_handle
        if active_progress is None:
            return
        active_progress.update(step)

    def adjust_progress_total(self, decrement: int) -> None:
        """Shrink the remaining total when work is skipped."""
        if decrement <= 0:
            return
        active_progress = self.progress_handle
        if active_progress is None:
            return
        active_progress.adjust_total(decrement)

    def select_samples(
        self, combinations: list[TraitMixCombination], k: int | None
    ) -> list[TraitMixCombination]:
        """Select which trait combinations will be executed."""
        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:
        """Public accessor for the computed concurrency limit."""
        return min(MAX_ALLOWED_CONCURRENCY, max(1, max_concurrency))

    async def execute_samples(
        self,
        samples: list[TraitMixCombination],
        num_exchanges: int,
        batch_delay: float,
        max_concurrency: int,
    ) -> list[SimulationResult]:
        """Execute the selected combinations with optional concurrency."""
        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 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 via AssistantClient."""
        try:
            opts = ASSISTANT_OPTS.get()
        except Exception:
            opts = None
        try:
            return await self.assistant_client.complete(
                messages,
                system_prompt,
                max_tokens=(opts.max_tokens if opts is not None else None),
                seed=(opts.seed if opts is not None else None),
            )
        except openai.RateLimitError:
            self.logger.exception("Error getting response")
            raise
        except Exception as exc:
            self.logger.exception("Error getting response")
            return f"Error: {exc!s}"

    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:
        """Return a masked preview of the Collinear API key for logging."""
        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:
        """Log unauthorized responses while masking the API key."""
        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:
        """Log unauthorized responses when appropriate."""
        if resp.status_code == HTTP_UNAUTHORIZED:
            self.log_unauthorized(resp)

    def resolve_traitmix_endpoint(self, endpoint: str) -> str:
        """Resolve an API endpoint relative to the configured base URL."""
        base = os.getenv("COLLINEAR_TRAITMIX_BASE_URL") or os.getenv("COLLINEAR_TRAITMIX_URL")
        base_url = normalize_base_url(base)
        normalized_endpoint = endpoint.lstrip("/")
        return urljoin(base_url, normalized_endpoint)

    def parse_batch_responses(
        self,
        resp: httpx.Response,
        expected_count: int,
    ) -> list[str]:
        """Parse the Steering batch response into a list of strings."""
        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:
        """Decide whether to skip payloads when the API returns 422."""
        collected_traits: set[str] = set()
        first_trait_dict: dict[str, object] | None = None
        if payloads:
            first_payload = payloads[0]
            if isinstance(first_payload, dict):
                trait_dict_obj = first_payload.get("trait_dict")
                if isinstance(trait_dict_obj, dict):
                    first_trait_dict = trait_dict_obj
        for payload in payloads:
            trait_dict = payload.get("trait_dict") if isinstance(payload, dict) else None
            if isinstance(trait_dict, dict):
                collected_traits.update(str(name) for name in trait_dict)

        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)

    async def request_traitmix(
        self, url: str, headers: dict[str, str], payload: object
    ) -> tuple[httpx.Response | None, str | None]:
        """Send the raw Steering API request."""
        try:
            resp, err = await self.steering_client.post_json(url, headers, payload)
        except Exception:
            self.logger.exception("User service error")
            raise
        if resp is None and err is not None:
            self.logger.error("User service error: %s", err)
        return resp, err

    def should_skip_trait_for_422(self, trait: str) -> tuple[bool, list[str]]:
        """Check whether a failing trait should be skipped after a 422."""
        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:
        """Log skipped traits after a 422 response."""
        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 InvalidTraitError(trait)

    def handle_unprocessable_or_skip_mixed(self, traits: set[str], _resp: httpx.Response) -> None:
        """Log skipped trait combinations after a 422 response."""
        available = set(self.list_traits())
        if not available:
            raise 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 InvalidTraitError(",".join(sorted(traits)))

    def read_json_or_error(self, resp: httpx.Response) -> tuple[object | None, str | None]:
        """Attempt to read JSON, returning an error string on failure."""
        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

    def user_characteristics_payload(self, combo: TraitMixCombination) -> dict[str, object]:
        """Public helper for extracting persona fields."""
        return user_characteristics_from_combo(combo)

    def _user_characteristics_payload(self, combo: TraitMixCombination) -> dict[str, object]:
        """Backward-compatible alias for notebooks that patched internals."""
        return self.user_characteristics_payload(combo)

    async def call_collinear_traitmix_api(
        self,
        *,
        trait: str,
        intensity: int | str,
        combo: TraitMixCombination,
        conversation: list[ChatCompletionMessageParam] | None = None,
    ) -> str:
        """Call the Steering API for a single trait/intensity."""
        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:
        """Call the Steering API with the provided trait dictionary."""
        conv = conversation if conversation is not None else []
        payload = build_traitmix_payload(
            trait_dict=trait_dict,
            conversation=conv,
            combo=combo,
            temperature=self.traitmix_temperature,
            max_tokens=self.traitmix_max_tokens,
            seed=self.traitmix_seed,
        )
        headers = {
            "Content-Type": "application/json",
            "API-Key": self.collinear_api_key,
        }
        url = self.resolve_traitmix_endpoint("steer_batch")
        responses = await self.call_batch_endpoint(url, [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,
        url: str,
        payloads: list[dict[str, object]],
        *,
        headers: dict[str, str],
    ) -> list[str]:
        """Invoke the Steering batch endpoint and return string responses."""
        if not payloads:
            return []
        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]:
        """Generate the next user turn and return whether to stop."""
        response = await self.generate_turn(combo, conversation, role=Role.USER)
        message: ChatCompletionUserMessageParam = {
            "role": Role.USER.value,
            "content": response,
        }
        conversation.append(message)
        stop_here = should_stop(response)
        if stop_here:
            conversation.pop()
        return stop_here, response

    async def run_assistant_turn(
        self,
        combo: TraitMixCombination,
        conversation: list[ChatCompletionMessageParam],
    ) -> str:
        """Generate the next assistant turn and append it to the dialogue."""
        response = await self.generate_turn(combo, conversation, role=Role.ASSISTANT)
        message: ChatCompletionAssistantMessageParam = {
            "role": Role.ASSISTANT.value,
            "content": response,
        }
        conversation.append(message)
        return 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 InvalidTraitError as e:
            raise BuildConversationError(
                completed_user_turns,
                invalid_trait=True,
                trait=e.trait,
            ) from e
        except Exception as e:
            raise BuildConversationError(completed_user_turns) from e
        else:
            if should_stop(final_response) or 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 using ``COLLINEAR_TRAITMIX_BASE_URL`` (or the
        legacy ``COLLINEAR_TRAITMIX_URL``) before falling back to the default
        production host.

        Network errors or unexpected payloads result in an empty list.
        """
        try:
            return self.steering_client.list_traits()
        except Exception:
            self.logger.exception("Failed to fetch traits")
            return []
