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

from hessian_recipes.trajectories import read_trajectories, validate_trajectory, write_trajectories


EXAMPLE = Path(__file__).resolve().parents[1] / "examples" / "agent_harness.py"
spec = importlib.util.spec_from_file_location("agent_harness_example", EXAMPLE)
example = importlib.util.module_from_spec(spec)
spec.loader.exec_module(example)


class TrajectoryTests(unittest.TestCase):
    def setUp(self):
        self.episode = example.build_episode()

    def test_external_tool_observation_is_context_not_a_training_response(self):
        row = validate_trajectory(self.episode)
        self.assertEqual(len(row["calls"]), 2)
        self.assertEqual(row["calls"][1]["context"][-1], {"role": "tool", "content": "42"})
        self.assertTrue(all(call["response_role"] == "assistant" for call in row["calls"]))
        # Each call retains its exact prompt, including the newly inserted tool observation.
        self.assertNotEqual(row["calls"][0]["prompt_token_ids"], row["calls"][1]["prompt_token_ids"])
        row["calls"][0]["response_loss_mask"][0] = 0
        self.assertEqual(self.episode["calls"][0]["response_loss_mask"][0], 1)

    def test_rejects_tool_output_mislabelled_as_trainable_response(self):
        self.episode["calls"][1]["response_role"] = "tool"
        with self.assertRaisesRegex(ValueError, "only assistant"):
            validate_trajectory(self.episode)

    def test_rejects_token_and_logprob_alignment_errors(self):
        mutations = (
            ("response_loss_mask", [1]),
            ("response_loss_mask", [True] * len(self.episode["calls"][0]["response_token_ids"])),
            ("response_token_ids", [True]),
            ("behavior_logprobs", [float("nan")] * len(self.episode["calls"][0]["response_token_ids"])),
            ("behavior_logprobs", [0.1] * len(self.episode["calls"][0]["response_token_ids"])),
        )
        for key, value in mutations:
            with self.subTest(key=key, value=value):
                row = copy.deepcopy(self.episode)
                row["calls"][0][key] = value
                with self.assertRaises(ValueError):
                    validate_trajectory(row)

    def test_rejects_nonfinite_reward_mutable_revision_and_unknown_fields(self):
        rows = [copy.deepcopy(self.episode) for _ in range(3)]
        rows[0]["reward"]["value"] = float("inf")
        rows[1]["policy"]["revision"] = "main"
        rows[2]["unrecognized"] = 1
        for row in rows:
            with self.assertRaises(ValueError):
                validate_trajectory(row)

    def test_jsonl_roundtrip_and_atomic_rejection(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "trace.jsonl"
            self.assertEqual(write_trajectories(path, [self.episode]), 1)
            self.assertEqual(read_trajectories(path), [self.episode])
            before = path.read_bytes()
            invalid = copy.deepcopy(self.episode)
            invalid["reward"]["value"] = float("nan")
            with self.assertRaises(ValueError):
                write_trajectories(path, [invalid])
            self.assertEqual(path.read_bytes(), before)

    def test_rejects_duplicate_episodes_and_mixed_policy_groups(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "trace.jsonl"
            with self.assertRaisesRegex(ValueError, "duplicate episode_id"):
                write_trajectories(path, [self.episode, self.episode])
            other = copy.deepcopy(self.episode)
            other["episode_id"] = "another-episode"
            other["policy"]["revision"] = "a-different-policy"
            with self.assertRaisesRegex(ValueError, "mismatch"):
                write_trajectories(path, [self.episode, other])

    def test_reader_identifies_bad_line_and_rejects_duplicate_json_keys(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "trace.jsonl"
            valid = json.dumps(self.episode)
            path.write_text(valid + "\n{broken\n", encoding="utf-8")
            with self.assertRaisesRegex(ValueError, "trace.jsonl:2:"):
                read_trajectories(path)
            path.write_text('{"schema_version":1,"schema_version":2}\n', encoding="utf-8")
            with self.assertRaisesRegex(ValueError, "duplicate JSON key"):
                read_trajectories(path)


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