"""Text generation and conversation management for the benchmark."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Optional

import numpy as np

from ray.llm._internal.serve.benchmark.models import WorkloadSpec

if TYPE_CHECKING:
    from transformers import PreTrainedTokenizerBase

logger = logging.getLogger(__name__)


class Conversation:
    """A single multi-turn conversation with a unique session ID."""

    def __init__(
        self,
        session_id: str,
        system_prompt: str,
        user_messages: list[str],
        num_turns: int,
    ):
        self.session_id = session_id
        self.system_prompt = system_prompt
        self.user_messages = user_messages
        self.num_turns = num_turns
        self._assistant_responses: list[str] = []

    def get_turn_messages(self, turn_idx: int) -> list[dict[str, str]]:
        """Build the messages list for turn `turn_idx` (0-indexed)."""
        messages: list[dict[str, str]] = []
        if self.system_prompt:
            messages.append({"role": "system", "content": self.system_prompt})

        for i in range(turn_idx + 1):
            messages.append({"role": "user", "content": self.user_messages[i]})
            if i < turn_idx:
                if i < len(self._assistant_responses):
                    messages.append(
                        {"role": "assistant", "content": self._assistant_responses[i]}
                    )
                else:
                    messages.append({"role": "assistant", "content": "(placeholder)"})
        return messages

    def inject_assistant_response(self, turn_idx: int, content: str) -> None:
        """Record the server's response for turn `turn_idx`."""
        if turn_idx == len(self._assistant_responses):
            self._assistant_responses.append(content)
        elif turn_idx < len(self._assistant_responses):
            self._assistant_responses[turn_idx] = content
        else:
            raise ValueError(
                f"Cannot inject response for turn {turn_idx}: "
                f"only {len(self._assistant_responses)} responses recorded."
            )


class TextGenerator:
    """Generates random text with exact token counts using a tokenizer."""

    def __init__(self, tokenizer: "PreTrainedTokenizerBase"):
        self._tokenizer = tokenizer
        self._vocab_size = tokenizer.vocab_size
        logger.info(
            "TextGenerator using tokenizer (vocab_size=%d) for exact token counts.",
            self._vocab_size,
        )

    def generate(self, num_tokens: int) -> str:
        if num_tokens <= 0:
            return ""
        return self._generate_exact(num_tokens)

    def generate_token_ids(self, num_tokens: int) -> list[int]:
        if num_tokens <= 0:
            return []
        return np.random.randint(0, self._vocab_size, size=num_tokens).tolist()

    def _generate_exact(self, target_tokens: int) -> str:
        tokenizer = self._tokenizer
        token_ids = np.random.randint(
            0, self._vocab_size, size=target_tokens + 20
        ).tolist()

        text = tokenizer.decode(token_ids, skip_special_tokens=True)
        actual_ids = tokenizer.encode(text, add_special_tokens=False)
        actual_len = len(actual_ids)

        if actual_len == target_tokens:
            return text

        if actual_len > target_tokens:
            trimmed_ids = actual_ids[:target_tokens]
            text = tokenizer.decode(trimmed_ids, skip_special_tokens=True)
            final_len = len(tokenizer.encode(text, add_special_tokens=False))
            if final_len != target_tokens:
                text = self._binary_search_trim(actual_ids, target_tokens)
            return text

        deficit = target_tokens - actual_len
        extra_ids = np.random.randint(0, self._vocab_size, size=deficit + 20).tolist()
        extra_text = tokenizer.decode(extra_ids, skip_special_tokens=True)
        combined = text + " " + extra_text
        combined_ids = tokenizer.encode(combined, add_special_tokens=False)

        if len(combined_ids) >= target_tokens:
            trimmed = combined_ids[:target_tokens]
            text = tokenizer.decode(trimmed, skip_special_tokens=True)
            final_len = len(tokenizer.encode(text, add_special_tokens=False))
            if final_len != target_tokens:
                text = self._binary_search_trim(combined_ids, target_tokens)
            return text

        while len(tokenizer.encode(combined, add_special_tokens=False)) < target_tokens:
            combined += " hello"
        combined_ids = tokenizer.encode(combined, add_special_tokens=False)
        return self._binary_search_trim(combined_ids, target_tokens)

    def _binary_search_trim(self, token_ids: list[int], target: int) -> str:
        tokenizer = self._tokenizer
        lo, hi = target, len(token_ids)
        best_text = tokenizer.decode(token_ids[:target], skip_special_tokens=True)

        while lo <= hi:
            mid = (lo + hi) // 2
            text = tokenizer.decode(token_ids[:mid], skip_special_tokens=True)
            actual = len(tokenizer.encode(text, add_special_tokens=False))
            if actual == target:
                return text
            elif actual < target:
                lo = mid + 1
            else:
                hi = mid - 1
                best_text = text

        for n in range(target, len(token_ids) + 1):
            text = tokenizer.decode(token_ids[:n], skip_special_tokens=True)
            if len(tokenizer.encode(text, add_special_tokens=False)) == target:
                return text
        return best_text


def conversation_factory(
    session_idx: int,
    spec: WorkloadSpec,
    shared_system_text: str,
    text_gen: Optional[TextGenerator],
) -> Conversation:
    """Create a single conversation on-demand (lazy generation)."""
    session_id = f"session-{session_idx:06d}"

    if spec.unique_s > 0 and text_gen is not None:
        unique_text = text_gen.generate(spec.unique_s)
        system_prompt = shared_system_text + " " + unique_text
    else:
        system_prompt = shared_system_text

    user_messages = (
        [text_gen.generate(spec.user_tokens) for _ in range(spec.num_turns)]
        if text_gen is not None
        else ["" for _ in range(spec.num_turns)]
    )

    return Conversation(
        session_id=session_id,
        system_prompt=system_prompt,
        user_messages=user_messages,
        num_turns=spec.num_turns,
    )