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

from __future__ import annotations

import logging
import os
import random
import time
from contextlib import contextmanager
from typing import TYPE_CHECKING
from typing import cast
from urllib.parse import urlparse
from urllib.parse import urlunparse

if TYPE_CHECKING:
    from collections.abc import Iterator

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

HTTP_UNAUTHORIZED = 401
HTTP_UNPROCESSABLE = 422
MIN_MASK_VISIBLE = 4


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).
    """

    USER_PROMPT_TEMPLATE: str | None = None

    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.client = openai.OpenAI(
            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,
        mix_traits: bool = False,
    ) -> 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).
            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).

        Returns:
            List of simulation results with conv_prefix and response.

        """
        with self._steer_settings(steer_temperature, steer_max_tokens):
            combinations = config.combinations(mix_traits=mix_traits)
            samples = self._select_samples(combinations, k)
            if not samples:
                return []
            return self._execute_samples(samples, num_exchanges, batch_delay)

    @contextmanager
    def _steer_settings(
        self, steer_temperature: float | None, steer_max_tokens: int | None
    ) -> Iterator[None]:
        prev_temp = self.steer_temperature
        prev_max = self.steer_max_tokens
        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)
            yield
        finally:
            self.steer_temperature = prev_temp
            self.steer_max_tokens = prev_max

    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 _execute_samples(
        self,
        samples: list[SteerCombination],
        num_exchanges: int,
        batch_delay: float,
    ) -> list[SimulationResult]:
        simulations: list[SimulationResult] = []
        for i, combo in enumerate(samples):
            try:
                if i > 0 and batch_delay > 0:
                    time.sleep(batch_delay)
                self.logger.info("=" * 40)
                conversation, final_response = 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.InvalidTraitError:
                self.logger.warning(
                    "Skipping simulation %d/%d due to invalid trait '%s'.",
                    i + 1,
                    len(samples),
                    combo.trait,
                )
                continue
            except Exception:
                self.logger.exception(f"Failed simulation {i + 1}/{len(samples)}")
                continue
        return simulations

    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 = 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()))
                time.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 ""

    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")
            system_prompt = self._build_user_system_prompt(combo)
            if len(combo.traits) == 1:
                trait = next(iter(combo.traits))
                intensity = next(iter(combo.traits.values()))
                response = self._call_collinear_steer_api(
                    _conversation=conversation,
                    _system_prompt=system_prompt,
                    trait=trait,
                    intensity=float(intensity),
                )
            else:
                response = self._call_collinear_steer_api_trait_dict(
                    _conversation=conversation,
                    _system_prompt=system_prompt,
                    trait_dict={k: float(v) for k, v in combo.traits.items()},
                )
        else:
            self.logger.info("Generating ASSISTANT turn")
            system_prompt = self.ASSISTANT_PROMPT_TEMPLATE
            response = self._call_with_retry(conversation, system_prompt)
        self.logger.info(response)

        return response

    def _build_user_system_prompt(self, combo: SteerCombination) -> str:
        """Construct a line-per-item USER system prompt.

        Returns a caller-provided template if set; otherwise builds a dynamic
        prompt listing present axes. Intensities remain only in the payload.
        """
        formatted = self._format_user_template_if_any(combo)
        if formatted is not None:
            return formatted

        lines: list[str] = [
            "You are the CUSTOMER in a chat with a customer service assistant.",
            (
                "Write only the next customer message in first person. "
                "Keep it concise (1-2 sentences)."
            ),
            "Do not include role names or quotes.",
            "",
        ]

        self._append_axis_lines(lines, combo)
        trait_line = self._traits_line(combo)
        if trait_line:
            lines.append(trait_line)
        lines.extend(self._language_note(combo))
        return "\n".join(lines)

    def _format_user_template_if_any(self, combo: SteerCombination) -> str | None:
        template: str | None = cast("str | None", getattr(self, "USER_PROMPT_TEMPLATE", None))
        if not template:
            return None
        age = combo.age or ""
        gender = combo.gender or ""
        occupation = combo.occupation or ""
        intent = combo.intent or ""
        location = combo.location or ""
        language = combo.language or ""
        if combo.traits:
            keys = list(combo.traits.keys())
            trait_text = keys[0] if len(keys) == 1 else " and ".join(keys[:2])
        else:
            trait_text = ""
        article = "an" if occupation[:1].lower() in "aeiou" else "a"
        try:
            return template.format(
                age=age,
                gender=gender,
                occupation=occupation,
                intent=intent,
                trait=trait_text,
                article=article,
                location=location,
                language=language,
            )
        except Exception as e:
            self.logger.debug("USER_PROMPT_TEMPLATE format failed: %s", e)
            return None

    @staticmethod
    def _maybe_add(lines: list[str], label: str, value: str | None) -> None:
        if value:
            lines.append(f"{label}: {value}")

    def _append_axis_lines(self, lines: list[str], combo: SteerCombination) -> None:
        self._maybe_add(lines, "Age", combo.age)
        self._maybe_add(lines, "Gender", combo.gender)
        self._maybe_add(lines, "Occupation", combo.occupation)
        self._maybe_add(lines, "Location", combo.location)
        self._maybe_add(lines, "Language", combo.language)
        self._maybe_add(lines, "Intent", combo.intent)

    @staticmethod
    def _traits_line(combo: SteerCombination) -> str | None:
        if not combo.traits:
            return None
        keys = list(combo.traits.keys())
        return f"Traits: {keys[0]}" if len(keys) == 1 else f"Traits: {keys[0]}, {keys[1]}"

    @staticmethod
    def _language_note(combo: SteerCombination) -> list[str]:
        return ["", "Write the message in the specified language."] if combo.language else []

    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."""

    def _steer_url(self) -> str:
        return os.getenv("COLLINEAR_STEER_URL", "https://steer.collinear.ai/steer")

    def _build_steer_payload(
        self,
        *,
        trait_dict: dict[str, float],
        full_conv: list[ChatCompletionMessageParam],
    ) -> dict[str, object]:
        return {
            "trait_dict": {k: float(v) for k, v in trait_dict.items()},
            "messages": full_conv,
            "temperature": float(self.steer_temperature),
            "max_tokens": int(self.steer_max_tokens),
        }

    def _request_steer(
        self, url: str, headers: dict[str, str], payload: dict[str, object]
    ) -> tuple[httpx.Response | None, str | None]:
        try:
            with httpx.Client(timeout=self.timeout) as client:
                self.logger.debug("POST %s (API-Key present=%s)", url, "API-Key" in headers)
                resp = 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

    def _call_collinear_steer_api(
        self,
        *,
        trait: str,
        intensity: float,
        conversation: list[ChatCompletionMessageParam] | None = None,
        system_prompt: str | None = None,
        _conversation: list[ChatCompletionMessageParam] | None = None,
        _system_prompt: str | None = None,
    ) -> str:
        conv = conversation if conversation is not None else (_conversation or [])
        sys = system_prompt if system_prompt is not None else (_system_prompt or "")
        sys_msg: ChatCompletionSystemMessageParam = {"role": "system", "content": sys}
        full_conv: list[ChatCompletionMessageParam] = [sys_msg, *conv]

        payload: dict[str, object] = self._build_steer_payload(
            trait_dict={trait: float(intensity)}, full_conv=full_conv
        )
        headers: dict[str, str] = {
            "Content-Type": "application/json",
            "API-Key": self.steer_api_key,
        }

        url = self._steer_url()
        resp, err = 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:
            self._handle_unprocessable_or_skip(trait, 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)

    def _call_collinear_steer_api_trait_dict(
        self,
        *,
        trait_dict: dict[str, float],
        conversation: list[ChatCompletionMessageParam] | None = None,
        system_prompt: str | None = None,
        _conversation: list[ChatCompletionMessageParam] | None = None,
        _system_prompt: str | None = None,
    ) -> str:
        conv = conversation if conversation is not None else (_conversation or [])
        sys = system_prompt if system_prompt is not None else (_system_prompt or "")
        sys_msg: ChatCompletionSystemMessageParam = {"role": "system", "content": sys}
        full_conv: list[ChatCompletionMessageParam] = [sys_msg, *conv]

        payload: dict[str, object] = self._build_steer_payload(
            trait_dict=trait_dict, full_conv=full_conv
        )
        headers: dict[str, str] = {
            "Content-Type": "application/json",
            "API-Key": self.steer_api_key,
        }

        url = self._steer_url()
        resp, err = 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:
            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)

    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 = ""
        for turn in range(1, total_turns + 1):
            is_user_turn = turn % 2 == 1
            role = Role.USER if is_user_turn else Role.ASSISTANT
            response = self._generate_turn(combo, conversation, role=role)
            if is_user_turn:
                u: ChatCompletionUserMessageParam = {"role": Role.USER.value, "content": response}
                conversation.append(u)
            else:
                a: ChatCompletionAssistantMessageParam = {
                    "role": Role.ASSISTANT.value,
                    "content": response,
                }
                conversation.append(a)
                if turn == total_turns:
                    final_response = response
        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 []
