from model_library.registry_utils import get_registry_model
import asyncio

from vals import RunParameters, TestResult
from vals.graphql_client import LocalEvalUploadInputType
from vals.sdk.suite import Suite
from vals.sdk.types import Check, QuestionAnswerPair, Test
from dotenv import load_dotenv

load_dotenv()

llm = get_registry_model("google/gemini-2.5-flash")


async def run_agent():
    tests = [
        Test(
            input_under_test="What is QSBS?",
            checks=[Check(operator="equals", criteria="QSBS")],
        ),
        Test(
            input_under_test="What is an 83 election?",
            checks=[Check(operator="equals", criteria="QSBS")],
        ),
    ]
    suite = Suite(
        title="Test Suite",
        global_checks=[Check(operator="grammar")],
        tests=tests,
    )

    await suite.create()

    if not suite.id:
        raise Exception("Failed to create suite")

    print("Creating empty run object...")
    run = await Suite.create_run(
        suite_id=suite.id,
        parameters=RunParameters(
            eval_model="google/gemini-2.5-flash",
            create_text_summary=False,
            run_confidence_evaluation=False,
        ),
        run_name="Vals Agent Example",
    )
    print(f"Successfully created run object with ID: {run.id}")

    print("Uploading QA pairs...")
    qa_pairs: list[QuestionAnswerPair] = []
    for test in suite.tests:
        query_result = await llm.query(test.input_under_test)
        qa_pairs.append(
            await QuestionAnswerPair.upload(
                qa_set_id=run.qa_set_id,
                test_id=test._id,
                query_result=query_result,
            )
        )

    print(f"Uploaded {len(qa_pairs)} QA pairs.")

    await run.wait_for_run_completion()

    for qa_pair in qa_pairs:
        assert qa_pair.id

        local_eval = LocalEvalUploadInputType(
            questionAnswerPairId=qa_pair.id,
            name="Local Custom Operator",
            score=1,
            feedback="Great job!",
        )

        print(f"Uploading local eval for {qa_pair.id}...")
        await TestResult.upload(run.qa_set_id, [local_eval])

    print("Done!")


if __name__ == "__main__":
    asyncio.run(run_agent())
