import json
from copy import deepcopy

import pytest

from vita.config import (
    DEFAULT_LLM_AGENT,
    DEFAULT_LLM_ARGS_AGENT,
    DEFAULT_LLM_ARGS_USER,
    DEFAULT_LLM_USER,
)
from vita.data_model.simulation import RunConfig
from vita.data_model.tasks import EnvAssertion, RewardType, Task, make_task
from vita.run import (
    get_options,
    get_tasks,
    load_tasks,
    run_domain,
    run_task,
    run_tasks,
)


@pytest.fixture
def run_config() -> RunConfig:
    """Test that we can get available options from the registry"""
    return RunConfig(
        domain="mock",
        agent="llm_agent",
        user="user_simulator",
        task_ids=["create_task_1"],
        llm_agent="gpt-3.5-turbo",
        llm_args_agent={},
        llm_user="gpt-3.5-turbo",
        llm_args_user={},
        num_trials=3,
        max_steps=20,
        max_errors=10,
        save_to=None,
        llm_review=False,
        max_concurrency=3,
    )





def test_get_options():
    """Test that we can get available options from the registry"""
    options = get_options()
    assert options.domains is not None


def test_load_tasks():
    """Test that we can load tasks from the domain"""
    tasks = load_tasks("mock")
    assert len(tasks) > 0


def test_get_tasks():
    """Test that we can load tasks from the domain by id"""
    tasks = get_tasks("mock", task_ids=["create_task_1"])
    assert len(tasks) == 1
    assert tasks[0].id == "create_task_1"


def test_simplified_run(domain_name: str):
    """Test that we can run a task with the mock domain"""

    def run_simple_task(user_instruction: str, domain_name: str):
        task = make_task(
            user_instructions=user_instruction,
            eval_criteria=None,
        )
        simulation = run_tasks(
            domain=domain_name,
            tasks=[task],
            agent="llm_agent",
            user="user_simulator",
            llm_agent=DEFAULT_LLM_AGENT,
            llm_args_agent=deepcopy(DEFAULT_LLM_ARGS_AGENT),
            llm_user=DEFAULT_LLM_USER,
            llm_args_user=deepcopy(DEFAULT_LLM_ARGS_USER),
            max_steps=5,
            max_errors=5,
            evaluation_type="all",
            console_display=False,
            max_concurrency=1,
        )
        return simulation

    simulation = run_simple_task(
        user_instruction="create a task called 'test' for user_1",
        domain_name=domain_name,
    )
    assert simulation is not None


def test_run_tasks_base(domain_name: str, base_task: Task):
    """Test running a task with the mock domain"""
    results = run_tasks(
        domain=domain_name,
        tasks=[base_task],
        agent="llm_agent",
        user="user_simulator",
        llm_agent="gpt-3.5-turbo",
        llm_args_agent={},
        llm_user="gpt-3.5-turbo",
        llm_args_user={},
        max_concurrency=1,
    )
    # Check that simulation ran and has expected structure
    assert len(results.simulations) == 1
    simulation = results.simulations[0]
    assert simulation.messages is not None
    assert len(simulation.messages) > 0
    assert simulation.start_time is not None
    assert simulation.end_time is not None
    assert simulation.reward_info.reward is not None


def test_run_task_base(domain_name: str, base_task: Task):
    """Test running a task with the mock domain"""
    simulation = run_task(
        domain=domain_name,
        task=base_task,
        agent="llm_agent",
        user="user_simulator",
        llm_agent="gpt-3.5-turbo",
        llm_args_agent={},
        llm_user="gpt-3.5-turbo",
        llm_args_user={},
        evaluation_type="all",
    )
    # Check that simulation ran and has expected structure
    assert simulation.messages is not None
    assert len(simulation.messages) > 0
    assert simulation.start_time is not None
    assert simulation.end_time is not None
    assert simulation.reward_info.reward is not None


def test_run_tasks_message_history(domain_name: str, task_with_message_history: Task):
    """Test running a task with message history"""
    print(task_with_message_history.model_dump_json(indent=2))
    simulation = run_task(
        domain=domain_name,
        task=task_with_message_history,
        agent="llm_agent",
        user="user_simulator",
        llm_agent="gpt-3.5-turbo",
        llm_args_agent={},
        llm_user="gpt-3.5-turbo",
        llm_args_user={},
    )
    assert simulation is not None








def test_run_tasks_env_assertions(domain_name: str, task_with_env_assertions: Task):
    """Test running a task with env assertions"""
    simulation = run_task(
        domain=domain_name,
        task=task_with_env_assertions,
        agent="llm_agent",
        user="user_simulator",
        llm_agent="gpt-3.5-turbo",
        llm_args_agent={},
        llm_user="gpt-3.5-turbo",
        llm_args_user={},
        evaluation_type="all",
    )
    # Check that simulation ran and has expected structure
    assert simulation.messages is not None
    assert len(simulation.messages) > 0
    assert simulation.start_time is not None
    assert simulation.end_time is not None
    # Since env_assertions field has been removed, just check basic structure
    assert simulation.reward_info is not None
    assert simulation.reward_info.reward is not None


def test_run_tasks_history_and_env_assertions(
    domain_name: str, task_with_history_and_env_assertions: Task
):
    """Test running a task with history and env assertions"""
    simulation = run_task(
        domain=domain_name,
        task=task_with_history_and_env_assertions,
        agent="llm_agent",
        user="user_simulator",
        llm_agent="gpt-3.5-turbo",
        llm_args_agent={},
        llm_user="gpt-3.5-turbo",
        llm_args_user={},
    )
    assert simulation is not None


def test_run_tasks_nl_assertions(domain_name: str):
    """Test running a task with the mock domain"""
    task = get_tasks(domain_name, task_ids=["create_task_1_nl_eval"])[0]
    simulation = run_task(
        domain=domain_name,
        task=task,
        agent="llm_agent",
        user="user_simulator",
        llm_agent="gpt-3.5-turbo",
        llm_args_agent={},
        llm_user="gpt-3.5-turbo",
        llm_args_user={},
        evaluation_type="nl_rubrics",
    )
    # Check that simulation ran and has expected structure
    assert simulation.messages is not None
    assert len(simulation.messages) > 0
    assert simulation.start_time is not None
    assert simulation.end_time is not None
    # Since nl_rubrics field has been removed, just check basic structure
    assert simulation.reward_info is not None
    assert simulation.reward_info.reward is not None


def test_run_tasks_action_checks(domain_name: str, task_with_action_checks: Task):
    """Test running a task with action checks"""
    simulation = run_task(
        domain=domain_name,
        task=task_with_action_checks,
        agent="llm_agent",
        user="user_simulator",
        llm_agent="gpt-3.5-turbo",
        llm_args_agent={},
        llm_user="gpt-3.5-turbo",
        llm_args_user={},
    )
    assert simulation is not None
    # Since action_checks field has been removed, just check basic structure
    assert simulation.reward_info is not None
    assert simulation.reward_info.reward is not None


def test_run_domain(run_config: RunConfig):
    """Test running a domain with the mock domain
    Requires environment manager to be running
    """
    # Note: The function returns Results, not a tuple
    simulation_results = run_domain(run_config)
    assert simulation_results is not None


def test_re_run_validation():
    """Test that re-run parameter validation works correctly"""
    
    # Test that re_run without re_evaluate_file raises an error
    with pytest.raises(ValueError, match="--re-run can only be used with --re-evaluate-file"):
        config = RunConfig(
            domain="mock",
            re_run=True,
            task_ids=["task_1"]
        )
        config.validate()
    
    # Test that re_run without task_ids raises an error
    with pytest.raises(ValueError, match="--re-run requires --task-ids to specify which tasks to re-run"):
        config = RunConfig(
            domain="mock",
            re_run=True,
            re_evaluate_file="some_file.json"
        )
        config.validate()
    
    # Test that valid re_run configuration passes validation
    config = RunConfig(
        domain="mock",
        re_run=True,
        re_evaluate_file="some_file.json",
        task_ids=["task_1", "task_2"]
    )
    config.validate()  # Should not raise an error






