import json
import unittest
from unittest.mock import MagicMock, Mock

from openai_resptools import ToolRegistry, ToolRunner
from openai_resptools.errors import RunnerMaxIterationsError, ToolExecutionError


class TestToolRunner(unittest.TestCase):
    def setUp(self):
        self.registry = ToolRegistry()
        self.client = MagicMock()
        self.client.responses.create = MagicMock()

        @self.registry.tool()
        def echo(msg: str) -> str:
            return f"echo: {msg}"

        @self.registry.tool()
        def fail():
            raise ValueError("Something went wrong")

    def test_run_basic_flow(self):
        # 1. First response: requests tool call
        # 2. Second response: final text

        call_item = {
            "type": "function_call",
            "name": "echo",
            "call_id": "call_123",
            "arguments": json.dumps({"msg": "hello"}),
        }

        resp1 = Mock()
        resp1.output = [call_item]
        resp1.output_text = ""

        resp2 = Mock()
        resp2.output = []
        resp2.output_text = "Final answer"

        self.client.responses.create.side_effect = [resp1, resp2]

        runner = ToolRunner(self.client, model="gpt-test", registry=self.registry)
        result = runner.run("Please echo hello")

        self.assertEqual(result.text, "Final answer")
        self.assertEqual(self.client.responses.create.call_count, 2)

        # Verify first call arguments (initial prompt)
        call1_args = self.client.responses.create.call_args_list[0][1]
        self.assertEqual(call1_args["input"][0]["content"], "Please echo hello")

        # Verify second call arguments (includes tool output)
        call2_args = self.client.responses.create.call_args_list[1][1]
        input_items = call2_args["input"]

        # Input should contain:
        # 0. User prompt
        # 1. Model tool call
        # 2. Tool output
        self.assertEqual(len(input_items), 3)
        self.assertEqual(input_items[1]["call_id"], "call_123")
        self.assertEqual(input_items[2]["type"], "function_call_output")
        self.assertEqual(input_items[2]["call_id"], "call_123")
        self.assertIn("echo: hello", input_items[2]["output"])

    def test_max_iterations(self):
        # Model keeps calling tools forever
        call_item = {
            "type": "function_call",
            "name": "echo",
            "call_id": "loop",
            "arguments": json.dumps({"msg": "loop"}),
        }

        resp = Mock()
        resp.output = [call_item]

        self.client.responses.create.return_value = resp

        runner = ToolRunner(
            self.client, model="gpt-test", registry=self.registry, max_iters=3
        )

        with self.assertRaises(RunnerMaxIterationsError):
            runner.run("start loop")

        self.assertEqual(self.client.responses.create.call_count, 3)

    def test_tool_execution_error_retry(self):
        # Tool fails, but we allow retries (max_retry > 0)
        # 1. Model calls 'fail' tool
        # 2. Tool returns error string
        # 3. Model sees error and calls 'echo' instead (simulated)
        # 4. Final text

        fail_item = {
            "type": "function_call",
            "name": "fail",
            "call_id": "fail_1",
            "arguments": "{}",
        }

        echo_item = {
            "type": "function_call",
            "name": "echo",
            "call_id": "echo_1",
            "arguments": json.dumps({"msg": "fixed"}),
        }

        resp1 = Mock(output=[fail_item])
        resp2 = Mock(output=[echo_item])
        resp3 = Mock(output=[], output_text="Done")

        self.client.responses.create.side_effect = [resp1, resp2, resp3]

        # max_retry=1 means we catch the first error and send it back as output
        runner = ToolRunner(
            self.client, model="gpt-test", registry=self.registry, max_retry=1
        )
        result = runner.run("try fail")

        self.assertEqual(result.text, "Done")

        # Verify the error output was sent
        call2_args = self.client.responses.create.call_args_list[1][1]
        input_items = call2_args["input"]
        error_output = input_items[-1]  # The last item should be the tool output

        self.assertEqual(error_output["type"], "function_call_output")
        self.assertIn("Error: ", error_output["output"])
        self.assertIn("Something went wrong", error_output["output"])

    def test_tool_execution_error_raise(self):
        # Tool fails, max_retry=0, should raise exception immediately
        fail_item = {
            "type": "function_call",
            "name": "fail",
            "call_id": "fail_1",
            "arguments": "{}",
        }

        resp1 = Mock(output=[fail_item])
        self.client.responses.create.return_value = resp1

        runner = ToolRunner(
            self.client, model="gpt-test", registry=self.registry, max_retry=0
        )

        with self.assertRaises(ToolExecutionError):
            runner.run("try fail")

    def test_input_normalization(self):
        # Test with list of dicts
        runner = ToolRunner(self.client, model="gpt-test", registry=self.registry)

        # Mock empty response
        self.client.responses.create.return_value = Mock(output=[], output_text="ok")

        input_list = [
            {"role": "system", "content": "sys"},
            {"role": "user", "content": "hi"},
        ]
        runner.run(input_list)

        call_args = self.client.responses.create.call_args[1]
        self.assertEqual(call_args["input"], input_list)


if __name__ == "__main__":
    unittest.main()
