import re
from typing import Optional, Tuple, Dict

from loguru import logger

from murmur.data_model.message import (
    Message,
    MultiToolMessage,
    SystemMessage,
    ToolCall,
    UserMessage,
)
from murmur.data_model.tasks import UserInstructions
from murmur.environment.tool import Tool
from murmur.user.base import (
    OUT_OF_SCOPE,
    STOP,
    TRANSFER,
    BaseUser,
    UserState,
    ValidUserInputMessage,
    is_valid_user_history_message,
)
from murmur.user.user_simulator import UserSimulator, get_global_user_sim_guidelines
from murmur.utils.llm_utils import generate

# Multi-user specific system prompt template
MULTI_USER_SYSTEM_PROMPT = """
{global_user_sim_guidelines}

<scenario>
<user_name>{user_name}</user_name>
<user_role>{user_role}</user_role>
{instructions}
</scenario>

IMPORTANT: Your name is {user_name}. You are the {user_role}. When communicating with the customer service agent or other users, always identify yourself clearly and maintain your unique persona and role. If multiple users are participating in the conversation, be aware that others may also be contributing to the discussion.
""".strip()


class MultiUserSimulator(UserSimulator):
    """
    Enhanced user simulator that supports named users and multi-user interactions.
    Each instance represents a distinct user with a unique name and role.
    """

    def __init__(
        self,
        user_name: str,
        user_role: Optional[str] = None,
        tools: Optional[list[Tool]] = None,
        instructions: Optional[UserInstructions] = None,
        llm: Optional[str] = None,
        llm_args: Optional[dict] = None,
        user_id: Optional[str] = None,
    ):
        super().__init__(tools=tools, instructions=instructions, llm=llm, llm_args=llm_args)
        self.user_name = user_name
        self.user_role = user_role or "customer"
        self.user_id = user_id

    @property
    def system_prompt(self) -> str:
        """
        The system prompt for the multi-user simulator.
        Includes user name and role information.
        """
        if self.instructions is None:
            logger.warning(f"No instructions provided for user simulator {self.user_name}")

        system_prompt = MULTI_USER_SYSTEM_PROMPT.format(
            global_user_sim_guidelines=self.global_simulation_guidelines,
            user_name=self.user_name,
            user_role=self.user_role,
            instructions=self.instructions or "",
        )
        return system_prompt

    def generate_next_message(
        self, message: ValidUserInputMessage, state: UserState
    ) -> Tuple[UserMessage, UserState]:
        """Generate the next message with the user's identity."""
        logger.debug(f"Generating next message for user {self.user_name}")
        user_message, new_state = super().generate_next_message(message, state)
        
        # Ensure the message content includes the user's identity if it's a text message
        if user_message.content and not user_message.is_tool_call():
            # Normalize any leading speaker tag at the start of lines to [user_id]:
            # Matches either [Name]: or bare Name: at line start (multiline)
            pattern = r"(?m)^\s*(?:\[[^\]\n]{0,30}\]|[^:\[\]\n]{0,30})?\s*:"
            new_content, replaced = re.subn(pattern, f"[{self.user_id}]:", user_message.content)
            if replaced > 0:
                user_message.content = new_content
            else:
                user_message.content = f"[{self.user_id}]: {user_message.content}"
        
        return user_message, new_state