from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from vals.graphql_client.enums import RunStatus
from vals.sdk.suite import Suite
from vals.sdk.types import QuestionAnswerPair


def make_mock_defaults():
    mock_defaults = MagicMock()
    mock_defaults.default_parameters.temperature = None
    mock_defaults.default_parameters.max_output_tokens = 2048
    mock_defaults.default_parameters.eval_model = "openai/gpt-4o"
    mock_defaults.default_parameters.model_under_test = "openai/gpt-4o"
    mock_defaults.default_parameters.maximum_threads = 10
    mock_defaults.default_parameters.eval_concurrency = 1
    return mock_defaults


def make_suite_with_mock_client() -> tuple[Suite, MagicMock]:
    """Create a Suite with a mocked graphql client."""
    suite = Suite(id="suite-123", title="test suite")
    mock_client = MagicMock()

    # Mock _create_empty_qa_set to return fake IDs
    mock_client.create_question_answer_set = AsyncMock(
        return_value=MagicMock(
            create_question_answer_set=MagicMock(
                question_answer_set=MagicMock(id="qa-set-456"),
                run=MagicMock(id="run-789"),
            )
        )
    )

    # Mock mark_question_answer_set_as_complete
    mock_client.mark_question_answer_set_as_complete = AsyncMock()

    # Mock start_run for the final start_run call
    mock_client.start_run = AsyncMock(
        return_value=MagicMock(start_run=MagicMock(run_id="run-789"))
    )

    # Mock get_default_parameters for RunParameters.to_graphql()
    mock_client.get_default_parameters = AsyncMock(return_value=make_mock_defaults())

    suite._client = mock_client
    return suite, mock_client


@pytest.mark.asyncio
async def test_on_run_created_called_for_callable_model():
    """on_run_created fires with run_id when model is a callable."""
    suite, mock_client = make_suite_with_mock_client()
    callback = MagicMock()

    async def dummy_model(test):
        return "answer"

    mock_run = MagicMock()
    mock_run.status = RunStatus.SUCCESS
    mock_run.refresh = AsyncMock()
    mock_run.wait_for_run_completion = AsyncMock()

    with (
        patch("vals.sdk.types.get_ariadne_client", return_value=mock_client),
        patch("vals.sdk.suite.get_ariadne_client", return_value=mock_client),
        patch.object(
            suite, "_safe_local_execution", new_callable=AsyncMock, return_value=[]
        ),
        patch(
            "vals.sdk.suite.Run.from_id", new_callable=AsyncMock, return_value=mock_run
        ),
    ):
        await suite.run(
            model=dummy_model,
            model_name="test-model",
            on_run_created=callback,
        )

    callback.assert_called_once_with("run-789")


@pytest.mark.asyncio
async def test_on_run_created_not_called_when_run_id_provided():
    """on_run_created should NOT fire when run_id and qa_set_id are pre-provided."""
    suite, mock_client = make_suite_with_mock_client()
    callback = MagicMock()

    async def dummy_model(test):
        return "answer"

    mock_run = MagicMock()
    mock_run.status = RunStatus.SUCCESS
    mock_run.refresh = AsyncMock()
    mock_run.wait_for_run_completion = AsyncMock()

    with (
        patch("vals.sdk.types.get_ariadne_client", return_value=mock_client),
        patch("vals.sdk.suite.get_ariadne_client", return_value=mock_client),
        patch.object(
            suite, "_safe_local_execution", new_callable=AsyncMock, return_value=[]
        ),
        patch(
            "vals.sdk.suite.Run.from_id", new_callable=AsyncMock, return_value=mock_run
        ),
    ):
        await suite.run(
            model=dummy_model,
            model_name="test-model",
            run_id="existing-run",
            qa_set_id="existing-qa-set",
            on_run_created=callback,
        )

    callback.assert_not_called()


@pytest.mark.asyncio
async def test_on_run_created_not_called_when_no_callback():
    """No error when on_run_created is None (default)."""
    suite, mock_client = make_suite_with_mock_client()

    async def dummy_model(test):
        return "answer"

    mock_run = MagicMock()
    mock_run.status = RunStatus.SUCCESS
    mock_run.refresh = AsyncMock()
    mock_run.wait_for_run_completion = AsyncMock()

    with (
        patch("vals.sdk.types.get_ariadne_client", return_value=mock_client),
        patch("vals.sdk.suite.get_ariadne_client", return_value=mock_client),
        patch.object(
            suite, "_safe_local_execution", new_callable=AsyncMock, return_value=[]
        ),
        patch(
            "vals.sdk.suite.Run.from_id", new_callable=AsyncMock, return_value=mock_run
        ),
    ):
        # Should not raise
        run = await suite.run(
            model=dummy_model,
            model_name="test-model",
        )

    assert run is mock_run


@pytest.mark.asyncio
async def test_on_run_created_called_for_list_model():
    """on_run_created fires with run_id when model is a list of QA pairs."""
    suite, mock_client = make_suite_with_mock_client()
    callback = MagicMock()

    qa_pairs = [
        QuestionAnswerPair(input_under_test="q1", llm_output="a1"),
    ]

    # Mock batch_add_question_answer_pairs
    mock_client.batch_add_question_answer_pairs = AsyncMock(
        return_value=MagicMock(
            batch_add_question_answer_pairs=MagicMock(question_answer_pairs=[])
        )
    )

    mock_run = MagicMock()
    mock_run.status = RunStatus.SUCCESS
    mock_run.refresh = AsyncMock()
    mock_run.wait_for_run_completion = AsyncMock()

    with (
        patch("vals.sdk.types.get_ariadne_client", return_value=mock_client),
        patch("vals.sdk.suite.get_ariadne_client", return_value=mock_client),
        patch(
            "vals.sdk.suite.Run.from_id", new_callable=AsyncMock, return_value=mock_run
        ),
    ):
        await suite.run(
            model=qa_pairs,
            model_name="test-model",
            on_run_created=callback,
        )

    callback.assert_called_once_with("run-789")
