from __future__ import annotations

import json
from pathlib import Path
import tempfile
from types import ModuleType, SimpleNamespace
import unittest
from unittest.mock import MagicMock, patch

from hessian_recipes import checkpoints


class CheckpointMergeTests(unittest.TestCase):
    def setUp(self):
        self.temporary = tempfile.TemporaryDirectory()
        self.addCleanup(self.temporary.cleanup)
        root = Path(self.temporary.name)
        self.base, self.adapter, self.output = root / "base", root / "adapter", root / "merged"
        self.base.mkdir()
        self.adapter.mkdir()
        (self.base / "config.json").write_text('{"model_type":"test"}')
        (self.base / "tokenizer_config.json").write_text('{}')
        (self.base / "model.safetensors").write_bytes(b"base weights")
        (self.adapter / "adapter_config.json").write_text(json.dumps({
            "peft_type": "LORA", "base_model_name_or_path": "unsafe/remote-pointer"}))
        (self.adapter / "adapter_model.safetensors").write_bytes(b"adapter weights")
        self.model, self.adapted, self.merged, self.tokenizer = (MagicMock() for _ in range(4))

        def save_model(path, **kwargs):
            destination = Path(path)
            destination.mkdir(parents=True)
            (destination / "config.json").write_text('{}')
            (destination / "model.safetensors").write_bytes(b"merged weights")

        def save_tokenizer(path):
            (Path(path) / "tokenizer_config.json").write_text('{}')
            (Path(path) / "tokenizer.json").write_text('{"test":true}')

        self.merged.save_pretrained.side_effect = save_model
        self.tokenizer.save_pretrained.side_effect = save_tokenizer
        self.adapted.merge_and_unload.return_value = self.merged
        self.torch = ModuleType("torch")
        self.torch.bfloat16, self.torch.float32, self.torch.float16 = "bf16", "fp32", "fp16"
        self.torch.cuda = SimpleNamespace(is_available=lambda: False)
        self.transformers = ModuleType("transformers")
        self.transformers.AutoModelForCausalLM = SimpleNamespace(from_pretrained=MagicMock(return_value=self.model))
        self.transformers.AutoTokenizer = SimpleNamespace(from_pretrained=MagicMock(return_value=self.tokenizer))
        self.peft = ModuleType("peft")
        self.peft.PeftModel = SimpleNamespace(from_pretrained=MagicMock(return_value=self.adapted))

    def merge(self, **kwargs):
        with patch.object(checkpoints, "require_versions", return_value={}), patch.dict(
            "sys.modules", {"torch": self.torch, "transformers": self.transformers, "peft": self.peft}
        ):
            return checkpoints.merge_adapter(self.base, self.adapter, self.output, **kwargs)

    def test_explicit_base_is_used_and_provenance_matches_saved_files(self):
        report = self.merge()
        self.assertEqual(report["status"], "completed")
        self.assertEqual(report["output"]["sha256"], checkpoints.checkpoint_sha256(self.output))
        self.assertFalse(report["adapter"]["base_pointer_used"])
        self.assertEqual(report["adapter"]["declared_base_model_name_or_path"], "unsafe/remote-pointer")
        self.peft.PeftModel.from_pretrained.assert_called_once_with(
            self.model, str(self.adapter), local_files_only=True, is_trainable=False)
        self.transformers.AutoModelForCausalLM.from_pretrained.assert_called_once_with(
            str(self.base), dtype="bf16", local_files_only=True, trust_remote_code=False, use_safetensors=True)
        self.adapted.merge_and_unload.assert_called_once_with(safe_merge=True)
        self.assertTrue(Path(report["manifest"]).is_file())
        self.assertIn("tokenizer.json", report["tokenizer_file_sha256"])
        self.assertFalse((self.output / "adapter_config.json").exists())

    def test_mutation_during_loading_leaves_no_output(self):
        def changed(*args, **kwargs):
            (self.base / "model.safetensors").write_bytes(b"changed")
            return self.model
        self.transformers.AutoModelForCausalLM.from_pretrained.side_effect = changed
        with self.assertRaisesRegex(ValueError, "changed while loading"):
            self.merge()
        self.assertFalse(self.output.exists())

    def test_failed_merge_leaves_no_output(self):
        self.adapted.merge_and_unload.side_effect = ValueError("non-finite adapter")
        with self.assertRaisesRegex(ValueError, "non-finite"):
            self.merge()
        self.assertFalse(self.output.exists())

    def test_existing_output_and_nested_output_are_rejected(self):
        self.output.mkdir()
        with self.assertRaisesRegex(ValueError, "already exists"):
            self.merge()
        with self.assertRaisesRegex(ValueError, "separate"):
            checkpoints.merge_adapter(self.base, self.adapter, self.base / "new")

    def test_unsupported_and_remote_inputs_fail_before_imports(self):
        for arguments in ({"dtype": "auto"}, {"device": "http://remote"}):
            with self.subTest(arguments=arguments), self.assertRaises(ValueError):
                self.merge(**arguments)
        with self.assertRaisesRegex(ValueError, "existing local"):
            checkpoints.merge_adapter("remote/model", self.adapter, self.output)

    def test_cuda_request_is_never_silently_moved_to_cpu(self):
        with self.assertRaisesRegex(RuntimeError, "CUDA was requested"):
            self.merge(device="cuda")
        self.transformers.AutoModelForCausalLM.from_pretrained.assert_not_called()


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