from __future__ import annotations

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

from hessian_recipes import training


class TrainingPlanTests(unittest.TestCase):
    def setUp(self):
        self.temporary = tempfile.TemporaryDirectory()
        self.addCleanup(self.temporary.cleanup)
        self.root = Path(self.temporary.name)
        (self.root / "data.jsonl").write_text(json.dumps({"prompt": "What is 8 + 5?", "completion": "13", "solution": "13"}) + "\n")
        self.config = {"schema_version": 1, "method": "sft",
                       "model": {"name": "test/model", "revision": "a" * 40},
                       "dataset": "data.jsonl", "output_dir": "output"}
        self.path = self.root / "recipe.json"

    def plan(self):
        self.path.write_text(json.dumps(self.config))
        return training.plan(self.path)

    def test_plan_resolves_paths_without_importing_model_libraries(self):
        result = self.plan()
        self.assertEqual(result["rows"], 1)
        self.assertEqual(result["dataset"], str(self.root / "data.jsonl"))
        self.assertEqual(result["config_sha256"], training.sha256(self.path))
        self.assertFalse(result["model"]["allow_download"])
        self.assertFalse((self.root / "output").exists())
        self.assertEqual(result["report_to"], "none")

    def test_mutable_hub_revisions_rejected(self):
        for revision in (None, "main", "v1.0", "a" * 39):
            with self.subTest(revision=revision):
                self.config["model"]["revision"] = revision
                with self.assertRaisesRegex(ValueError, "immutable"):
                    self.plan()

    def test_misspelled_or_network_trainer_options_rejected(self):
        for key in ("learing_rate", "report_to", "push_to_hub", "use_vllm"):
            with self.subTest(key=key):
                self.config["training"] = {key: True}
                with self.assertRaisesRegex(ValueError, "Unknown training"):
                    self.plan()

    def test_grpo_requires_valid_generation_groups(self):
        self.config["method"] = "grpo"
        with self.assertRaisesRegex(ValueError, "generation batch"):
            self.plan()
        self.config["training"] = {"per_device_train_batch_size": 2, "num_generations": 2}
        self.assertEqual(self.plan()["method"], "grpo")
        self.config["training"]["num_generations"] = 1
        with self.assertRaisesRegex(ValueError, "at least 2"):
            self.plan()

    def test_nonfinite_and_boolean_hyperparameters_rejected(self):
        for settings in ({"max_steps": True}, {"learning_rate": float("nan")},
                         {"learning_rate": -1}, {"max_length": 0}, {"max_steps": 1.5}):
            with self.subTest(settings=settings):
                self.config["training"] = settings
                with self.assertRaises(ValueError):
                    self.plan()

    def test_duplicate_prompts_rejected(self):
        with (self.root / "data.jsonl").open("a") as stream:
            stream.write(json.dumps({"prompt": " What is 8 + 5? ", "completion": "13"}) + "\n")
        with self.assertRaisesRegex(ValueError, "Duplicate training prompt"):
            self.plan()

    def test_local_checkpoint_and_offline_kwargs(self):
        (self.root / "model").mkdir()
        self.config["model"] = {"name": "./model"}
        result = self.plan()
        self.assertEqual(result["model"]["source"], "local")
        self.assertEqual(training.model_kwargs(result["model"]),
                         {"local_files_only": True, "trust_remote_code": False})

    def test_typo_duplicate_json_key_rejected(self):
        self.path.write_text('{"method":"sft","method":"grpo"}')
        with self.assertRaisesRegex(ValueError, "Duplicate JSON key"):
            training.plan(self.path)

    def test_duplicate_dataset_keys_rejected(self):
        (self.root / "data.jsonl").write_text('{"prompt":"p","completion":"13","completion":"14"}\n')
        with self.assertRaisesRegex(ValueError, "Duplicate JSON key"):
            self.plan()

    def test_changed_dataset_is_rejected_before_runtime_imports(self):
        baseline = self.plan()
        (self.root / "data.jsonl").write_text('{"prompt":"changed","completion":"14"}\n')
        with patch.object(training, "plan", return_value=baseline):
            with self.assertRaisesRegex(ValueError, "Dataset changed"):
                training.run(self.path)
        self.assertFalse((self.root / "output").exists())

    def test_changed_configuration_is_rejected_before_runtime_imports(self):
        baseline = self.plan()
        self.path.write_text('{}')
        with patch.object(training, "plan", return_value=baseline):
            with self.assertRaisesRegex(ValueError, "Configuration changed"):
                training.run(self.path)
        self.assertFalse((self.root / "output").exists())

    def test_lora_validation(self):
        self.config["lora"] = {"r": 0}
        with self.assertRaisesRegex(ValueError, "lora.r"):
            self.plan()
        self.config["lora"] = {"r": 8, "target_modules": "all-linear"}
        self.assertEqual(self.plan()["lora"]["r"], 8)

    def test_thinking_mode_is_explicit_and_strict(self):
        self.config["chat_template_kwargs"] = {"enable_thinking": True}
        self.assertEqual(self.plan()["chat_template_kwargs"], {"enable_thinking": True})
        for invalid in (None, {"enable_thinking": "false"}, {"enable_thinking": 0}, {"tools": []}):
            self.config["chat_template_kwargs"] = invalid
            with self.subTest(value=invalid), self.assertRaises(ValueError):
                self.plan()

    def test_non_thinking_targets_cannot_include_reasoning(self):
        self.config["chat_template_kwargs"] = {"enable_thinking": False}
        for target in ("<think>Calculate 8+5.</think>13", "<think>unclosed", "13</think>"):
            (self.root / "data.jsonl").write_text(json.dumps({"prompt": "8+5?", "completion": target}) + "\n")
            with self.subTest(target=target), self.assertRaisesRegex(ValueError, "Non-thinking"):
                self.plan()
        (self.root / "data.jsonl").write_text(json.dumps({"prompt": "8+5?", "completion": "<think>\n\n</think>13"}) + "\n")
        self.assertEqual(self.plan()["rows"], 1)


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