import importlib.util
import json
import unittest
from pathlib import Path


SCRIPT = Path(__file__).resolve().parents[1] / "eval" / "qwen_native_closed_loop" / "run_closed_loop.py"
SPEC = importlib.util.spec_from_file_location("run_closed_loop", SCRIPT)
MODULE = importlib.util.module_from_spec(SPEC)
assert SPEC.loader is not None
SPEC.loader.exec_module(MODULE)


class ScriptedChat:
    def __init__(self):
        self.turn = 0

    def __call__(self, messages, tools):
        self.turn += 1
        self.assert_tool_schema(tools)
        if self.turn == 1:
            return {
                "role": "assistant",
                "content": "",
                "tool_calls": [{
                    "id": "call-1",
                    "type": "function",
                    "function": {
                        "name": "read_release_manifest",
                        "arguments": json.dumps({
                            "release_id": "aug3",
                            "component": "qwen3-coder-next",
                        }),
                    },
                }],
            }
        if self.turn == 2:
            manifest = json.loads(messages[-1]["content"])
            return {
                "role": "assistant",
                "content": "",
                "tool_calls": [{
                    "id": "call-2",
                    "type": "function",
                    "function": {
                        "name": "run_check_plan",
                        "arguments": json.dumps({
                            "component": manifest["component"],
                            "checks": manifest["required_checks"],
                            "run_nonce": manifest["run_nonce"],
                        }),
                    },
                }],
            }
        receipt = json.loads(messages[-1]["content"])["receipt"]
        return {"role": "assistant", "content": f"LAUNCH_PREFLIGHT_PASS:{receipt}"}

    @staticmethod
    def assert_tool_schema(tools):
        assert [tool["function"]["name"] for tool in tools] == [
            "read_release_manifest",
            "run_check_plan",
        ]


class SingleToolChat:
    def __init__(self):
        self.turn = 0

    def __call__(self, messages, tools):
        self.turn += 1
        if self.turn == 1:
            return {
                "role": "assistant",
                "content": "",
                "tool_calls": [{
                    "id": "single-1",
                    "type": "function",
                    "function": {
                        "name": "read_release_manifest",
                        "arguments": json.dumps({
                            "release_id": "aug3",
                            "component": "qwen3-coder-next",
                        }),
                    },
                }],
            }
        nonce = json.loads(messages[-1]["content"])["run_nonce"]
        return {"role": "assistant", "content": f"MANIFEST_NONCE:{nonce}"}


class RecoveryChat:
    def __init__(self):
        self.turn = 0

    def __call__(self, messages, tools):
        self.turn += 1
        if self.turn == 1:
            return {
                "role": "assistant",
                "content": "",
                "tool_calls": [{
                    "id": "recovery-1",
                    "type": "function",
                    "function": {
                        "name": "read_release_manifest",
                        "arguments": json.dumps({
                            "release_id": "missing",
                            "component": "qwen3-coder-next",
                        }),
                    },
                }],
            }
        if self.turn == 2:
            error = json.loads(messages[-1]["content"])
            assert error["error"] == "manifest_not_found"
            return {
                "role": "assistant",
                "content": "",
                "tool_calls": [{
                    "id": "recovery-2",
                    "type": "function",
                    "function": {
                        "name": "read_release_manifest",
                        "arguments": json.dumps({
                            "release_id": "aug3",
                            "component": "qwen3-coder-next",
                        }),
                    },
                }],
            }
        nonce = json.loads(messages[-1]["content"])["run_nonce"]
        return {"role": "assistant", "content": f"RECOVERED:{nonce}"}


class ClosedLoopTests(unittest.TestCase):
    def test_single_tool_dynamic_result_passes(self):
        result = MODULE.run_single_tool_roundtrip(SingleToolChat(), nonce="fixed-test-nonce")
        self.assertTrue(result["passed"], result)
        self.assertEqual(result["final_answer"], "MANIFEST_NONCE:fixed-test-nonce")

    def test_two_tool_roundtrip_passes(self):
        result = MODULE.run_closed_loop(ScriptedChat(), nonce="fixed-test-nonce")
        self.assertTrue(result["passed"], result)
        self.assertEqual([event["tool"] for event in result["tool_executions"]], [
            "read_release_manifest",
            "run_check_plan",
        ])
        self.assertTrue(result["final_answer"].startswith("LAUNCH_PREFLIGHT_PASS:"))

    def test_tool_error_recovery_passes(self):
        result = MODULE.run_tool_error_recovery(RecoveryChat(), nonce="fixed-test-nonce")
        self.assertTrue(result["passed"], result)
        self.assertEqual(len(result["tool_executions"]), 2)
        self.assertEqual(result["tool_executions"][0]["result"]["error"], "manifest_not_found")
        self.assertEqual(result["final_answer"], "RECOVERED:fixed-test-nonce")

    def test_wrong_first_arguments_fail(self):
        def wrong_chat(messages, tools):
            return {
                "role": "assistant",
                "content": "",
                "tool_calls": [{
                    "id": "bad-1",
                    "type": "function",
                    "function": {
                        "name": "read_release_manifest",
                        "arguments": json.dumps({
                            "release_id": "aug4",
                            "component": "qwen3-coder-next",
                        }),
                    },
                }],
            }

        result = MODULE.run_closed_loop(wrong_chat, nonce="fixed-test-nonce")
        self.assertFalse(result["passed"])
        self.assertIn("args mismatch", result["error"])


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