"""
Test script for chatCompletion GraphQL mutation.
"""

import asyncio
import base64
import json

from gql import Client, gql
from gql.transport.aiohttp import AIOHTTPTransport

MUTATION = gql("""
mutation ChatCompletionMutation($input: ChatCompletionInput!) {
  chatCompletion(input: $input) {
    repetitions {
      repetitionNumber
      content
      toolCalls {
        id
        function {
          name
          arguments
        }
      }
      span {
        id
      }
      errorMessage
      evaluations {
        name
        label
        score
        annotatorKind
        explanation
        metadata
      }
    }
  }
}
""")


async def main() -> None:
    transport = AIOHTTPTransport(url="http://localhost:6006/graphql")

    # The evaluator ID from test_create_llm_evaluator.py
    llm_evaluator_id = base64.b64encode(b"LLMEvaluator:1").decode("utf-8")
    built_in_evaluator_id = "QnVpbHRJbkV2YWx1YXRvcjotMjAwNTY2NTgzMgo="

    evaluators = [
        {
            "id": llm_evaluator_id,
            "inputMapping": {
                "pathMapping": {
                    "input": "$.input",
                    "output": "$.output",
                },
            },
        },
        {
            "id": built_in_evaluator_id,
            "inputMapping": {
                "pathMapping": {
                    "text": "$.output",
                },
                "literalMapping": {
                    "words": "Paris",
                    "case_sensitive": True,
                },
            },
        },
    ]

    # Build promptVersion (model identity + template); clientOptions optional.
    messages = [
        {"content": "You are a chatbot", "role": "SYSTEM"},
        {"content": "{{question}}", "role": "USER"},
    ]
    prompt_version = {
        "templateFormat": "MUSTACHE",
        "template": {
            "messages": [
                {
                    "role": m["role"],
                    "content": [{"text": {"text": m["content"]}}],
                }
                for m in messages
            ]
        },
        "modelProvider": "OPENAI",
        "modelName": "gpt-4o-mini",
        "invocationParameters": {},
        "tools": None,
        "responseFormat": None,
    }

    variables = {
        "input": {
            "promptVersion": prompt_version,
            "repetitions": 1,
            "template": {
                "variables": {"question": "What is the capital of France?"},
                "format": "MUSTACHE",
            },
            "evaluators": evaluators,
        }
    }

    async with Client(transport=transport, fetch_schema_from_transport=False) as session:
        result = await session.execute(MUTATION, variable_values=variables)
        print(json.dumps(result, indent=2))


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