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

from __future__ import annotations

import logging
import os
import random
import time
from typing import cast

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.persona import PersonaCombination
from collinear.schemas.persona import PersonaConfig
from collinear.schemas.persona import Role
from collinear.schemas.persona import SimulationResult

HTTP_UNAUTHORIZED = 401
MIN_MASK_VISIBLE = 4


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

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

    USER_PROMPT_TEMPLATE = (
        "You are a {age} year old {gender}, "
        "who works as {article} {occupation}. "
        "You are {trait} and are asking a customer service "
        "representative about: {intent}.\n\n"
        "You are the CUSTOMER in a chat with a customer service assistant.\n"
        "Write only the next customer message in first person, staying in character.\n"
        "Keep it concise (1-2 sentences). Do not include role names or quotes."
    )

    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,
        *,
        persona_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.
            persona_api_key: Persona service API key (sent as ``X-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 persona_api_key:
            raise ValueError("persona_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.persona_api_key = persona_api_key
        self.timeout = timeout
        self.max_retries = max_retries
        self.rate_limit_retries = rate_limit_retries
        self.logger = logging.getLogger("collinear")
        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: PersonaConfig,
        k: int = 10,
        num_exchanges: int = 2,
        batch_delay: float = 0.1,
    ) -> list[SimulationResult]:
        """Run k simulations with the given configuration.

        Args:
            config: Persona configuration
            k: Number of simulations to run
            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)

        Returns:
            List of simulation results with conv_prefix and response

        """
        combinations = config.combinations()
        samples = random.sample(combinations, min(k, len(combinations)))
        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,
                        persona=combo,
                    )
                )
                self.logger.info(f"Completed simulation {i + 1}/{len(samples)}")
            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: PersonaCombination,
        conversation: list[ChatCompletionMessageParam],
        role: Role,
    ) -> str:
        """Generate a single turn in the conversation."""
        if role is Role.USER:
            self.logger.info("Generating USER turn")
            article = "an" if combo.occupation[0].lower() in "aeiou" else "a"
            system_prompt = self.USER_PROMPT_TEMPLATE.format(
                age=combo.age,
                gender=combo.gender,
                article=article,
                occupation=combo.occupation,
                trait=combo.trait,
                intent=combo.intent,
            )
            response = self._call_collinear_persona_api(
                _conversation=conversation,
                _system_prompt=system_prompt,
                trait=combo.trait,
                intensity=combo.intensity,
            )
        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.persona_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(
            "Persona API unauthorized (401). X-API-Key preview=%s. Body=%s",
            self._mask_key_preview(),
            resp.text,
        )

    def _call_collinear_persona_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] = {
            "strength": float(intensity),
            "trait": trait,
            "messages": full_conv,
            "temperature": 0.7,
            "max_tokens": 256,
        }

        headers: dict[str, str] = {
            "Content-Type": "application/json",
            "X-API-Key": self.persona_api_key,
        }

        url = os.getenv(
            "COLLINEAR_STEER_URL",
            "https://6ddc3d50eee4.ngrok-free.app/steer",
        )
        try:
            with httpx.Client(timeout=self.timeout) as client:
                self.logger.debug(
                    "POST %s (X-API-Key present=%s)",
                    url,
                    "X-API-Key" in headers,
                )
                resp = client.post(url, headers=headers, json=payload)
                if resp.status_code == HTTP_UNAUTHORIZED:
                    self._log_unauthorized(resp)
                resp.raise_for_status()
                raw: object = resp.json()
        except Exception as e:
            self.logger.exception("User service error")
            return (
                "Error: Persona API call failed. "
                f"Details: {e!s}. Check PERSONA_API_KEY and COLLINEAR_STEER_URL."
            )

        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: PersonaCombination, 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
