from __future__ import annotations

import asyncio
import json
import os
import unittest

from erniebot_agent.tools.calculator_tool import CalculatorTool

import erniebot

erniebot.api_type = "aistudio"
erniebot.access_token = os.environ["AISTUDIO_ACCESS_TOKEN"]


class TestCalculator(unittest.TestCase):
    def setUp(self) -> None:
        self.tool = CalculatorTool()

    def run_query(self, query):
        response = erniebot.ChatCompletion.create(
            model="ernie-bot",
            messages=[
                {
                    "role": "user",
                    "content": query,
                }
            ],
            functions=[self.tool.function_call_schema()],
            stream=False,
        )
        result = response.get_result()
        if isinstance(result, str):
            return result

        assert result["name"] == "CalculatorTool"
        arguments = json.loads(result["arguments"])
        result = asyncio.run(self.tool(**arguments))
        return result

    def test_add(self):
        result = self.run_query("3 加四等于多少")
        self.assertEqual(result["formula_result"], 7)

    def test_complex_formula(self):
        result = self.run_query("3乘以五 再加10等于多少")
        self.assertEqual(result["formula_result"], 25)