import re
from unittest.mock import Mock, patch

import pandas as pd
import pytest

import mlflow
from mlflow.exceptions import MlflowException
from mlflow.genai.datasets.evaluation_dataset import EvaluationDataset
from mlflow.genai.simulators import (
    BaseSimulatedUserAgent,
    ConversationSimulator,
    SimulatedUserAgent,
    SimulatorContext,
)
from mlflow.genai.simulators.prompts import DEFAULT_PERSONA
from mlflow.genai.simulators.simulator import _MAX_METADATA_LENGTH
from mlflow.tracing.constant import TraceMetadataKey


def create_mock_evaluation_dataset(inputs: list[dict[str, object]]) -> Mock:
    mock_dataset = Mock(spec=EvaluationDataset)
    mock_dataset.to_df.return_value = pd.DataFrame({"inputs": inputs})
    return mock_dataset


def test_simulated_user_agent_generate_initial_message():
    with patch("mlflow.genai.simulators.simulator.invoke_model_without_tracing") as mock_invoke:
        mock_invoke.return_value = "Hello, I have a question about ML."

        agent = SimulatedUserAgent()
        context = SimulatorContext(
            goal="Learn about MLflow",
            persona="You are a beginner who asks curious questions.",
            conversation_history=[],
            turn=0,
        )

        message = agent.generate_message(context)

        assert message == "Hello, I have a question about ML."
        mock_invoke.assert_called_once()

        call_args = mock_invoke.call_args
        messages = call_args.kwargs["messages"]
        prompt = messages[0].content

        assert "Learn about MLflow" in prompt
        assert "beginner" in prompt


def test_simulated_user_agent_generate_followup_message():
    with patch("mlflow.genai.simulators.simulator.invoke_model_without_tracing") as mock_invoke:
        mock_invoke.return_value = "Can you tell me more?"

        agent = SimulatedUserAgent()
        context = SimulatorContext(
            goal="Learn about MLflow",
            persona="A helpful user",
            conversation_history=[
                {"role": "user", "content": "Hello"},
                {"role": "assistant", "content": "Hi there!"},
            ],
            turn=1,
        )

        message = agent.generate_message(context)

        assert message == "Can you tell me more?"
        mock_invoke.assert_called_once()

        call_args = mock_invoke.call_args
        messages = call_args.kwargs["messages"]
        prompt = messages[0].content

        assert "Hi there!" in prompt


def test_simulated_user_agent_default_persona():
    with patch("mlflow.genai.simulators.simulator.invoke_model_without_tracing") as mock_invoke:
        mock_invoke.return_value = "Test message"

        agent = SimulatedUserAgent()
        context = SimulatorContext(
            goal="Learn about ML",
            persona=DEFAULT_PERSONA,
            conversation_history=[],
            turn=0,
        )

        message = agent.generate_message(context)

        assert message == "Test message"

        call_args = mock_invoke.call_args
        messages = call_args.kwargs["messages"]
        prompt = messages[0].content

        assert "inquisitive user" in prompt.lower()


def test_conversation_simulator_basic_simulation(
    simple_test_case, mock_predict_fn, simulation_mocks
):
    # Each turn: generate_message + _check_goal_achieved
    simulation_mocks["invoke"].side_effect = [
        "What is MLflow?",  # turn 0 generate_message
        '{"rationale": "Goal not achieved yet", "result": "no"}',  # turn 0 goal check
        "Can you explain more?",  # turn 1 generate_message
        '{"rationale": "Goal not achieved yet", "result": "no"}',  # turn 1 goal check
    ]

    simulator = ConversationSimulator(
        test_cases=[simple_test_case],
        max_turns=2,
    )

    all_traces = simulator.simulate(mock_predict_fn)

    assert len(all_traces) == 1  # 1 test case
    assert len(all_traces[0]) == 2  # 2 traces
    assert all(t is simulation_mocks["trace"] for t in all_traces[0])
    assert simulation_mocks["invoke"].call_count == 4  # 2 turns * 2 calls each
    assert len(simulation_mocks["context_calls"]) == 2  # one per turn


def test_conversation_simulator_max_turns_stopping(
    simple_test_case, mock_predict_fn, simulation_mocks
):
    simulation_mocks["invoke"].side_effect = [
        "Test message",  # turn 0 generate_message
        '{"rationale": "Not yet", "result": "no"}',  # turn 0 goal check
        "Test message",  # turn 1 generate_message
        '{"rationale": "Not yet", "result": "no"}',  # turn 1 goal check
        "Test message",  # turn 2 generate_message
        '{"rationale": "Not yet", "result": "no"}',  # turn 2 goal check
    ]

    simulator = ConversationSimulator(
        test_cases=[simple_test_case],
        max_turns=3,
    )

    all_traces = simulator.simulate(mock_predict_fn)

    assert len(all_traces) == 1  # 1 test case
    assert len(all_traces[0]) == 3  # 3 traces
    assert simulation_mocks["invoke"].call_count == 6  # 3 turns * 2 calls each


def test_conversation_simulator_empty_response_stopping(simple_test_case, simulation_mocks):
    simulation_mocks["invoke"].return_value = "Test message"

    def empty_predict_fn(input=None, **kwargs):
        return {
            "output": [
                {
                    "id": "msg_123",
                    "type": "message",
                    "role": "assistant",
                    "content": [{"type": "output_text", "text": ""}],
                }
            ]
        }

    simulator = ConversationSimulator(
        test_cases=[simple_test_case],
        max_turns=5,
    )

    all_traces = simulator.simulate(empty_predict_fn)

    assert len(all_traces) == 1
    assert len(all_traces[0]) == 1  # Only 1 trace before stopping
    # Only generate_message called, goal check not called due to empty response
    assert simulation_mocks["invoke"].call_count == 1


def test_conversation_simulator_goal_achieved_stopping(
    simple_test_case, mock_predict_fn, simulation_mocks
):
    simulation_mocks["invoke"].side_effect = [
        "Test message",  # turn 0 generate_message
        '{"rationale": "Goal achieved!", "result": "yes"}',  # turn 0 goal check -> stop
    ]

    simulator = ConversationSimulator(
        test_cases=[simple_test_case],
        max_turns=5,
    )

    all_traces = simulator.simulate(mock_predict_fn)

    assert len(all_traces) == 1
    # Only 1 trace before goal was achieved
    assert len(all_traces[0]) == 1
    # 2 calls: generate_message + goal check
    assert simulation_mocks["invoke"].call_count == 2
    # Verify goal check was the second call with goal check prompt
    goal_check_call = simulation_mocks["invoke"].call_args_list[1]
    goal_check_prompt = goal_check_call.kwargs["messages"][0].content
    assert "achieved" in goal_check_prompt.lower()


def test_conversation_simulator_context_passing(test_case_with_context, simulation_mocks):
    simulation_mocks["invoke"].side_effect = [
        "Test message",
        '{"rationale": "Not achieved", "result": "no"}',
    ]

    captured_kwargs = {}

    def capturing_predict_fn(input=None, **kwargs):
        captured_kwargs.update(kwargs)
        return {
            "output": [
                {
                    "id": "msg_123",
                    "type": "message",
                    "role": "assistant",
                    "content": [{"type": "output_text", "text": "Response"}],
                }
            ]
        }

    simulator = ConversationSimulator(
        test_cases=[test_case_with_context],
        max_turns=1,
    )

    all_traces = simulator.simulate(capturing_predict_fn)

    assert len(all_traces) == 1
    assert len(all_traces[0]) == 1
    # Verify context was passed to predict_fn
    assert captured_kwargs.get("user_id") == "U001"
    assert captured_kwargs.get("session_id") == "S001"


def test_conversation_simulator_mlflow_session_id_passed_to_predict_fn(
    simple_test_case, simulation_mocks
):
    simulation_mocks["invoke"].side_effect = [
        "Test message",
        '{"rationale": "Not yet", "result": "no"}',
        "Test message 2",
        '{"rationale": "Not yet", "result": "no"}',
    ]

    captured_session_ids = []

    def capturing_predict_fn(input=None, **kwargs):
        captured_session_ids.append(kwargs.get("mlflow_session_id"))
        return {
            "output": [
                {
                    "id": "msg_123",
                    "type": "message",
                    "role": "assistant",
                    "content": [{"type": "output_text", "text": "Response"}],
                }
            ]
        }

    simulator = ConversationSimulator(
        test_cases=[simple_test_case],
        max_turns=2,
    )

    all_traces = simulator.simulate(capturing_predict_fn)

    assert len(all_traces) == 1
    assert len(all_traces[0]) == 2
    # Verify mlflow_session_id was passed to predict_fn
    assert len(captured_session_ids) == 2
    assert all(sid is not None for sid in captured_session_ids)
    assert all(sid.startswith("sim-") for sid in captured_session_ids)
    # Verify session ID is consistent across all turns in the same conversation
    assert captured_session_ids[0] == captured_session_ids[1]


def test_conversation_simulator_multiple_test_cases(
    simple_test_case, test_case_with_persona, mock_predict_fn, simulation_mocks
):
    # 2 test cases * 2 turns each * 2 calls per turn = 8 calls
    simulation_mocks["invoke"].side_effect = [
        "Test message",
        '{"rationale": "Not yet", "result": "no"}',
        "Test message",
        '{"rationale": "Not yet", "result": "no"}',
        "Test message",
        '{"rationale": "Not yet", "result": "no"}',
        "Test message",
        '{"rationale": "Not yet", "result": "no"}',
    ]

    simulator = ConversationSimulator(
        test_cases=[simple_test_case, test_case_with_persona],
        max_turns=2,
    )

    all_traces = simulator.simulate(mock_predict_fn)

    assert len(all_traces) == 2  # 2 test cases
    assert len(all_traces[0]) == 2  # 2 traces for first test case
    assert len(all_traces[1]) == 2  # 2 traces for second test case


@pytest.mark.parametrize(
    ("test_cases", "expected_error"),
    [
        ([], "test_cases cannot be empty"),
        ([{"persona": "test"}], r"indices \[0\].*'goal' field"),
        (
            [{"goal": "valid"}, {"persona": "missing goal"}],
            r"indices \[1\].*'goal' field",
        ),
        (
            [{"persona": "a"}, {"goal": "valid"}, {"persona": "b"}],
            r"indices \[0, 2\].*'goal' field",
        ),
    ],
    ids=[
        "empty_test_cases",
        "missing_goal",
        "second_case_missing_goal",
        "multiple_missing_goals",
    ],
)
def test_conversation_simulator_validation(test_cases, expected_error):
    with pytest.raises(ValueError, match=expected_error):
        ConversationSimulator(
            test_cases=test_cases,
            max_turns=2,
        )


@pytest.mark.parametrize(
    ("test_cases", "expected_error"),
    [
        (
            [{"goal": "test", "context": {"input": "foo"}}],
            r"indices \[0\].*reserved",
        ),
        (
            [{"goal": "test", "context": {"messages": []}}],
            r"indices \[0\].*reserved",
        ),
        (
            [{"goal": "test", "context": {"mlflow_session_id": "abc"}}],
            r"indices \[0\].*reserved",
        ),
        (
            [{"goal": "ok"}, {"goal": "test", "context": {"input": "bar"}}],
            r"indices \[1\].*reserved",
        ),
    ],
    ids=[
        "context_has_input",
        "context_has_messages",
        "context_has_mlflow_session_id",
        "second_case_has_reserved_key",
    ],
)
def test_conversation_simulator_rejects_reserved_context_keys(test_cases, expected_error):
    with pytest.raises(ValueError, match=expected_error):
        ConversationSimulator(
            test_cases=test_cases,
            max_turns=2,
        )


@pytest.mark.parametrize(
    ("test_cases", "expected_error"),
    [
        ([{"goal": "test", "context": "foo"}], r"indices \[0\].*'context' as a dict"),
        ([{"goal": "test", "context": ["foo"]}], r"indices \[0\].*'context' as a dict"),
        ([{"goal": "ok"}, {"goal": "test", "context": 1}], r"indices \[1\].*'context' as a dict"),
    ],
    ids=["context_string", "context_list", "second_case_context_int"],
)
def test_conversation_simulator_rejects_invalid_context_types(test_cases, expected_error):
    with pytest.raises(ValueError, match=expected_error):
        ConversationSimulator(
            test_cases=test_cases,
            max_turns=2,
        )


def test_conversation_simulator_accepts_dataframe_with_missing_context_values():
    test_cases_df = pd.DataFrame([
        {"goal": "Debug an error", "context": {"user_id": "U001"}},
        {"goal": "Learn about MLflow"},
    ])

    simulator = ConversationSimulator(
        test_cases=test_cases_df,
        max_turns=2,
    )

    assert len(simulator.test_cases) == 2
    assert simulator.test_cases[0]["context"] == {"user_id": "U001"}


@pytest.mark.parametrize(
    "inputs",
    [
        [{"goal": "Learn about MLflow"}],
        [{"goal": "Debug issue", "persona": "Engineer"}],
        [{"goal": "Ask questions", "persona": "Student", "context": {"id": "1"}}],
        [{"goal": "Learn ML", "simulation_guidelines": "Be concise"}],
        [{"goal": "Learn ML", "simulation_guidelines": ["Be concise", "Ask follow-ups"]}],
        [
            {
                "goal": "Debug deployment",
                "persona": "Engineer",
                "context": {"env": "prod"},
                "simulation_guidelines": "Focus on logs",
            }
        ],
    ],
)
def test_conversation_simulator_evaluation_dataset_valid(inputs):
    mock_dataset = create_mock_evaluation_dataset(inputs)
    simulator = ConversationSimulator(test_cases=mock_dataset, max_turns=2)
    assert len(simulator.test_cases) == len(inputs)
    assert simulator.test_cases == inputs


@pytest.mark.parametrize(
    "inputs",
    [
        [{"request": "What is MLflow?"}],
        [{"inputs": {"query": "Help me"}, "expected_response": "Sure!"}],
        [{"inputs": {"question": "How to log?", "answer": "Use mlflow.log"}}],
        [],
    ],
)
def test_conversation_simulator_evaluation_dataset_invalid(inputs):
    mock_dataset = create_mock_evaluation_dataset(inputs)
    with pytest.raises(ValueError, match="conversational test cases with a 'goal' field"):
        ConversationSimulator(test_cases=mock_dataset, max_turns=2)


def test_reassignment_with_valid_test_cases(simple_test_case):
    simulator = ConversationSimulator(test_cases=[simple_test_case], max_turns=2)
    new_test_cases = [
        {"goal": "New goal"},
    ]
    simulator.test_cases = new_test_cases
    assert simulator.test_cases == new_test_cases
    assert len(simulator.test_cases) == 1


def test_reassignment_with_dataframe(simple_test_case):
    simulator = ConversationSimulator(test_cases=[simple_test_case], max_turns=2)
    df = pd.DataFrame([{"goal": "Goal from DataFrame", "persona": "Analyst"}])
    simulator.test_cases = df
    assert simulator.test_cases == [{"goal": "Goal from DataFrame", "persona": "Analyst"}]


@pytest.mark.parametrize(
    ("invalid_test_cases", "expected_error"),
    [
        ([], "test_cases cannot be empty"),
        ([{"persona": "no goal here"}], r"indices \[0\].*'goal' field"),
    ],
)
def test_reassignment_with_invalid_test_cases_raises_error(
    simple_test_case, invalid_test_cases, expected_error
):
    simulator = ConversationSimulator(test_cases=[simple_test_case], max_turns=2)
    original_test_cases = simulator.test_cases
    with pytest.raises(ValueError, match=expected_error):
        simulator.test_cases = invalid_test_cases
    assert simulator.test_cases == original_test_cases


def test_simulator_context_is_first_turn():
    context_first = SimulatorContext(
        goal="Test goal",
        persona="Test persona",
        conversation_history=[],
        turn=0,
    )
    assert context_first.is_first_turn is True

    context_later = SimulatorContext(
        goal="Test goal",
        persona="Test persona",
        conversation_history=[{"role": "user", "content": "Hello"}],
        turn=1,
    )
    assert context_later.is_first_turn is False


def test_simulator_context_formatted_history():
    context_empty = SimulatorContext(
        goal="Test goal",
        persona="Test persona",
        conversation_history=[],
        turn=0,
    )
    assert context_empty.formatted_history is None

    context_with_history = SimulatorContext(
        goal="Test goal",
        persona="Test persona",
        conversation_history=[
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "Hi there!"},
        ],
        turn=1,
    )
    assert context_with_history.formatted_history == "user: Hello\nassistant: Hi there!"


def test_simulator_context_last_assistant_response():
    context_empty = SimulatorContext(
        goal="Test goal",
        persona="Test persona",
        conversation_history=[],
        turn=0,
    )
    assert context_empty.last_assistant_response is None

    context_with_history = SimulatorContext(
        goal="Test goal",
        persona="Test persona",
        conversation_history=[
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "Hi there!"},
        ],
        turn=1,
    )
    assert context_with_history.last_assistant_response == "Hi there!"


def test_simulator_context_is_frozen():
    context = SimulatorContext(
        goal="Test goal",
        persona="Test persona",
        conversation_history=[],
        turn=0,
    )
    with pytest.raises(AttributeError, match="cannot assign to field"):
        context.goal = "New goal"


def test_simulator_context_with_simulation_guidelines():
    context = SimulatorContext(
        goal="Test goal",
        persona="Test persona",
        conversation_history=[],
        turn=0,
        simulation_guidelines="Be concise and ask clarifying questions",
    )
    assert context.simulation_guidelines == "Be concise and ask clarifying questions"


def test_custom_user_agent_class(simple_test_case, mock_predict_fn, simulation_mocks):
    class CustomUserAgent(BaseSimulatedUserAgent):
        def generate_message(self, context: SimulatorContext) -> str:
            return f"Custom message for: {context.goal}"

    simulation_mocks["invoke"].return_value = '{"rationale": "Goal achieved!", "result": "yes"}'

    simulator = ConversationSimulator(
        test_cases=[simple_test_case],
        max_turns=2,
        user_agent_class=CustomUserAgent,
    )

    all_traces = simulator.simulate(mock_predict_fn)

    assert len(all_traces) == 1
    assert len(all_traces[0]) == 1


def test_user_agent_class_default(simple_test_case):
    simulator = ConversationSimulator(
        test_cases=[simple_test_case],
        max_turns=2,
    )
    assert simulator.user_agent_class is SimulatedUserAgent


def test_user_agent_class_receives_context(simple_test_case, mock_predict_fn, simulation_mocks):
    captured_contexts = []

    class ContextCapturingAgent(BaseSimulatedUserAgent):
        def generate_message(self, context: SimulatorContext) -> str:
            captured_contexts.append(context)
            return f"Message for turn {context.turn}"

    simulation_mocks["invoke"].return_value = '{"rationale": "Not yet", "result": "no"}'

    simulator = ConversationSimulator(
        test_cases=[simple_test_case],
        max_turns=2,
        user_agent_class=ContextCapturingAgent,
    )

    simulator.simulate(mock_predict_fn)

    assert len(captured_contexts) == 2
    assert captured_contexts[0].turn == 0
    assert captured_contexts[0].is_first_turn is True
    assert captured_contexts[0].goal == simple_test_case["goal"]
    assert captured_contexts[1].turn == 1
    assert captured_contexts[1].is_first_turn is False


def test_conversation_simulator_sets_simulation_metadata(mock_predict_fn_with_context):
    long_goal = "A" * 500
    long_persona = "B" * 500
    context = {"user_id": "U001", "session_id": "S001"}

    with patch("mlflow.genai.simulators.simulator.invoke_model_without_tracing") as mock_invoke:
        mock_invoke.side_effect = [
            "Test message",
            '{"rationale": "Not achieved", "result": "no"}',
            "Follow up message",
            '{"rationale": "Goal achieved!", "result": "yes"}',
        ]

        simulator = ConversationSimulator(
            test_cases=[{"goal": long_goal, "persona": long_persona, "context": context}],
            max_turns=2,
        )

        all_traces = simulator.simulate(mock_predict_fn_with_context)
        first_test_case_traces = all_traces[0]

        assert len(first_test_case_traces) == 2

        for trace in first_test_case_traces:
            metadata = trace.info.trace_metadata

            assert TraceMetadataKey.TRACE_SESSION in metadata
            assert metadata["mlflow.simulation.goal"] == long_goal[:_MAX_METADATA_LENGTH]
            assert metadata["mlflow.simulation.persona"] == long_persona[:_MAX_METADATA_LENGTH]


def test_conversation_simulator_uses_default_persona(mock_predict_fn):
    with patch("mlflow.genai.simulators.simulator.invoke_model_without_tracing") as mock_invoke:
        mock_invoke.side_effect = [
            "Test message",
            '{"rationale": "Goal achieved!", "result": "yes"}',
        ]

        simulator = ConversationSimulator(
            test_cases=[{"goal": "Test goal"}],
            max_turns=1,
        )

        all_traces = simulator.simulate(mock_predict_fn)

        trace = all_traces[0][0]
        metadata = trace.info.trace_metadata

        assert metadata["mlflow.simulation.goal"] == "Test goal"
        assert metadata["mlflow.simulation.persona"] == DEFAULT_PERSONA
        assert TraceMetadataKey.TRACE_SESSION in metadata


def test_conversation_simulator_logs_expectations_to_first_trace(mock_predict_fn):
    expectations = {"expected_topic": "MLflow", "expected_sentiment": "positive"}

    with patch("mlflow.genai.simulators.simulator.invoke_model_without_tracing") as mock_invoke:
        mock_invoke.side_effect = [
            "Test message",
            '{"rationale": "Not achieved", "result": "no"}',
            "Follow up message",
            '{"rationale": "Goal achieved!", "result": "yes"}',
        ]

        simulator = ConversationSimulator(
            test_cases=[{"goal": "Test goal", "expectations": expectations}],
            max_turns=2,
        )

        all_traces = simulator.simulate(mock_predict_fn)

        assert len(all_traces[0]) == 2

        first_trace = all_traces[0][0]
        expectation_assessments = [
            a for a in first_trace.info.assessments if a.expectation is not None
        ]

        assert len(expectation_assessments) == 2
        for assessment in expectation_assessments:
            assert assessment.name in expectations
            assert assessment.expectation.value == expectations[assessment.name]
            assert TraceMetadataKey.TRACE_SESSION in assessment.metadata

        second_trace = all_traces[0][1]
        second_trace_assessments = second_trace.info.assessments
        second_expectation_assessments = [
            a for a in second_trace_assessments if a.expectation is not None
        ]
        assert len(second_expectation_assessments) == 0


def test_invoke_llm_with_prompt_only():
    with patch("mlflow.genai.simulators.simulator.invoke_model_without_tracing") as mock_invoke:
        mock_invoke.return_value = "LLM response"

        agent = SimulatedUserAgent()
        result = agent.invoke_llm("Test prompt")

        assert result == "LLM response"
        mock_invoke.assert_called_once()
        call_args = mock_invoke.call_args
        messages = call_args.kwargs["messages"]
        assert len(messages) == 1
        assert messages[0].role == "user"
        assert messages[0].content == "Test prompt"


def test_invoke_llm_with_system_prompt():
    with patch("mlflow.genai.simulators.simulator.invoke_model_without_tracing") as mock_invoke:
        mock_invoke.return_value = "LLM response with system"

        agent = SimulatedUserAgent()
        result = agent.invoke_llm("Test prompt", system_prompt="System instructions")

        assert result == "LLM response with system"
        mock_invoke.assert_called_once()
        call_args = mock_invoke.call_args
        messages = call_args.kwargs["messages"]
        assert len(messages) == 2
        assert messages[0].role == "system"
        assert messages[0].content == "System instructions"
        assert messages[1].role == "user"
        assert messages[1].content == "Test prompt"


def test_invalid_user_agent_class_raises_type_error(simple_test_case):
    class NotAUserAgent:
        pass

    with pytest.raises(TypeError, match="must be a subclass of BaseSimulatedUserAgent"):
        ConversationSimulator(
            test_cases=[simple_test_case],
            max_turns=2,
            user_agent_class=NotAUserAgent,
        )


def test_conversation_simulator_digest_is_deterministic():
    test_cases = [
        {"goal": "Learn about MLflow"},
        {"goal": "Debug deployment", "persona": "Data scientist"},
        {"goal": "Setup", "context": {"env": "prod"}},
    ]
    simulator1 = ConversationSimulator(test_cases=test_cases, max_turns=2)
    simulator2 = ConversationSimulator(test_cases=test_cases, max_turns=2)

    digest1 = simulator1._compute_test_case_digest()
    digest2 = simulator2._compute_test_case_digest()

    assert digest1 == digest2
    assert isinstance(digest1, str)
    assert len(digest1) == 8


@pytest.mark.parametrize(
    ("test_cases_1", "test_cases_2"),
    [
        # Different goals
        ([{"goal": "Goal A"}], [{"goal": "Goal B"}]),
        # Adding persona changes digest
        ([{"goal": "Goal"}], [{"goal": "Goal", "persona": "Engineer"}]),
        # Different order
        ([{"goal": "A"}, {"goal": "B"}], [{"goal": "B"}, {"goal": "A"}]),
    ],
    ids=["different_goals", "added_persona", "different_order"],
)
def test_conversation_simulator_digest_differs_for_different_test_cases(test_cases_1, test_cases_2):
    simulator1 = ConversationSimulator(test_cases=test_cases_1, max_turns=2)
    simulator2 = ConversationSimulator(test_cases=test_cases_2, max_turns=2)

    assert simulator1._compute_test_case_digest() != simulator2._compute_test_case_digest()


def test_conversation_simulator_get_dataset_name_default():
    test_cases = [{"goal": "Learn about MLflow"}]
    simulator = ConversationSimulator(test_cases=test_cases, max_turns=2)

    assert simulator._get_dataset_name() == "conversational_dataset"


def test_conversation_simulator_get_dataset_name_from_evaluation_dataset():
    inputs = [{"goal": "Learn about MLflow"}]
    mock_dataset = create_mock_evaluation_dataset(inputs)
    mock_dataset.name = "my_custom_dataset"

    simulator = ConversationSimulator(test_cases=mock_dataset, max_turns=2)

    assert simulator._get_dataset_name() == "my_custom_dataset"


def test_simulate_creates_run_when_no_parent_run(
    tmp_path, simple_test_case, mock_predict_fn, simulation_mocks
):
    simulation_mocks["invoke"].side_effect = [
        "Test message",
        '{"rationale": "Goal achieved!", "result": "yes"}',
    ]

    mlflow.set_tracking_uri(f"sqlite:///{tmp_path}/mlflow.db")
    mlflow.set_experiment("test-experiment")

    simulator = ConversationSimulator(test_cases=[simple_test_case], max_turns=1)
    simulator.simulate(mock_predict_fn)

    runs = mlflow.search_runs()
    assert len(runs) == 1
    run_name = runs.iloc[0]["tags.mlflow.runName"]
    assert re.match(r"^simulation-[0-9a-f]{8}$", run_name)


def test_simulate_uses_parent_run_when_exists(
    tmp_path, simple_test_case, mock_predict_fn, simulation_mocks
):
    simulation_mocks["invoke"].side_effect = [
        "Test message",
        '{"rationale": "Goal achieved!", "result": "yes"}',
    ]

    mlflow.set_tracking_uri(f"sqlite:///{tmp_path}/mlflow.db")
    mlflow.set_experiment("test-experiment")

    with mlflow.start_run(run_name="parent-run") as parent_run:
        parent_run_id = parent_run.info.run_id

        simulator = ConversationSimulator(test_cases=[simple_test_case], max_turns=1)
        simulator.simulate(mock_predict_fn)

        assert mlflow.active_run().info.run_id == parent_run_id

    runs = mlflow.search_runs()
    assert len(runs) == 1
    assert runs.iloc[0]["tags.mlflow.runName"] == "parent-run"


def test_simulate_run_name_format(tmp_path, simple_test_case, mock_predict_fn, simulation_mocks):
    simulation_mocks["invoke"].side_effect = [
        "Test message",
        '{"rationale": "Goal achieved!", "result": "yes"}',
    ]

    mlflow.set_tracking_uri(f"sqlite:///{tmp_path}/mlflow.db")
    mlflow.set_experiment("test-experiment")

    simulator = ConversationSimulator(test_cases=[simple_test_case], max_turns=1)
    simulator.simulate(mock_predict_fn)

    runs = mlflow.search_runs()
    run_name = runs.iloc[0]["tags.mlflow.runName"]

    assert run_name.startswith("simulation-")
    hex_part = run_name[len("simulation-") :]
    assert len(hex_part) == 8
    assert re.match(r"^[0-9a-f]+$", hex_part)


def test_conversation_simulator_completions_messages_format(simple_test_case, simulation_mocks):
    simulation_mocks["invoke"].side_effect = [
        "Test message",
        '{"rationale": "Goal achieved!", "result": "yes"}',
    ]

    captured_messages_snapshots = []

    def capturing_predict_fn(messages: list[dict[str, str]] | None = None, **kwargs):
        captured_messages_snapshots.append(list(messages) if messages else None)
        return {
            "output": [
                {
                    "id": "msg_123",
                    "type": "message",
                    "role": "assistant",
                    "content": [{"type": "output_text", "text": "Response"}],
                }
            ]
        }

    simulator = ConversationSimulator(test_cases=[simple_test_case], max_turns=1)
    simulator.simulate(capturing_predict_fn)

    assert len(captured_messages_snapshots) == 1
    assert captured_messages_snapshots[0][0]["role"] == "user"
    assert captured_messages_snapshots[0][0]["content"] == "Test message"


def test_conversation_simulator_rejects_both_input_and_messages(simple_test_case, simulation_mocks):
    simulation_mocks["invoke"].return_value = "Test message"

    def invalid_predict_fn(input: list[dict[str, str]], messages: list[dict[str, str]], **kwargs):
        return {
            "output": [
                {"role": "assistant", "content": [{"type": "output_text", "text": "Response"}]}
            ]
        }

    simulator = ConversationSimulator(test_cases=[simple_test_case], max_turns=1)

    with pytest.raises(MlflowException, match="cannot have both 'messages' and 'input' parameters"):
        simulator.simulate(invalid_predict_fn)


def test_conversation_simulator_rejects_neither_input_nor_messages(
    simple_test_case, simulation_mocks
):
    simulation_mocks["invoke"].return_value = "Test message"

    def invalid_predict_fn(**kwargs):
        return None

    simulator = ConversationSimulator(test_cases=[simple_test_case], max_turns=1)

    with pytest.raises(MlflowException, match="must accept either 'messages' or 'input'"):
        simulator.simulate(invalid_predict_fn)


def test_simulated_user_agent_with_simulation_guidelines():
    with patch("mlflow.genai.simulators.simulator.invoke_model_without_tracing") as mock_invoke:
        mock_invoke.return_value = "I have a question about ML pipelines."

        agent = SimulatedUserAgent()
        context = SimulatorContext(
            goal="Learn about ML pipelines",
            persona=DEFAULT_PERSONA,
            conversation_history=[],
            turn=0,
            simulation_guidelines="Ask clarifying questions before proceeding",
        )

        message = agent.generate_message(context)

        assert message == "I have a question about ML pipelines."
        mock_invoke.assert_called_once()

        call_args = mock_invoke.call_args
        messages = call_args.kwargs["messages"]
        prompt = messages[0].content

        assert "Learn about ML pipelines" in prompt
        assert "Ask clarifying questions before proceeding" in prompt
        assert "simulation_guidelines" in prompt


def test_simulated_user_agent_followup_with_simulation_guidelines():
    with patch("mlflow.genai.simulators.simulator.invoke_model_without_tracing") as mock_invoke:
        mock_invoke.return_value = "Let me clarify something first."

        agent = SimulatedUserAgent()
        context = SimulatorContext(
            goal="Learn about ML",
            persona=DEFAULT_PERSONA,
            conversation_history=[
                {"role": "user", "content": "Hello"},
                {"role": "assistant", "content": "Hi there!"},
            ],
            turn=1,
            simulation_guidelines="Be thorough and ask follow-up questions",
        )

        message = agent.generate_message(context)

        assert message == "Let me clarify something first."
        mock_invoke.assert_called_once()

        call_args = mock_invoke.call_args
        messages = call_args.kwargs["messages"]
        prompt = messages[0].content

        assert "Be thorough and ask follow-up questions" in prompt
        assert "simulation_guidelines" in prompt


def test_conversation_simulator_with_simulation_guidelines(mock_predict_fn):
    test_case = {
        "goal": "Learn about ML pipelines",
        "simulation_guidelines": "Ask clarifying questions before proceeding",
    }

    with patch("mlflow.genai.simulators.simulator.invoke_model_without_tracing") as mock_invoke:
        mock_invoke.side_effect = [
            "Test message with simulation_guidelines",
            '{"rationale": "Goal achieved!", "result": "yes"}',
        ]

        simulator = ConversationSimulator(
            test_cases=[test_case],
            max_turns=2,
        )

        all_traces = simulator.simulate(mock_predict_fn)

        assert len(all_traces) == 1
        assert len(all_traces[0]) == 1

        # Verify simulation_guidelines are in the generate_message prompt
        generate_call = mock_invoke.call_args_list[0]
        prompt = generate_call.kwargs["messages"][0].content
        assert "Ask clarifying questions before proceeding" in prompt

        trace = all_traces[0][0]
        metadata = trace.info.trace_metadata
        assert "mlflow.simulation.simulation_guidelines" in metadata
        assert (
            metadata["mlflow.simulation.simulation_guidelines"]
            == "Ask clarifying questions before proceeding"
        )
        assert TraceMetadataKey.TRACE_SESSION in metadata