"""
Simple script to test the chat completion GraphQL subscription.
"""

import asyncio
import base64
import json

from gql import Client, gql
from gql.transport.websockets import WebsocketsTransport

SUBSCRIPTION = gql("""
subscription PlaygroundOutputSubscription($input: ChatCompletionInput!) {
  chatCompletion(input: $input) {
    __typename
    repetitionNumber
    ... on TextChunk {
      content
    }
    ... on ToolCallChunk {
      id
      function {
        name
        arguments
      }
    }
    ... on ChatCompletionSubscriptionResult {
      span {
        id
      }
    }
    ... on ChatCompletionSubscriptionError {
      message
    }
    ... on EvaluationChunk {
      experimentRunEvaluation {
        name
        label
        score
        metadata
      }
    }
  }
}
""")


async def main() -> None:
    transport = WebsocketsTransport(
        url="ws://localhost:6006/graphql",
        subprotocols=["graphql-transport-ws"],
    )

    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": "Warriors",
                    "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": "who won the nba finals in 2015?"},
                "format": "MUSTACHE",
            },
            "evaluators": evaluators,
        }
    }

    async with Client(transport=transport, fetch_schema_from_transport=False) as session:
        async for result in session.subscribe(SUBSCRIPTION, variable_values=variables):
            payload = result["chatCompletion"]
            print(json.dumps(payload))


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