"""Unit tests for agentic upload helpers."""

from types import SimpleNamespace

import pytest
from model_library.base import QueryResult, QueryResultMetadata

from vals.graphql_client.input_types import LocalEvalUploadInputType
from vals.graphql_client.upload_test_result import (
    UploadTestResultUploadLocalEvaluationTestResults,
    UploadTestResultUploadLocalEvaluationTestResultsResultJson,
)
from vals.sdk.types import Check, QuestionAnswerPair, Test, TestResult


class _FakeAriadneClient:
    def __init__(self):
        self.create_calls = []
        self.upload_calls = []
        self.create_response: SimpleNamespace | None = None
        self.upload_response: SimpleNamespace | None = None

    async def create_question_answer_pair(self, qa_set_id, qa_pair):
        self.create_calls.append((qa_set_id, qa_pair))
        return self.create_response

    async def upload_test_result(self, qa_set_id, local_evals):
        self.upload_calls.append((qa_set_id, local_evals))
        return self.upload_response


@pytest.mark.asyncio
async def test_question_answer_pair_upload(monkeypatch):
    fake_client = _FakeAriadneClient()

    fake_metadata = SimpleNamespace(
        in_tokens=10,
        out_tokens=5,
        duration_seconds=1.5,
        cost=None,
    )
    fake_graphql_qa_pair = SimpleNamespace(
        id="qa-123",
        input_under_test="prompt",
        llm_output="answer",
        context={"subject": "QSBS"},
        output_context={"debug": True},
        metadata=fake_metadata,
        file_ids=["file-1"],
        local_evals=[],
        test=SimpleNamespace(id="test-123"),
        status="SUCCESS",
    )

    fake_client.create_response = SimpleNamespace(
        create_question_answer_pair=SimpleNamespace(
            question_answer_pair=fake_graphql_qa_pair
        )
    )

    monkeypatch.setattr("vals.sdk.types.get_ariadne_client", lambda: fake_client)

    test = Test(input_under_test="prompt", checks=[Check(operator="equals")])
    test._id = "test-123"
    test._file_ids = ["file-1"]
    test.context = {"subject": "QSBS"}

    query_result = QueryResult(
        output_text="answer",
        metadata=QueryResultMetadata(in_tokens=10, out_tokens=5, duration_seconds=1.5),
        raw={"debug": True},
    )

    qa_pair = await QuestionAnswerPair.upload(
        qa_set_id="qa-set-1",
        test_id=test._id,
        query_result=query_result,
    )

    assert qa_pair.id == "qa-123"
    assert qa_pair.file_ids == ["file-1"]
    assert qa_pair.context == {"subject": "QSBS"}
    assert qa_pair.output_context == query_result.raw
    assert qa_pair.metadata and qa_pair.metadata.in_tokens == 10


def _build_graphql_test_result():
    result_json = (
        UploadTestResultUploadLocalEvaluationTestResultsResultJson.model_validate(
            {
                "operator": "equals",
                "criteria": "expected",
                "modifiers": {
                    "optional": False,
                    "severity": 1.0,
                    "displayMetrics": True,
                    "examples": [],
                    "extractor": None,
                    "conditional": None,
                    "category": None,
                },
                "autoEval": 1.0,
                "feedback": "Looks good",
                "evalCont": 1.0,
                "isGlobal": False,
                "severity": 1.0,
                "customOperator": None,
            }
        )
    )

    return UploadTestResultUploadLocalEvaluationTestResults.model_validate(
        {
            "id": "tr-1",
            "llmOutput": "answer",
            "resultJson": [result_json.model_dump(by_alias=True)],
            "qaPair": {
                "context": {"subject": "QSBS"},
                "outputContext": {"debug": True},
                "errorMessage": "",
            },
            "test": {
                "id": "test-123",
                "inputUnderTest": "prompt",
                "context": {"subject": "QSBS"},
                "tags": [],
                "crossVersionId": "cv-1",
                "goldenOutput": "",
                "fileIds": [],
                "checks": [
                    {
                        "operator": "equals",
                        "criteria": "expected",
                        "modifiers": {
                            "optional": False,
                            "severity": None,
                            "examples": [],
                            "extractor": None,
                            "conditional": None,
                            "category": None,
                            "displayMetrics": False,
                        },
                    }
                ],
                "testSuite": {"id": "suite-1"},
            },
            "metadata": {
                "inTokens": 10,
                "outTokens": 5,
                "durationSeconds": 1.5,
                "cost": None,
                "tokenRetryMetadata": None,
            },
            "passPercentage": 1.0,
            "passPercentageWithWeight": 1.0,
        }
    )


@pytest.mark.asyncio
async def test_test_result_upload(monkeypatch):
    fake_client = _FakeAriadneClient()
    graphql_test_result = _build_graphql_test_result()
    fake_client.upload_response = SimpleNamespace(
        upload_local_evaluation=SimpleNamespace(test_results=[graphql_test_result])
    )
    monkeypatch.setattr("vals.sdk.types.get_ariadne_client", lambda: fake_client)

    local_eval = LocalEvalUploadInputType(
        questionAnswerPairId="qa-123",
        name="Local Eval",
        score=1,
        feedback="Great job!",
    )

    result = await TestResult.upload("qa-set-1", [local_eval])

    assert result.llm_output == "answer"
    assert len(result.check_results) == 1
    assert result.metadata and result.metadata.in_tokens == 10
    assert fake_client.upload_calls[0][0] == "qa-set-1"
    assert fake_client.upload_calls[0][1][0] is local_eval


@pytest.mark.asyncio
async def test_test_result_upload_rejects_multiple_ids():
    local_eval_one = LocalEvalUploadInputType(
        questionAnswerPairId="qa-1",
        name="Eval 1",
        score=1,
        feedback="",
    )
    local_eval_two = LocalEvalUploadInputType(
        questionAnswerPairId="qa-2",
        name="Eval 2",
        score=1,
        feedback="",
    )

    with pytest.raises(ValueError):
        await TestResult.upload("qa-set-1", [local_eval_one, local_eval_two])


@pytest.mark.asyncio
async def test_test_result_upload_empty_response(monkeypatch):
    fake_client = _FakeAriadneClient()
    fake_client.upload_response = SimpleNamespace(upload_local_evaluation=None)
    monkeypatch.setattr("vals.sdk.types.get_ariadne_client", lambda: fake_client)

    local_eval = LocalEvalUploadInputType(
        questionAnswerPairId="qa-123",
        name="Local Eval",
        score=1,
        feedback="Great job!",
    )

    with pytest.raises(ValueError):
        await TestResult.upload("qa-set-1", [local_eval])


@pytest.mark.asyncio
async def test_test_result_upload_no_test_results(monkeypatch):
    fake_client = _FakeAriadneClient()
    fake_client.upload_response = SimpleNamespace(
        upload_local_evaluation=SimpleNamespace(test_results=[])
    )
    monkeypatch.setattr("vals.sdk.types.get_ariadne_client", lambda: fake_client)

    local_eval = LocalEvalUploadInputType(
        questionAnswerPairId="qa-123",
        name="Local Eval",
        score=1,
        feedback="Great job!",
    )

    with pytest.raises(ValueError):
        await TestResult.upload("qa-set-1", [local_eval])
