"""User persona simulator.

Generates user messages based on Identity/Goal/Personality prompts
using an LLM to simulate realistic user behavior.
"""

from collections.abc import Awaitable
from collections.abc import Callable
import re
import uuid

import dspy
from pydantic import BaseModel
from pydantic import ValidationError

from voicetest.llm import _invoke_callback
from voicetest.llm import call_llm
from voicetest.models.results import Message
from voicetest.util.cache import try_evict_last_call
from voicetest.util.retry import EmptyLLMOutputError
from voicetest.util.retry import OnErrorCallback


# Callback type for token updates: receives token string and source ("agent" or "user")
OnTokenCallback = Callable[[str, str], Awaitable[None] | None]


class UserSimSignature(dspy.Signature):
    """Simulate a real phone caller. Respond briefly and naturally, the way people actually talk
    on the phone. No introductions, pleasantries, or formalities unless the persona demands it.
    On the first turn, say a standard greeting, hello etc unless instructed otherwise in persona."""

    persona: str = dspy.InputField(desc="User persona (identity, goal, personality)")
    conversation_history: str = dspy.InputField(
        desc="Prior conversation (excluding latest agent message)"
    )
    current_agent_message: str = dspy.InputField(
        desc="The agent's most recent message to respond to"
    )
    turn_number: int = dspy.InputField(desc="Current turn number (1 = first turn)")

    message: str = dspy.OutputField(
        desc="User's next spoken response — short and natural like a real phone call."
    )


class SimulatorResponse(BaseModel):
    """Response from user simulator."""

    message: str


class UserSimulator:
    """LLM-based user persona simulator.

    Generates realistic user messages based on a persona definition
    following the Identity/Goal/Personality format.
    """

    def __init__(self, user_prompt: str, model: str):
        """Initialize the simulator.

        Args:
            user_prompt: Persona definition in Identity/Goal/Personality format.
            model: LLM model to use for generation.
        """
        self.user_prompt = user_prompt
        self.model = model

        # Mock mode for testing without LLM calls
        self._mock_mode = False
        self._mock_responses: list[SimulatorResponse] = []
        self._mock_index = 0

    async def generate(
        self,
        transcript: list[Message],
        on_token: OnTokenCallback | None = None,
        on_error: OnErrorCallback | None = None,
    ) -> SimulatorResponse | None:
        """Generate next user message based on conversation so far.

        Args:
            transcript: Conversation history.
            on_token: Optional callback for streaming tokens.
            on_error: Optional callback for retryable errors.

        Returns:
            SimulatorResponse with the generated message, or None if mock responses exhausted.
        """
        # Mock mode for testing
        if self._mock_mode:
            if self._mock_index >= len(self._mock_responses):
                return None
            response = self._mock_responses[self._mock_index]
            self._mock_index += 1
            return response

        return await self._generate_with_llm(transcript, on_token, on_error)

    async def _generate_with_llm(
        self,
        transcript: list[Message],
        on_token: OnTokenCallback | None = None,
        on_error: OnErrorCallback | None = None,
    ) -> SimulatorResponse | None:
        """Generate response using LLM.

        Args:
            transcript: Conversation history.
            on_token: Optional callback for streaming tokens.
            on_error: Optional callback for retryable errors.

        Returns:
            SimulatorResponse with the generated message.
        """
        user_prompt = self.user_prompt
        turn_number = len([m for m in transcript if m.role == "user"]) + 1

        # Split transcript into history and current agent message
        current_agent_message = ""
        history_transcript = transcript
        if transcript and transcript[-1].role == "assistant":
            current_agent_message = transcript[-1].content
            history_transcript = transcript[:-1]

        conversation_history = self._format_transcript(history_transcript)

        # Wrap on_token to add source="user"
        async def user_token_callback(token: str) -> None:
            await _invoke_callback(on_token, token, "user")

        # One retry if the LLM returns None for `message`. Pydantic's ValidationError
        # on SimulatorResponse is the signal — we translate to EmptyLLMOutputError
        # so the failure surfaces with a clear diagnostic.
        #
        # Retry uses a random salt AND no_cache=True so it (a) never reads a cached
        # entry and (b) never writes one (no accumulation of zombie retry entries).
        # On ValidationError we also evict the poisoned cache entry from the first
        # call so future runs don't hit it.
        last_error: ValidationError | None = None
        for attempt in range(2):
            is_retry = attempt > 0
            cache_salt = uuid.uuid4().hex if is_retry else None
            result = await call_llm(
                self.model,
                UserSimSignature,
                on_token=user_token_callback if on_token else None,
                stream_field="message" if on_token else None,
                on_error=on_error,
                cache_salt=cache_salt,
                no_cache=is_retry,
                predictor_class=dspy.Predict,
                persona=user_prompt,
                conversation_history=conversation_history,
                current_agent_message=current_agent_message or "(agent has not spoken yet)",
                turn_number=turn_number,
            )

            try:
                return SimulatorResponse(message=result.message)
            except ValidationError as e:
                last_error = e
                # Evict the poisoned cache entry (no-op if retry call had no_cache=True)
                lm = getattr(result, "_voicetest_lm", None)
                if lm is not None:
                    try_evict_last_call(lm)

        raise EmptyLLMOutputError(field_name="message", model=self.model) from last_error

    def _format_transcript(self, transcript: list[Message]) -> str:
        """Format transcript for LLM input."""
        if not transcript:
            return "(conversation not started)"

        lines = []
        for msg in transcript:
            lines.append(f"{msg.role.upper()}: {msg.content}")
        return "\n".join(lines)

    def _parse_persona(self) -> dict[str, str]:
        """Parse Identity/Goal/Personality sections from user prompt.

        Returns:
            Dictionary with 'identity', 'goal', and 'personality' keys.
        """
        sections: dict[str, str] = {}

        # Pattern to match ## Section headers
        pattern = r"##\s*(Identity|Goal|Personality)\s*\n(.*?)(?=##|\Z)"
        matches = re.findall(pattern, self.user_prompt, re.IGNORECASE | re.DOTALL)

        for section_name, content in matches:
            sections[section_name.lower()] = content.strip()

        return sections
