from __future__ import annotations

import json
from pathlib import Path
import tempfile
import unittest

import kvcache


def metadata() -> dict[str, object]:
    return {
        "model_revision": "tiny-lab-model@0123456789abcdef",
        "tokenizer_sha256": "7f" * 32,
        "rope_theta": 10_000.0,
        "layers": 4,
        "kv_heads": 2,
        "head_dim": 8,
        "dtype": "f16",
        "layout": "layer-major-k-then-v",
        "sequence_start": 0,
        "sequence_length": 37,
    }


def tensor_metadata() -> dict[str, object]:
    value = metadata()
    value.update({
        "layers": 1,
        "kv_heads": 1,
        "head_dim": 2,
        "sequence_length": 2,
        "tensor_manifest": {
            "version": 1,
            "byte_order": "little",
            "batch_size": 1,
            "cache_class": "dynamic",
            "attention_type": "causal",
            "tensors": [
                {
                    "name": "layer.0.key",
                    "kind": "key",
                    "layer": 0,
                    "shape": [1, 1, 2, 2],
                    "strides": [8, 8, 4, 2],
                    "offset": 0,
                    "nbytes": 8,
                    "dtype": "f16",
                },
                {
                    "name": "layer.0.value",
                    "kind": "value",
                    "layer": 0,
                    "shape": [1, 1, 2, 2],
                    "strides": [8, 8, 4, 2],
                    "offset": 8,
                    "nbytes": 8,
                    "dtype": "f16",
                },
            ],
        },
    })
    return value


class CacheTests(unittest.TestCase):
    def test_round_trip_tensor_manifest_with_exact_payload_size(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "real-tensors.kvc"
            payload = bytes(range(16))
            kvcache.write_cache(path, tensor_metadata(), payload)
            self.assertEqual((tensor_metadata(), payload), kvcache.read_cache(path, tensor_metadata()))
            receipt = kvcache.inspect_cache(path)
            self.assertTrue(receipt["tensor_manifest_verified"])
            self.assertEqual(2, receipt["tensor_count"])

    def test_rejects_tensor_manifest_payload_size_mismatch(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "wrong-size.kvc"
            with self.assertRaisesRegex(kvcache.CacheFormatError, "tensor manifest declares 16 payload bytes"):
                kvcache.write_cache(path, tensor_metadata(), b"too short")
            self.assertFalse(path.exists())

    def test_rejects_non_contiguous_tensor_offsets(self) -> None:
        value = tensor_metadata()
        value["tensor_manifest"]["tensors"][1]["offset"] = 9  # type: ignore[index]
        with tempfile.TemporaryDirectory() as root:
            with self.assertRaisesRegex(kvcache.CacheFormatError, "contiguous"):
                kvcache.write_cache(Path(root) / "gap.kvc", value, bytes(17))

    def test_rejects_tensor_shape_that_disagrees_with_cache_identity(self) -> None:
        value = tensor_metadata()
        value["tensor_manifest"]["tensors"][0]["shape"] = [1, 1, 1, 4]  # type: ignore[index]
        with tempfile.TemporaryDirectory() as root:
            with self.assertRaisesRegex(kvcache.CacheFormatError, "shape must equal"):
                kvcache.write_cache(Path(root) / "wrong-shape.kvc", value, bytes(16))

    def test_rejects_non_contiguous_tensor_strides(self) -> None:
        value = tensor_metadata()
        value["tensor_manifest"]["tensors"][0]["strides"] = [1, 1, 1, 1]  # type: ignore[index]
        with tempfile.TemporaryDirectory() as root:
            with self.assertRaisesRegex(kvcache.CacheFormatError, "C-contiguous byte strides"):
                kvcache.write_cache(Path(root) / "wrong-strides.kvc", value, bytes(16))

    def test_legacy_payload_remains_inspectable_but_is_not_tensor_verified(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "legacy.kvc"
            kvcache.write_cache(path, metadata(), b"legacy")
            receipt = kvcache.inspect_cache(path)
            self.assertNotIn("tensor_manifest_verified", receipt)
            self.assertNotIn("tensor_count", receipt)

    def test_round_trip_raw_tensor_bytes(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            payload = bytes(range(255))
            kvcache.write_cache(path, metadata(), payload)
            self.assertEqual((metadata(), payload), kvcache.read_cache(path, metadata()))

    def test_inspection_receipt(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            kvcache.write_cache(path, metadata(), b"kv")
            receipt = kvcache.inspect_cache(path)
            self.assertEqual("KVC1", receipt["format"])
            self.assertEqual(2, receipt["payload_bytes"])
            json.dumps(receipt)

    def test_rejects_wrong_model(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            kvcache.write_cache(path, metadata(), b"kv")
            wrong = metadata()
            wrong["model_revision"] = "other@revision"
            with self.assertRaises(kvcache.CacheCompatibilityError):
                kvcache.read_cache(path, wrong)

    def test_rejects_wrong_position(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            kvcache.write_cache(path, metadata(), b"kv")
            wrong = metadata()
            wrong["sequence_start"] = 10
            with self.assertRaises(kvcache.CacheCompatibilityError):
                kvcache.read_cache(path, wrong)

    def test_expected_identity_keyword_accepts_a_valid_subset(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            kvcache.write_cache(path, metadata(), b"kv")
            self.assertEqual(
                (metadata(), b"kv"),
                kvcache.read_cache(
                    path,
                    expected_identity={"model_revision": metadata()["model_revision"]},
                ),
            )
            with self.assertRaises(kvcache.CacheCompatibilityError):
                kvcache.read_cache(path, expected_identity={"dtype": "bf16"})

    def test_rejects_corruption(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            kvcache.write_cache(path, metadata(), b"kv")
            raw = bytearray(path.read_bytes())
            raw[-1] ^= 1
            path.write_bytes(raw)
            with self.assertRaises(kvcache.CacheFormatError):
                kvcache.read_cache(path)

    def test_rejects_truncation(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            kvcache.write_cache(path, metadata(), b"kv")
            path.write_bytes(path.read_bytes()[:-1])
            with self.assertRaises(kvcache.CacheFormatError):
                kvcache.read_cache(path)

    def test_rejects_incomplete_identity(self) -> None:
        value = metadata()
        value.pop("tokenizer_sha256")
        with tempfile.TemporaryDirectory() as root:
            with self.assertRaises(kvcache.CacheFormatError):
                kvcache.write_cache(Path(root) / "idea.kvc", value, b"kv")

    def test_atomic_replacement(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            kvcache.write_cache(path, metadata(), b"first")
            kvcache.write_cache(path, metadata(), b"second")
            self.assertEqual(b"second", kvcache.read_cache(path)[1])
            self.assertEqual([], list(Path(root).glob("*.tmp")))


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