from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from vals.graphql_client.enums import RunStatus
from vals.graphql_client.input_types import (
    MetadataType,
    QuestionAnswerPairInputType,
)
from vals.sdk.suite import Suite
from vals.sdk.types import Test


def make_qa_pair(index: int) -> QuestionAnswerPairInputType:
    return QuestionAnswerPairInputType(
        inputUnderTest=f"input-{index}",
        llmOutput=f"output-{index}",
        metadata=MetadataType(inTokens=0, outTokens=0, durationSeconds=0),
        testId=f"test-{index}",
        status="success",
    )


@pytest.mark.asyncio
async def test_periodic_uploader_batches_at_most_200():
    num_tests = 350
    suite = Suite(
        id="suite-123",
        title="test suite",
        tests=[
            Test(input_under_test=f"input-{i}", checks=[]) for i in range(num_tests)
        ],
    )
    mock_client = MagicMock()

    mock_client.get_run_status = AsyncMock(
        return_value=MagicMock(run=MagicMock(status=RunStatus.IN_PROGRESS))
    )

    batch_calls: list[list[QuestionAnswerPairInputType]] = []

    async def capture_batch(run_id, qa_set_id, batch):
        batch_calls.append(list(batch))
        mock_response = MagicMock()
        mock_response.batch_add_question_answer_pairs.question_answer_pairs = []
        return mock_response

    mock_client.batch_add_question_answer_pairs = AsyncMock(side_effect=capture_batch)

    suite._client = mock_client

    async def mock_process_single_test(test, model_function, is_simple):
        return make_qa_pair(0)

    param_input = MagicMock()
    param_input.maximum_threads = 10

    async def dummy_model(inp):
        return "answer"

    with (
        patch.object(
            Suite,
            "from_id",
            new_callable=AsyncMock,
            return_value=suite,
        ),
        patch.object(
            suite,
            "_process_single_test",
            new_callable=AsyncMock,
            side_effect=mock_process_single_test,
        ),
    ):
        await suite._generate_qa_pairs_from_function(
            model_function=dummy_model,
            parameter_input=param_input,
            run_id="run-123",
            qa_set_id="qa-set-456",
        )

    assert len(batch_calls) > 0
    for batch in batch_calls:
        assert len(batch) <= 200

    total_uploaded = sum(len(b) for b in batch_calls)
    assert total_uploaded == num_tests
