import json
import multiprocessing
import random
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import string
from typing import Optional

from loguru import logger

from murmur.agent.llm_agent import LLMAgent, LLMGTAgent, LLMSoloAgent, LLMProtectAIAgent, LLMPromptGuardAgent
from murmur.agent.llm_murmur_agent import LLMMurmurAgent
from murmur.data_model.simulation import (AgentInfo, Info, Results, RunConfig,
                                        SimulationRun, UserInfo)
from murmur.data_model.tasks import Task
from murmur.environment.environment import Environment, EnvironmentInfo
from murmur.evaluator.evaluator import EvaluationType, evaluate_simulation, evaluate_multi_task_simulation
from murmur.metrics.agent_metrics import compute_metrics
from murmur.orchestrator.orchestrator import Orchestrator
from murmur.orchestrator.multi_user_orchestrator import MultiUserOrchestrator
from murmur.registry import RegistryInfo, registry
from murmur.user.user_simulator import DummyUser, get_global_user_sim_guidelines
from murmur.user.multi_user_simulator import MultiUserSimulator
from murmur.user.user_simulator import UserSimulator
from murmur.utils.display import ConsoleDisplay, Text
from murmur.utils.pydantic_utils import get_pydantic_hash
from murmur.utils.utils import DATA_DIR, get_commit_hash, get_now, show_dict_diff


def categorize_tasks_by_user_count(tasks: list[Task]) -> dict[int, list[Task]]:
    """
    Categorize tasks by the number of users they contain.
    
    Args:
        tasks: List of tasks to categorize
        
    Returns:
        Dictionary mapping user count to list of tasks
    """
    categories = {}
    for task in tasks:
        # Handle both string and StructuredUserInstructions types
        if hasattr(task.user_scenario.instructions, 'user_names') and task.user_scenario.instructions.user_names:
            user_count = len(task.user_scenario.instructions.user_names)
        else:
            # Default to 1 user if user_names is not available or empty
            user_count = 1
            
        if user_count not in categories:
            categories[user_count] = []
        categories[user_count].append(task)
    return categories


def generate_task_combinations(tasks: list[Task], multi_user: bool, num_tasks: int) -> list[list[Task]]:
    """
    Generate combinations of tasks with one task from each user category.
    Categories are: 1 user, 2 users, and more than 2 users.
    
    Args:
        tasks: List of tasks to generate combinations from
        
    Returns:
        List of task combinations, where each combination has one task from each category
    """
    # shuffle tasks and make groups of tasks
    random.shuffle(tasks)
    if multi_user and num_tasks is not None:
        combinations = [tasks[i:i+num_tasks] for i in range(0, len(tasks), num_tasks)]
        # add all remaining tasks to the last combination
        combinations[-1].extend(tasks[num_tasks*len(combinations):])
    else:
        # Default to individual tasks
        combinations = [tasks[i:i+1] for i in range(0, len(tasks), 1)]
    
    logger.info(f"Generated {len(combinations)} task combinations")
    
    return combinations


def get_options() -> RegistryInfo:
    """
    Returns options for the simulator.
    """
    return registry.get_info()


def get_environment_info(
    domain_name: str, include_tool_info: bool = False
) -> EnvironmentInfo:
    """Get information about the environment for a registered Domain"""
    global registry
    env_constructor = registry.get_env_constructor(domain_name)
    return env_constructor().get_info(include_tool_info=include_tool_info)


def load_tasks(task_set_name: str) -> list[Task]:
    """
    Loads the tasks for the given domain.
    """
    global registry
    task_loader = registry.get_tasks_loader(task_set_name)
    tasks = task_loader()
    return tasks


def get_tasks(
    task_set_name: str,
    task_ids: Optional[list[str]] = None,
    num_tasks: Optional[int] = None,
) -> list[Task]:
    """
    Loads the tasks for the given domain.
    """
    if task_ids is None:
        tasks = load_tasks(task_set_name=task_set_name)
    else:
        tasks = [
            task for task in load_tasks(task_set_name=task_set_name) if task.id in task_ids
        ]
    if task_ids is not None and len(tasks) != len(task_ids):
        missing_tasks = set(task_ids) - set([task.id for task in tasks])
        raise ValueError(
            f"Not all tasks were found for task set {task_set_name}: {missing_tasks}"
        )
    return tasks


def extract_multi_user_info(task: Task) -> dict:
    """
    Extract multi-user information from task's user scenario.
    Returns a dictionary with user information for multi-user scenarios.
    """
    if not hasattr(task.user_scenario.instructions, 'user_names') or task.user_scenario.instructions.user_names is None:
        # Fallback: extract from known_info if available
        known_info = getattr(task.user_scenario.instructions, 'known_info', None)
        if known_info and ('you are' in known_info.lower() or 'your user id' in known_info.lower()):
            # Try to extract user names from known_info
            # This is a simple parser - might need to be more sophisticated
            lines = known_info.split('\n')
            user_names = []
            user_ids = {}
            for line in lines:
                if 'you are' in line.lower():
                    # line = line.lower()
                    # Extract names after "You are"
                    parts = line.split('You are', 1)[1].strip()
                    if ',' in parts:
                        names = [name.strip() for name in parts.split(',')]
                        user_names.extend(names)
                    elif '.' in parts:
                        names = [name.strip() for name in parts.split('.')]
                        user_names.extend(names)
                    else:
                        user_names.append(parts.strip())
                elif 'your user id' in line.lower():
                    # Extract user ID
                    parts = line.split('is', 1)
                    if len(parts) > 1:
                        user_id = parts[1].strip().rstrip('.')
                        if user_names:
                            user_ids[user_names[0]] = user_id
            
            if not user_names:
                user_names = ['User1', 'User2']  # Default fallback
            
            return {
                'user_names': user_names,
                'user_roles': {name: 'customer' for name in user_names},
                'per_user_instructions': {}
            }
    
    # Use structured multi-user information if available
    instructions = task.user_scenario.instructions
    return {
        'user_names': getattr(instructions, 'user_names', ['User1']),
        'user_roles': getattr(instructions, 'user_roles', {}),
        'per_user_instructions': getattr(instructions, 'per_user_instructions', {})
    }


def make_run_name(config: RunConfig) -> str:
    """
    Make a run name from the run config
    """
    clean_llm_agent_name = config.llm_agent.split("/")[-1]
    agent_name = f"{config.agent}_{clean_llm_agent_name}"

    clean_llm_user_name = config.llm_user.split("/")[-1]
    user_name = f"{config.user}_{clean_llm_user_name}"

    return f"{get_now()}_{config.domain}_{agent_name}_{user_name}"


def run_domain(config: RunConfig) -> Results:
    """
    Run simulations for a domain
    """

    config.validate()
    ConsoleDisplay.display_run_config(config)
    if config.task_set_name is None:
        task_set_name = config.domain
    else:
        task_set_name = config.task_set_name
    tasks = get_tasks(task_set_name, config.task_ids, config.num_tasks)
    if "gt" in config.agent:
        total_num_tasks = len(tasks)
        tasks = [task for task in tasks if LLMGTAgent.check_valid_task(task)]
        num_tasks = len(tasks)
        console_text = Text(text=f"Running {num_tasks} out of {total_num_tasks} tasks for GT agent.", style="bold green")
        ConsoleDisplay.console.print(console_text)
    if "solo" in config.agent:
        total_num_tasks = len(tasks)
        tasks = [task for task in tasks if LLMSoloAgent.check_valid_task(task)]
        num_tasks = len(tasks)
        console_text = Text(text=f"Running {num_tasks} out of {total_num_tasks} tasks for solo agent.", style="bold green")
        ConsoleDisplay.console.print(console_text)

    num_trials = config.num_trials
    save_to = config.save_to
    if save_to is None:
        save_to = make_run_name(config)
    save_to = DATA_DIR / "simulations" / f"{save_to}.json"
    simulation_results = run_tasks(
        domain=config.domain,
        tasks=tasks,
        agent=config.agent,
        user=config.user,
        llm_agent=config.llm_agent,
        llm_args_agent=config.llm_args_agent,
        llm_user=config.llm_user,
        llm_args_user=config.llm_args_user,
        num_trials=num_trials,
        max_steps=config.max_steps,
        max_errors=config.max_errors,
        save_to=save_to,
        console_display=True,
        evaluation_type=EvaluationType.ALL,
        max_concurrency=config.max_concurrency,
        seed=config.seed,
        log_level=config.log_level,
        multi_user=config.multi_user,
        allow_multiple_user_messages=config.allow_multiple_user_messages,
        num_tasks=config.num_tasks,
        injection_task=config.injection_task,
    )
    metrics = compute_metrics(simulation_results)
    ConsoleDisplay.display_agent_metrics(metrics)

    return simulation_results


def run_tasks(
    domain: str,
    tasks: list[Task],
    agent: str,
    user: str,
    llm_agent: Optional[str] = None,
    llm_args_agent: Optional[dict] = None,
    llm_user: Optional[str] = None,
    llm_args_user: Optional[dict] = None,
    num_trials: int = 1,
    max_steps: int = 100,
    max_errors: int = 10,
    save_to: Optional[str | Path] = None,
    console_display: bool = True,
    evaluation_type: EvaluationType = EvaluationType.ALL,
    max_concurrency: int = 1,
    seed: Optional[int] = 300,
    log_level: Optional[str] = "INFO",
    multi_user: bool = False,
    allow_multiple_user_messages: bool = False,
    num_tasks: int = 3,
    injection_task: Optional[str] = None,
) -> Results:
    """
    Runs tasks for a given domain.
    Automatically categorizes tasks by number of users (1, 2, or more than 2) and generates
    combinations with one task from each category. Each combination is run as a separate simulation.
    If llm_as_judge is True, the LLM will be used to annotate the simulation run.
    Calculates the reward for the simulation run.
    Args:
        domain (str): The domain to run the simulation on.
        tasks (list[Task]): The tasks to run. Will be categorized by user count.
        agent (str): The agent to run the simulation on.
        user (str): The user to run the simulation on.
        llm_agent (str): The model to use for the agent.
        llm_args_agent (dict): The arguments to pass to the LLM for the agent.
        llm_user (str): The model to use for the user.
        llm_args_user (dict): The arguments to pass to the LLM for the user.
        max_steps (int): The maximum number of steps to run the simulation.
        max_errors (int): The maximum number of errors to allow in the simulation.
        save_to (str | Path): The path to json file where to save the simulation results. If the file already exists, it will try to resume the run.
        evaluation_type (EvaluationType): The type of evaluation to use.
        max_concurrency (int): The maximum number of concurrent simulations to run.
        seed (int): The seed to use for the simulation.
        log_level (str): The log level to use.
        multi_user (bool): Whether to run the simulation in multi-user mode.
        allow_multiple_user_messages (bool): Whether to allow multiple user messages.
        num_tasks (int): The number of tasks to run in each combination.
    Returns:
        The simulation results and the annotations (if llm_review is True).
    """
    if isinstance(save_to, str):
        save_to = Path(save_to)
    # Set log level from config
    logger.remove()
    logger.add(lambda msg: print(msg), level=log_level)
    if len(tasks) == 0:
        raise ValueError("No tasks to run")
    if num_trials <= 0:
        raise ValueError("Number of trials must be greater than 0")
    if max_steps <= 0:
        raise ValueError("Max steps must be greater than 0")
    if max_errors <= 0:
        raise ValueError("Max errors must be greater than 0")

    random.seed(seed)

    seeds = [random.randint(0, 1000000) for _ in range(num_trials)]
    if "seed" in llm_args_agent:
        logger.warning("Each trial will modify the seed for the agent")

    if "seed" in llm_args_user:
        logger.warning("Each trial will modify the seed for the user")

    lock = multiprocessing.Lock()

    info = get_info(
        domain=domain,
        agent=agent,
        user=user,
        llm_agent=llm_agent,
        llm_args_agent=llm_args_agent,
        llm_user=llm_user,
        llm_args_user=llm_args_user,
        num_trials=num_trials,
        max_steps=max_steps,
        max_errors=max_errors,
        seed=seed,
    )
    # Store the original tasks for reference, but we'll run combinations
    simulation_results = Results(
        info=info,
        tasks=tasks,
        simulations=[],
    )
    done_runs = set()
    if save_to is not None:
        # If save_to already exists, check if the user wants to resume the run.
        if save_to.exists():
            # response = (
            #     ConsoleDisplay.console.input(
            #         "[yellow]File [bold]{}[/bold] already exists. Do you want to resume the run? (y/n)[/yellow] ".format(
            #             save_to
            #         )
            #     )
            #     .lower()
            #     .strip()
            # )
            # if response != "y":
            #     raise FileExistsError(
            #         f"File {save_to} already exists. Please delete it or use a different save_to name."
            #     )
            with open(save_to, "r") as fp:
                prev_simulation_results = Results.model_validate_json(fp.read())
                # Check if the run config has changed
                if get_pydantic_hash(prev_simulation_results.info) != get_pydantic_hash(
                    simulation_results.info
                ):
                    diff = show_dict_diff(
                        prev_simulation_results.info.model_dump(),
                        simulation_results.info.model_dump(),
                    )
                    ConsoleDisplay.console.print(
                        f"The run config has changed.\n\n{diff}\n\nDo you want to resume the run? (y/n)"
                    )
                    response = (
                        ConsoleDisplay.console.input(
                            "[yellow]File [bold]{}[/bold] already exists. Do you want to resume the run? (y/n)[/yellow] ".format(
                                save_to
                            )
                        )
                        .lower()
                        .strip()
                    )
                    if response != "y":
                        raise ValueError(
                            "The run config has changed. Please delete the existing file or use a different save_to name."
                        )
                # # Check if the task set has changed
                # if not all(
                #     get_pydantic_hash(task) == get_pydantic_hash(prev_task)
                #     for task, prev_task in zip(
                #         sorted(simulation_results.tasks, key=lambda x: x.id),
                #         sorted(prev_simulation_results.tasks, key=lambda x: x.id),
                #     )
                # ):
                #     raise ValueError(
                #         "The task set has changed. Please delete the existing file or use a different save_to name."
                #     )
                # Check which of the runs have already been done
                done_runs = set(
                    [
                        (sim.trial, ",".join(sim.task_ids), sim.seed)
                        for sim in prev_simulation_results.simulations
                    ]
                )
                simulation_results: Results = prev_simulation_results
                
                console_text = Text(text=f"Resuming run from {len(done_runs)} runs. {(len(tasks) // num_tasks + (len(tasks) % num_tasks != 0)) * num_trials - len(done_runs)} runs remaining.", style="bold yellow")
                ConsoleDisplay.console.print(console_text)
        # Create new save file
        else:
            # Check if save_to exists and create parent directories if needed
            if not save_to.parent.exists():
                save_to.parent.mkdir(parents=True, exist_ok=True)
            logger.info(f"Saving simulation batch to {save_to}")
            with open(save_to, "w") as fp:
                fp.write(simulation_results.model_dump_json(indent=2))

    def _save(simulation: SimulationRun):
        if save_to is None:
            return
        with lock:
            with open(save_to, "r") as fp:
                ckpt = json.load(fp)
            ckpt["simulations"].append(simulation.model_dump())
            with open(save_to, "w") as fp:
                json.dump(ckpt, fp, indent=2)

    def _run(tasks: list[Task], trial: int, seed: int, progress_str: str, injection_task_arg: Optional[str] = None) -> SimulationRun:
        console_text = Text(text=f"{progress_str}. Running tasks {', '.join([task.id for task in tasks])}, trial {trial + 1}", style="bold green")
        ConsoleDisplay.console.print(console_text)
        try:
            simulation = run_task_groups(
                domain=domain,
                tasks=tasks,
                agent=agent,
                user=user,
                llm_agent=llm_agent,
                llm_args_agent=llm_args_agent,
                llm_user=llm_user,
                llm_args_user=llm_args_user,
                max_steps=max_steps,
                max_errors=max_errors,
                evaluation_type=evaluation_type,
                seed=seed,
                multi_user=multi_user,
                allow_multiple_user_messages=allow_multiple_user_messages,
                injection_task=injection_task_arg,
            )
            simulation.trial = trial
            if console_display:
                ConsoleDisplay.display_simulation(simulation, show_details=False)
            _save(simulation)
        except Exception as e:
            logger.error(f"Error running tasks {', '.join([task.id for task in tasks])}, trial {trial}: {e}")
            raise e
        return simulation

    # Generate task combinations based on user count
    task_combinations = generate_task_combinations(tasks, multi_user, num_tasks)
    
    if not task_combinations:
        logger.error("No valid task combinations could be generated. Please check that tasks exist for all user categories (1, 2, and more than 2 users).")
        return simulation_results
    
    args = []
    for trial in range(num_trials):
        for i, task_combination in enumerate(task_combinations):
            if (trial, ",".join([task.id for task in task_combination]), seeds[trial]) in done_runs:
                console_text = Text(text=f"Skipping task {', '.join([task.id for task in task_combination])}, trial {trial} because it has already been run.", style="bold yellow")
                ConsoleDisplay.console.print(console_text)
                # find the simulation in simulation_results.simulations
                simulation = next(sim for sim in simulation_results.simulations if sim.task_ids == [task.id for task in task_combination])
                reevaluate = evaluate_multi_task_simulation(
                    simulation=simulation,
                    tasks=task_combination,
                    task_termination_reasons=simulation.task_termination_reasons,
                    evaluation_type=evaluation_type,
                    solo_mode=False,
                    domain=domain,
                    injection_task=injection_task,
                )
                simulation.task_rewards = reevaluate
                _save(simulation)
                continue
            progress_str = f"{i}/{len(task_combinations)} (trial {trial + 1}/{num_trials})"
            args.append((task_combination, trial, seeds[trial], progress_str, injection_task))

    for arg in args:
        try:
            res = _run(*arg)
            simulation_results.simulations.append(res)
        except Exception as e:
            logger.error(f"Error running task {', '.join([task.id for task in arg[0]])} , trial {arg[1]}: {e}")
            continue
        
    ConsoleDisplay.console.print(
        "\n✨ [bold green]Successfully completed all simulations![/bold green]\n"
        "To review the simulations, run: [bold blue]murmur view[/bold blue]"
    )
    return simulation_results


def run_task_groups(
    domain: str,
    tasks: list[Task],  # Changed from single task to list of tasks
    agent: str,
    user: str,
    llm_agent: Optional[str] = None,
    llm_args_agent: Optional[dict] = None,
    llm_user: Optional[str] = None,
    llm_args_user: Optional[dict] = None,
    max_steps: int = 100,
    max_errors: int = 10,
    evaluation_type: EvaluationType = EvaluationType.ALL,
    seed: Optional[int] = None,
    multi_user: bool = False,
    allow_multiple_user_messages: bool = False,
    injection_task: Optional[str] = None,
) -> SimulationRun:
    """
    Runs a list of tasks for a given domain with non-collaborative multi-user setup.
    Each task in the list will have its own group of users working collaboratively on that task.
    Users from different tasks do not share context, but the agent maintains context across all tasks.
     
     Args:
         domain (str): The domain to run the simulation on.
         tasks (list[Task]): The list of tasks to run non-collaboratively.
         agent (str): The agent to run the simulation on.
         user (str): The user to run the simulation on.
         llm_agent (str): The model to use for the agent.
         llm_args_agent (dict): The arguments to pass to the LLM for the agent.
         llm_user (str): The model to use for the user.
         llm_args_user (dict): The arguments to pass to the LLM for the user.
         max_steps (int): The maximum number of steps to run the simulation.
         max_errors (int): The maximum number of errors to allow in the simulation.
         evaluation_type (EvaluationType): The type of evaluation to use.
         seed (int): The seed to use for the simulation.
     Returns:
         The simulation run.
    """

    if max_steps <= 0:
        raise ValueError("Max steps must be greater than 0")
    if max_errors <= 0:
        raise ValueError("Max errors must be greater than 0")
    global registry
    
    # Validate that we have tasks and they're all multi-user
    if not tasks:
        raise ValueError("At least one task must be provided")
    
    task_ids = [task.id for task in tasks]
    logger.info(
        f"STARTING NON-COLLABORATIVE MULTI-TASK SIMULATION: Domain: {domain}, Tasks: {task_ids} with injection task: {injection_task}, Agent: {agent}, User: {user}"
    )
    
    environment_constructor = registry.get_env_constructor(domain)
    environment = environment_constructor()
    AgentConstructor = registry.get_agent_constructor(agent)

    solo_mode = False
    if issubclass(AgentConstructor, LLMAgent):
        agent = AgentConstructor(
            tools=environment.get_tools(),
            domain_policy=environment.get_policy(),
            llm=llm_agent,
            llm_args=llm_args_agent,
        )
    elif issubclass(AgentConstructor, LLMGTAgent):
        agent = AgentConstructor(
            tools=environment.get_tools(),
            domain_policy=environment.get_policy(),
            llm=llm_agent,
            llm_args=llm_args_agent,
            task=tasks[0],  # Use first task for backward compatibility
        )
    elif issubclass(AgentConstructor, LLMSoloAgent):
        solo_mode = True
        environment: Environment = environment_constructor(solo_mode=True)
        user_tools = environment.get_user_tools() if environment.user_tools else []
        agent = AgentConstructor(
            tools=environment.get_tools() + user_tools,
            domain_policy=environment.get_policy(),
            llm=llm_agent,
            llm_args=llm_args_agent,
            task=tasks[0],  # Use first task for backward compatibility
        )
    elif issubclass(AgentConstructor, LLMMurmurAgent):
        agent = AgentConstructor(
            tools=environment.get_tools(),
            domain_policy=environment.get_policy(),
            llm=llm_agent,
            llm_args=llm_args_agent,
        )
    else:
        raise ValueError(
            f"Unknown agent type: {AgentConstructor}. Should be LLMAgent or LLMSoloAgent"
        )
    try:
        user_tools = environment.get_user_tools()
    except Exception:
        user_tools = None

    # For non-collaborative multi-task setup, we always use multi-user mode
    # Create task to users mapping and collect all users
    task_to_users = {}  # task_id -> list of user_ids
    all_users = {}      # global user_id -> BaseUser mapping for the agent
    
    for task in tasks:
        # Extract multi-user information from each task
        multi_user_info = extract_multi_user_info(task)
        user_names = multi_user_info['user_names']
        user_roles = multi_user_info.get('user_roles', {})
        
        # Create users for this specific task
        task_user_ids = []
        for user_name in user_names:
            user_role = user_roles.get(user_name, 'customer')
            
            # Create unique user_id by combining random 5 letter string and user_name
            unique_user_id = f"{''.join(random.choices(string.ascii_letters, k=5))}_{user_name}"
            
            # Use MultiUserSimulator for multi-user scenarios
            if user == "user_simulator":
                if multi_user:
                    user_obj = MultiUserSimulator(
                        user_name=user_name,
                        user_role=user_role,
                        tools=user_tools,
                        instructions=str(task.user_scenario),
                        llm=llm_user,
                        llm_args=llm_args_user,
                        user_id=unique_user_id,
                    )
                else:
                    user_obj = UserSimulator(
                        tools=user_tools,
                        instructions=str(task.user_scenario),
                        llm=llm_user,
                        llm_args=llm_args_user,
                    )

            else:
                # Fallback to regular user constructor with user_name if supported
                UserConstructor = registry.get_user_constructor(user)
                try:
                    user_obj = UserConstructor(
                        user_name=user_name,
                        user_role=user_role,
                        tools=user_tools,
                        instructions=str(task.user_scenario),
                        llm=llm_user,
                        llm_args=llm_args_user,
                    )
                except TypeError:
                    # Fallback for constructors that don't support user_name
                    user_obj = UserConstructor(
                        tools=user_tools,
                        instructions=str(task.user_scenario),
                        llm=llm_user,
                        llm_args=llm_args_user,
                    )
            
            task_user_ids.append(unique_user_id)
            all_users[unique_user_id] = user_obj
        
        task_to_users[task.id] = task_user_ids
    
    # Handle injection task if provided
    injection_task_obj = None
    
    if injection_task:
        from murmur.registry import registry
        injection_tasks_loader = registry.get_injection_tasks_loader(domain)
        injection_tasks = injection_tasks_loader()
        
        # Find the specific injection task by ID
        injection_task_obj = None
        for task in injection_tasks:
            if task.id == injection_task:
                injection_task_obj = task
                break
        
        if injection_task_obj:
            logger.info(f"Using Task-based injection task: {injection_task}")
            print(injection_task_obj)
    
    # Use MultiUserOrchestrator with non-collaborative setup
    orchestrator = MultiUserOrchestrator(
        domain=domain,
        agent=agent,
        users=all_users,
        environment=environment,
        tasks=tasks,  # Pass the list of tasks
        task_to_users=task_to_users,  # Pass task to users mapping
        max_steps=max_steps,
        max_errors=max_errors,
        seed=seed,
        solo_mode=solo_mode,
        allow_multiple_user_messages=allow_multiple_user_messages,
        common_task=False,  # Non-collaborative mode
        injection_task_obj=injection_task_obj
    )
    simulation = orchestrator.run()

    # For multi-task evaluation, use the new multi-task evaluator
    task_rewards = evaluate_multi_task_simulation(
        simulation=simulation,
        tasks=tasks,
        task_termination_reasons=simulation.task_termination_reasons,
        evaluation_type=evaluation_type,
        solo_mode=solo_mode,
        domain=domain,
        injection_task=injection_task_obj,
    )
    
    simulation.task_rewards = task_rewards  # Store all task evaluations

    return simulation


def get_info(
    domain: str,
    agent: str,
    user: str,
    llm_agent: Optional[str] = None,
    llm_args_agent: Optional[dict] = None,
    llm_user: Optional[str] = None,
    llm_args_user: Optional[dict] = None,
    num_trials: int = 1,
    max_steps: int = 100,
    max_errors: int = 10,
    seed: Optional[int] = None,
) -> Info:
    user_info = UserInfo(
        implementation=user,
        llm=llm_user,
        llm_args=llm_args_user,
        global_simulation_guidelines=get_global_user_sim_guidelines(),
    )
    agent_info = AgentInfo(
        implementation=agent,
        llm=llm_agent,
        llm_args=llm_args_agent,
    )
    environment_info = get_environment_info(
        domain, include_tool_info=False
    )  # NOTE: Not saving tool info to avoid clutter.
    return Info(
        git_commit=get_commit_hash(),
        num_trials=num_trials,
        max_steps=max_steps,
        max_errors=max_errors,
        user_info=user_info,
        agent_info=agent_info,
        environment_info=environment_info,
        seed=seed,
    )
