import copy
import importlib.util
import json
import os
from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch


ROOT = Path(__file__).resolve().parents[1]


def load(name, path):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


harness = load("hessian_lightning_harness", ROOT / "examples/agent_lightning/hessian_lightning_harness.py")
launcher = load("hessian_lightning_launcher", ROOT / "examples/agent_lightning/launch.py")


class LightningHarnessTests(unittest.IsolatedAsyncioTestCase):
    async def test_real_tool_loop_preserves_tool_observation(self):
        contexts = []

        async def policy(messages):
            contexts.append(copy.deepcopy(messages))
            if len(contexts) == 1:
                return {"role": "assistant", "content": None, "tool_calls": [{
                    "id": "tool-1", "type": "function", "function": {
                        "name": "multiply", "arguments": json.dumps({"a": 6, "b": 7})}}]}
            return {"role": "assistant", "content": "<answer>42</answer>"}

        self.assertEqual(await harness.external_harness("6 times 7?", policy), "<answer>42</answer>")
        self.assertEqual(contexts[1][-1], {"role": "tool", "tool_call_id": "tool-1", "content": "42"})

    async def test_malformed_arguments_receive_feedback_and_loop_is_bounded(self):
        count = 0

        async def policy(messages):
            nonlocal count
            count += 1
            if count > 1:
                self.assertIn("Tool error", messages[-1]["content"])
            return {"role": "assistant", "content": None, "tool_calls": [{
                "id": f"tool-{count}", "function": {"name": "multiply", "arguments": "[]"}}]}

        self.assertEqual(await harness.external_harness("multiply", policy), "")
        self.assertEqual(count, 3)

    async def test_agent_uses_upstream_environment_and_reports_scalar_reward(self):
        requests = []

        def post(url, key, payload):
            requests.append((url, payload))
            return {"choices": [{"message": {"role": "assistant", "content": "<answer>42</answer>"}}]}

        with patch.dict(os.environ, {"AGL_KEY": "local-test", "AGL_OPENAI_BASE_URL": "http://localhost/proxy/v1",
                                     "AGL_EVENT_URL": "http://localhost/events", "QUESTION": "6*7", "RESULT": "42"}), patch.object(harness, "_post", post):
            await harness.Agent().run()
        self.assertEqual(requests[-1], ("http://localhost/events", {"event_type": "reward", "data": {"value": 1.0}}))


class LightningSourceTests(unittest.TestCase):
    def test_source_check_requires_exact_revision_and_expected_signatures(self):
        with tempfile.TemporaryDirectory() as directory:
            upstream = Path(directory)
            for path, code in {
                "examples/calc_x/train_calc_agent.py": "def build_config(*, model, agl_base_url, agl_key, config_overrides): pass",
                "agentlightning/verl/entrypoint.py": "def run_ppo(config, train_dataset, val_dataset): pass",
            }.items():
                file = upstream / path
                file.parent.mkdir(parents=True, exist_ok=True)
                file.write_text(code)
            with patch.object(launcher.subprocess, "check_output", side_effect=[launcher.REVISION + "\n", ""]):
                self.assertEqual(launcher.check_source(upstream)["revision"], launcher.REVISION)
            with patch.object(launcher.subprocess, "check_output", return_value="wrong-revision"):
                with self.assertRaisesRegex(ValueError, "Expected upstream"):
                    launcher.check_source(upstream)