from __future__ import annotations

import json
from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch

from hessian_recipes import distillation
from hessian_recipes.distillation import exact_reward, final_answer, generate, generation_plan, prepare, verified


class VerifierTests(unittest.TestCase):
    def test_exact_or_unique_final_answer(self):
        self.assertTrue(verified(" 13\n", "13"))
        self.assertTrue(verified("8 + 5 = 13. <answer>13</answer>\n", "13"))
        self.assertTrue(verified([{"role": "assistant", "content": "<answer>13</answer>"}], "13"))

    def test_reward_hacking_and_ambiguous_answers_fail(self):
        invalid = ["The answer is 13", "13 or 14", "<answer>12</answer><answer>13</answer>",
                   "<answer>13</answer> actually 12", "<answer><b>13</b></answer>",
                   "<answer></answer>", "<answer>13", "", "013", "13.0",
                   [{"role": "tool", "content": "13"}],
                   [{"role": "assistant", "content": "13"}, {"role": "assistant", "content": "14"}]]
        for completion in invalid:
            with self.subTest(completion=completion):
                self.assertFalse(verified(completion, "13"))
        self.assertIsNone(final_answer(None))

    def test_rewards_preserve_batch_alignment(self):
        self.assertEqual(exact_reward(["1", "<answer>3</answer>", "2"], ["1", "2", "2"]), [1.0, 0.0, 1.0])
        with self.assertRaises(ValueError):
            exact_reward(["1"], [])


class DistillationTests(unittest.TestCase):
    def setUp(self):
        self.temporary = tempfile.TemporaryDirectory()
        self.addCleanup(self.temporary.cleanup)
        self.root = Path(self.temporary.name)
        self.source, self.output = self.root / "teacher.jsonl", self.root / "student.jsonl"

    def write(self, rows):
        self.source.write_text("\n".join(json.dumps(row) for row in rows) + "\n")

    def test_filters_deduplicates_and_records_provenance(self):
        row = {"prompt": "8+5?", "solution": "13", "completion": "<answer>13</answer>"}
        self.write([row, row, row | {"completion": "12"}, row | {"prompt": "empty?", "completion": ""}])
        result = prepare(self.source, self.output)
        self.assertEqual(result["accepted"], 1)
        self.assertEqual(result["rejected"], {"duplicate_prompt": 1, "incorrect_or_malformed_final_answer": 2})
        saved = json.loads(self.output.read_text())
        self.assertEqual(saved["provenance"]["source_sha256"], result["input_sha256"])
        self.assertTrue(self.output.with_suffix(".jsonl.manifest.json").exists())
        with self.assertRaisesRegex(ValueError, "already exists"):
            prepare(self.source, self.output)

    def test_conflicting_references_write_no_partial_output(self):
        row = {"prompt": "8+5?", "solution": "13", "completion": "13"}
        self.write([row, row | {"solution": "14"}])
        with self.assertRaisesRegex(ValueError, "Conflicting references"):
            prepare(self.source, self.output)
        self.assertFalse(self.output.exists())

    def test_no_correct_sequences_write_no_output(self):
        self.write([{"prompt": "8+5?", "solution": "13", "completion": "12"}])
        with self.assertRaisesRegex(ValueError, "No teacher sequences"):
            prepare(self.source, self.output)
        self.assertFalse(self.output.exists())

    def test_generation_plan_is_offline_and_pinned(self):
        self.write([{"prompt": "8+5?", "solution": "13"}])
        path = self.root / "generate.json"
        path.write_text(json.dumps({"schema_version": 1, "model": {"name": "a/b", "revision": "b" * 40},
                                    "dataset": "teacher.jsonl", "output": "raw.jsonl"}))
        result = generation_plan(path)
        self.assertFalse(result["model"]["allow_download"])
        self.assertEqual(result["device"], "cpu")
        self.assertFalse(result["do_sample"])
        self.assertFalse((self.root / "raw.jsonl").exists())

    def test_generation_rejects_data_changed_after_plan(self):
        self.write([{"prompt": "8+5?", "solution": "13"}])
        path = self.root / "generate.json"
        path.write_text(json.dumps({"schema_version": 1, "model": {"name": "a/b", "revision": "b" * 40},
                                    "dataset": "teacher.jsonl", "output": "raw.jsonl"}))
        baseline = generation_plan(path)
        self.write([{"prompt": "changed?", "solution": "14"}])
        with patch.object(distillation, "generation_plan", return_value=baseline):
            with self.assertRaisesRegex(ValueError, "Dataset changed"):
                generate(path)
        self.assertFalse((self.root / "raw.jsonl").exists())

    def test_generation_records_and_validates_thinking_mode(self):
        self.write([{"prompt": "8+5?", "solution": "13"}])
        path = self.root / "generate.json"
        config = {"schema_version": 1, "model": {"name": "a/b", "revision": "b" * 40},
                  "dataset": "teacher.jsonl", "output": "raw.jsonl",
                  "chat_template_kwargs": {"enable_thinking": False}}
        path.write_text(json.dumps(config))
        self.assertEqual(generation_plan(path)["chat_template_kwargs"], {"enable_thinking": False})
        config["chat_template_kwargs"] = {"enable_thinking": 1}
        path.write_text(json.dumps(config))
        with self.assertRaisesRegex(ValueError, "boolean"):
            generation_plan(path)


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