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

from __future__ import annotations

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

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

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

if TYPE_CHECKING:
    from collections.abc import Iterator

HTTP_UNAUTHORIZED = 401
HTTP_UNPROCESSABLE = 422
MIN_MASK_VISIBLE = 4
DEFAULT_STEER_SEED = -1


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

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

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

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


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


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


_PROGRESS_FACTORY: _ProgressFactory | None = _load_progress_factory()


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

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

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

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

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

        """
        if not assistant_model_name:
            raise ValueError("model_name is required")
        if not steer_api_key:
            raise ValueError("steer_api_key is required")
        self.assistant_model_url = assistant_model_url
        self.assistant_model_api_key = assistant_model_api_key
        self.assistant_model_name = assistant_model_name
        self.steer_api_key = steer_api_key
        self.timeout = timeout
        self.max_retries = max_retries
        self.rate_limit_retries = rate_limit_retries
        self.logger = logging.getLogger("collinear")
        self.steer_temperature: float = 0.7
        self.steer_max_tokens: int = 256
        self.steer_seed: int = DEFAULT_STEER_SEED
        self._progress_bar: _ProgressBar | None = None
        self._user_turns_completed: int = 0
        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,
        steer_seed: int | None = None,
        mix_traits: bool = False,
        progress: bool = True,
    ) -> list[SimulationResult]:
        """Run simulations with the given configuration.

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

        Returns:
            List of simulation results with conv_prefix and response.

        """
        with self._steer_settings(
            steer_temperature=steer_temperature,
            steer_max_tokens=steer_max_tokens,
            steer_seed=steer_seed,
        ):
            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 self._execute_samples(samples, num_exchanges, batch_delay)

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

    @contextmanager
    def _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="Steer queries", unit="query")
        except Exception:
            self.logger.exception("Failed to initialize progress bar; continuing without it.")
            yield
            return
        self._progress_bar = bar
        try:
            yield
        finally:
            self._progress_bar = None
            try:
                bar.close()
            except Exception:
                self.logger.debug("Failed to close progress bar", exc_info=True)

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

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

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

    def _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)
                self._user_turns_completed = 0
                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:
                remaining = max(0, num_exchanges - self._user_turns_completed)
                if remaining:
                    self._adjust_progress_total(remaining)
                self.logger.warning(
                    "Skipping simulation %d/%d due to invalid trait '%s'.",
                    i + 1,
                    len(samples),
                    combo.trait,
                )
                continue
            except Exception:
                remaining = max(0, num_exchanges - self._user_turns_completed)
                if remaining:
                    self._adjust_progress_total(remaining)
                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")
            if len(combo.traits) == 1:
                trait = next(iter(combo.traits))
                intensity = next(iter(combo.traits.values()))
                response = self._call_collinear_steer_api(
                    _conversation=conversation,
                    trait=trait,
                    intensity=intensity,
                    combo=combo,
                )
            else:
                response = self._call_collinear_steer_api_trait_dict(
                    _conversation=conversation,
                    trait_dict={k: int(v) for k, v in combo.traits.items()},
                    combo=combo,
                )
        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 _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, int],
        conversation: list[ChatCompletionMessageParam],
        combo: SteerCombination,
    ) -> dict[str, object]:
        user_characteristics = self._user_characteristics_payload(combo)
        messages = [
            {"role": str(msg["role"]), "content": str(msg["content"])} for msg in conversation
        ]

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

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

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

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

        return payload

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

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

    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: int,
        combo: SteerCombination,
        conversation: list[ChatCompletionMessageParam] | None = None,
        _conversation: list[ChatCompletionMessageParam] | None = None,
    ) -> str:
        conv = conversation if conversation is not None else (_conversation or [])

        return self._call_collinear_steer_api_trait_dict(
            trait_dict={trait: int(intensity)},
            combo=combo,
            conversation=conv,
        )

    def _call_collinear_steer_api_trait_dict(
        self,
        *,
        trait_dict: dict[str, int],
        combo: SteerCombination,
        conversation: list[ChatCompletionMessageParam] | None = None,
        _conversation: list[ChatCompletionMessageParam] | None = None,
    ) -> str:
        conv = conversation if conversation is not None else (_conversation or [])

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

        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:
            if len(trait_dict) == 1:
                trait = next(iter(trait_dict))
                self._handle_unprocessable_or_skip(trait, resp)
            else:
                self._handle_unprocessable_or_skip_mixed(set(trait_dict.keys()), resp)

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

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

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

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

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