"""
Tests for the product gallery's matching logic.

The embedding model is injected, so these run without torch or a GPU: what is
being tested is the retrieval contract - best view wins, several views per SKU
count as one answer, and a stranger is rejected rather than named.
"""

from __future__ import annotations

import os
import sys
import unittest

import numpy as np

current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
if os.path.exists(os.path.join(project_root, "stock")):
    os.chdir(project_root)
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

from gallery import (  # noqa: E402
    DINOV2_MODEL,
    DINOV2_SHA256,
    GalleryEntry,
    ProductGallery,
    catalog_digest,
    dinov2_weights_path,
    head_class_set,
    load_dinov2,
)

import gallery as gallery_module  # noqa: E402


def a_gallery(views, threshold=0.55):
    """A gallery with hand-made embeddings: no model, no files."""
    skus = sorted({sku for sku, _ in views})
    gallery = ProductGallery("does-not-exist", skus, threshold=threshold)
    embeddings = []
    entries = []
    for index, (sku, vector) in enumerate(views):
        vector = np.asarray(vector, dtype=np.float32)
        vector /= np.linalg.norm(vector)
        embeddings.append(vector)
        entries.append(GalleryEntry(sku_id=sku, view=f"{sku}_{index}", embedding=vector))
    gallery._entries = entries
    gallery._matrix = np.stack(embeddings)
    gallery._available = True
    return gallery


def unit(vector):
    vector = np.asarray(vector, dtype=np.float32)
    return (vector / np.linalg.norm(vector)).tolist()


class Backbone:
    """A deterministic stand-in for DINOv2: the prepared crop's channel means,
    tiled to the width the real backbone returns."""

    def __init__(self):
        self.calls = 0

    def __call__(self, tensor):
        self.calls += 1
        return tensor.mean(dim=(2, 3)).repeat(1, gallery_module.DINOV2_EMBED_DIM // 3)


def a_query():
    """A unit vector in the backbone's space, for match() without a model."""
    query = np.zeros(gallery_module.DINOV2_EMBED_DIM, dtype=np.float32)
    query[0] = 1.0
    query[1] = 0.05
    return query / np.linalg.norm(query)


class UnreachableBackbone:
    def __call__(self, tensor):
        raise AssertionError("the gallery loaded a backbone it should have cached")


def write_view(path, seed):
    import cv2

    image = np.random.RandomState(seed).randint(0, 256, (24, 32, 3), dtype=np.uint8)
    cv2.imwrite(path, image)


def a_store(directory, names=("sku_a", "sku_a_1", "sku_b")):
    """A store dir: one JPEG per view plus the stock.csv the catalog is read from."""
    with open(os.path.join(directory, "stock.csv"), "w", encoding="utf-8") as handle:
        handle.write("SKU_ID,product_name\n")
    for index, name in enumerate(names):
        write_view(os.path.join(directory, f"{name}.jpeg"), index)


def a_stored_gallery(directory, names, backbone, image_size=32):
    gallery = ProductGallery(directory, list(names), device="cpu", image_size=image_size)
    gallery._load_model = lambda: backbone
    return gallery


class MatchingTests(unittest.TestCase):
    def test_the_nearest_view_names_the_product(self):
        gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0, 1, 0]))])
        gallery.embed = lambda crop: np.array([0.9, 0.1, 0.0], dtype=np.float32)
        match = gallery.match(object())
        self.assertFalse(match.rejected)
        self.assertEqual(match.sku_id, "sku_a")
        self.assertAlmostEqual(match.similarity, 0.9, places=3)

    def test_several_views_of_one_sku_are_one_answer(self):
        # The second view of sku_b is closer than any view of sku_a, but sku_a must
        # still win on its best view: a SKU is not outvoted by having more photos.
        gallery = a_gallery([
            ("sku_a", unit([1, 0, 0])),
            ("sku_b", unit([0, 1, 0])),
            ("sku_b", unit([0.2, 0.98, 0])),
        ])
        gallery.embed = lambda crop: np.array([1.0, 0.05, 0.0], dtype=np.float32)
        match = gallery.match(object())
        self.assertEqual(match.sku_id, "sku_a")
        self.assertEqual(match.runner_up, "sku_b")

    def test_a_stranger_is_rejected_not_named(self):
        gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0, 1, 0]))])
        gallery.embed = lambda crop: np.array([0.0, 0.0, 1.0], dtype=np.float32)
        match = gallery.match(object())
        self.assertTrue(match.rejected)
        self.assertIsNone(match.sku_id)

    def test_the_threshold_moves_the_rejection_line(self):
        gallery = a_gallery([("sku_a", unit([1, 0, 0]))], threshold=0.99)
        gallery.embed = lambda crop: np.array([0.9, 0.4, 0.0], dtype=np.float32)
        self.assertTrue(gallery.match(object()).rejected)
        self.assertFalse(gallery.match(object(), threshold=0.5).rejected)

    def test_margin_separates_a_confident_match_from_a_close_one(self):
        gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0.9, 0.44, 0]))])
        gallery.embed = lambda crop: np.array([1.0, 0.1, 0.0], dtype=np.float32)
        confident = gallery.match(object())
        gallery.embed = lambda crop: np.array([0.95, 0.3, 0.0], dtype=np.float32)
        close = gallery.match(object())
        self.assertGreater(confident.margin, close.margin)

    def test_the_per_sku_max_does_not_depend_on_entry_order(self):
        # Three SKUs, several views each, enrolled out of order: a SKU is scored
        # by its best view, and the answer cannot move because a row moved.
        gallery = a_gallery([
            ("sku_b", unit([0, 1, 0])),
            ("sku_a", unit([1, 0, 0])),
            ("sku_c", unit([0, 0, 1])),
            ("sku_b", unit([0.3, 0.95, 0])),
            ("sku_a", unit([0.98, 0.2, 0])),
        ])
        gallery.embed = lambda crop: np.array([1.0, 0.05, 0.0], dtype=np.float32)
        straight = gallery.match(object())
        self.assertEqual(straight.sku_id, "sku_a")
        self.assertEqual(straight.runner_up, "sku_b")

        for order in ([4, 0, 3, 1, 2], [2, 3, 1, 4, 0]):
            gallery._entries = [gallery._entries[index] for index in order]
            gallery._matrix = gallery._matrix[order]
            shuffled = gallery.match(object())
            self.assertEqual(shuffled.sku_id, straight.sku_id)
            self.assertEqual(shuffled.runner_up, straight.runner_up)
            self.assertAlmostEqual(shuffled.similarity, straight.similarity, places=6)
            self.assertAlmostEqual(shuffled.runner_up_similarity, straight.runner_up_similarity, places=6)

    def test_a_tie_keeps_the_first_enrolled_sku(self):
        # Exactly equidistant: first appearance decides, as it did before the
        # per-SKU maximum became a segmented reduction.
        gallery = a_gallery([("sku_a", unit([1, 0, 0])), ("sku_b", unit([0, 1, 0]))])
        gallery.embed = lambda crop: np.array([1.0, 1.0, 0.0], dtype=np.float32)
        match = gallery.match(object())
        self.assertFalse(match.rejected)
        self.assertEqual(match.sku_id, "sku_a")
        self.assertEqual(match.runner_up, "sku_b")
        self.assertAlmostEqual(match.similarity, match.runner_up_similarity, places=6)
        self.assertAlmostEqual(match.margin, 0.0, places=6)

    def test_an_unloaded_gallery_rejects_everything(self):
        gallery = ProductGallery("does-not-exist", ["sku_a"])
        match = gallery.match(object())
        self.assertTrue(match.rejected)


class EnrollmentTests(unittest.TestCase):
    def test_views_are_found_by_sku_prefix(self):
        import tempfile

        with tempfile.TemporaryDirectory() as directory:
            for name in ("sku_a.jpeg", "sku_a_1.jpeg", "sku_a-2.jpeg", "sku_ab.jpeg", "notes.txt"):
                open(os.path.join(directory, name), "wb").close()
            gallery = ProductGallery(directory, ["sku_a", "sku_ab"])
            self.assertEqual(len(gallery._views_for("sku_a")), 3,
                             "one view, two extra views, and not sku_ab")
            self.assertEqual(len(gallery._views_for("sku_ab")), 1)


class FeatureCacheTests(unittest.TestCase):
    """The cached matrix must be this store's, or it must not be used."""

    def setUp(self):
        import hashlib
        import tempfile

        self.directory = tempfile.TemporaryDirectory()
        self.addCleanup(self.directory.cleanup)
        a_store(self.directory.name)
        # A cache is only served when the pinned backbone is on disk, so point the
        # pin at a stand-in file the module's expected digest is told about. The
        # check itself - hash this file, compare - is what is under test.
        self.weights = os.path.join(self.directory.name, "dinov2.safetensors")
        self.write(self.weights, b"not really a checkpoint")
        os.environ["DORAYAKI_DINOV2_WEIGHTS"] = self.weights
        self.addCleanup(os.environ.pop, "DORAYAKI_DINOV2_WEIGHTS", None)
        self.addCleanup(setattr, gallery_module, "DINOV2_SHA256", gallery_module.DINOV2_SHA256)
        gallery_module.DINOV2_SHA256 = hashlib.sha256(b"not really a checkpoint").hexdigest()

    def write(self, path, content):
        with open(path, "wb") as handle:
            handle.write(content)
        return path

    def real_gallery(self):
        """A gallery with the real _load_model: the server's own readiness path."""
        return ProductGallery(self.directory.name, ["sku_a", "sku_b"], device="cpu", image_size=32)

    def loaded(self, names=("sku_a", "sku_b"), image_size=32):
        backbone = Backbone()
        gallery = a_stored_gallery(self.directory.name, names, backbone, image_size=image_size)
        self.assertTrue(gallery.load())
        return gallery, backbone

    def cache_payload(self, gallery):
        """The rows and matrix a loaded gallery left on disk."""
        with np.load(gallery._cache_path(), allow_pickle=False) as payload:
            return (
                [str(value) for value in payload["skus"].tolist()],
                [str(value) for value in payload["views"].tolist()],
                payload["matrix"],
            )

    def rewrite_cache(self, gallery, skus, views, matrix):
        """The same valid key with different contents: only the rows can catch it."""
        path = gallery._cache_path()
        with np.load(path, allow_pickle=False) as payload:
            key = str(payload["key"].item())
        np.savez(path, key=key, skus=np.array(skus), views=np.array(views), matrix=matrix)

    def re_embeds(self, first, why):
        """The one right answer to a rewritten cache: build it from the pixels."""
        backbone = Backbone()
        gallery = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), backbone)
        self.assertTrue(gallery.load(), why)
        self.assertEqual(backbone.calls, 1, f"{why} must not be served from cache")
        np.testing.assert_array_equal(gallery._matrix, first._matrix)
        self.assertEqual([entry.sku_id for entry in gallery._entries], ["sku_a", "sku_a", "sku_b"])

    def test_a_second_start_answers_the_same_from_the_cache(self):
        first, backbone = self.loaded()
        self.assertEqual(backbone.calls, 1)
        self.assertEqual(first.view_count(), 3)
        self.assertTrue(os.path.isfile(first._cache_path()))

        cached = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), UnreachableBackbone())
        self.assertTrue(cached.load())
        self.assertEqual(cached.view_count(), first.view_count())
        np.testing.assert_array_equal(cached._matrix, first._matrix)
        self.assertEqual([entry.view for entry in cached._entries],
                         [entry.view for entry in first._entries])
        query = a_query()
        first.embed = cached.embed = lambda crop: query
        self.assertEqual(vars(cached.match(object())), vars(first.match(object())))

    def test_changed_view_bytes_re_embed(self):
        first, _ = self.loaded()
        before = first._matrix.copy()
        write_view(os.path.join(self.directory.name, "sku_b.jpeg"), seed=99)

        second, backbone = self.loaded()
        self.assertEqual(backbone.calls, 1, "a replaced view must not be served from cache")
        self.assertFalse(np.array_equal(second._matrix, before))

    def test_a_reordered_catalog_does_not_borrow_another_mapping(self):
        self.loaded()
        swapped, backbone = self.loaded(names=("sku_b", "sku_a"))
        self.assertEqual(backbone.calls, 1)
        self.assertEqual([entry.sku_id for entry in swapped._entries], ["sku_b", "sku_a", "sku_a"])

    def test_another_image_size_is_never_served(self):
        small, _ = self.loaded(image_size=32)
        large, backbone = self.loaded(image_size=64)
        self.assertEqual(backbone.calls, 1)
        self.assertNotEqual(small._cache_path(), large._cache_path())

    def test_a_corrupt_cache_re_embeds_the_same_matrix(self):
        first, _ = self.loaded()
        with open(first._cache_path(), "wb") as handle:
            handle.write(b"not an npz at all")

        second, backbone = self.loaded()
        self.assertEqual(backbone.calls, 1)
        np.testing.assert_array_equal(second._matrix, first._matrix)

    def test_swapped_cache_rows_are_not_served(self):
        # A current key says nothing about the rows inside the file. The two
        # sku_a views are adjacent, so a swap has to cross SKUs to mean
        # anything - and the matrix that moves with them is the dangerous part:
        # a vector that would then be filed under the wrong product.
        first, _ = self.loaded()
        skus, views, matrix = self.cache_payload(first)
        self.assertEqual(list(zip(skus, views)),
                         [("sku_a", "sku_a.jpeg"), ("sku_a", "sku_a_1.jpeg"), ("sku_b", "sku_b.jpeg")])
        self.rewrite_cache(first, skus[::-1], views[::-1], matrix[::-1])
        self.re_embeds(first, "rows swapped with their vectors")

    def test_mislabelled_cache_rows_are_not_served(self):
        # The views and the matrix stay where they are, so every row names a
        # product it was not embedded for.
        first, _ = self.loaded()
        skus, views, matrix = self.cache_payload(first)
        self.rewrite_cache(first, [skus[2], skus[1], skus[0]], views, matrix)
        self.re_embeds(first, "labels swapped around the views")

    def test_a_duplicated_cache_row_is_not_served(self):
        # The first view written twice: one image would answer for two products.
        first, _ = self.loaded()
        skus, views, matrix = self.cache_payload(first)
        self.rewrite_cache(first, skus[:1] * 2 + skus[2:], views[:1] * 2 + views[2:], matrix)
        self.re_embeds(first, "a row repeated")

    def test_a_foreign_sku_row_is_not_served(self):
        # A row from another store's catalog: a SKU this gallery never enrolled.
        first, _ = self.loaded()
        skus, views, matrix = self.cache_payload(first)
        self.rewrite_cache(first, skus[1:] + ["other_store"], views[1:] + ["other_store.jpeg"], matrix)
        self.re_embeds(first, "another store's row")

    def test_a_cache_may_omit_a_view_it_could_not_decode(self):
        # The writer drops a view it could not read and keeps everyone else's
        # order, so such a cache is still this store's - and still worth serving
        # rather than re-embedding the whole catalog on every start.
        first, _ = self.loaded()
        skus, views, matrix = self.cache_payload(first)
        self.rewrite_cache(first, skus[:1] + skus[2:], views[:1] + views[2:], matrix[[0, 2]])

        backbone = Backbone()
        gallery = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), backbone)
        self.assertTrue(gallery.load())
        self.assertEqual(backbone.calls, 0, "a readable shorter cache is still a hit")
        self.assertEqual([entry.sku_id for entry in gallery._entries], ["sku_a", "sku_b"])
        np.testing.assert_array_equal(gallery._matrix, first._matrix[[0, 2]])

    def test_a_nonnumeric_cache_matrix_is_not_served(self):
        # Only the float matrix _cache_write produces is a cache. A complex
        # matrix is finite and casts without complaint, so it answers with a
        # different matrix than the one on disk; integer, boolean and text
        # matrices are not embeddings at all.
        first, _ = self.loaded()
        skus, views, matrix = self.cache_payload(first)
        others = {
            "complex": matrix.astype(np.complex128),
            "timedelta": np.zeros_like(matrix, dtype="timedelta64[s]"),
            "integer": (matrix * 8).astype(np.int32),
            "boolean": np.ones_like(matrix, dtype=bool),
            "text": np.zeros(matrix.shape, dtype="<U4"),
        }
        for name, other in others.items():
            with self.subTest(matrix=name):
                self.rewrite_cache(first, skus, views, other)
                self.re_embeds(first, f"a {name} matrix")

    def test_a_loaded_gallery_reloads_when_its_source_moves(self):
        gallery, backbone = self.loaded()
        self.assertTrue(gallery.load())
        self.assertEqual(backbone.calls, 1, "an unchanged store must not re-embed")

        write_view(os.path.join(self.directory.name, "sku_b.jpeg"), seed=99)
        self.assertTrue(gallery.load())
        self.assertEqual(backbone.calls, 2)

        with open(os.path.join(self.directory.name, "stock.csv"), "a", encoding="utf-8") as handle:
            handle.write("sku_c,new product\n")
        self.assertTrue(gallery.load())
        self.assertEqual(backbone.calls, 3, "stock.csv is part of what an embedding is pinned to")
        self.assertEqual(gallery.view_count(), 3)

    def test_a_valid_cache_needs_the_pinned_backbone_on_disk(self):
        # The regression: a cache hit never opens load_dinov2, so a deleted or
        # altered checkpoint used to read as a loaded gallery and then raise on
        # the first match - after the server had reported itself ready.
        self.loaded()
        gallery = self.real_gallery()
        self.assertTrue(gallery.load(), "the cache serves while the pin holds")

        os.environ["DORAYAKI_DINOV2_WEIGHTS"] = os.path.join(self.directory.name, "gone.safetensors")
        missing = self.real_gallery()
        self.assertFalse(missing.load())
        self.assertFalse(missing.available())

        os.environ["DORAYAKI_DINOV2_WEIGHTS"] = self.write(
            os.path.join(self.directory.name, "altered.safetensors"), b"a different checkpoint"
        )
        altered = self.real_gallery()
        self.assertFalse(altered.load())
        self.assertFalse(altered.available())

    def test_an_altered_backbone_re_embeds_rather_than_serves(self):
        # With a backbone that can still load, an altered file must re-embed.
        self.loaded()
        self.assertTrue(
            a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), UnreachableBackbone()).load(),
            "the cache still serves on the pinned file",
        )
        self.write(self.weights, b"a different checkpoint")
        backbone = Backbone()
        rebuilt = a_stored_gallery(self.directory.name, ("sku_a", "sku_b"), backbone)
        self.assertTrue(rebuilt.load())
        self.assertEqual(backbone.calls, 1, "an unpinned cache must not be served")


class OfflineWeightsTests(unittest.TestCase):
    def test_a_missing_weights_file_names_the_path_and_does_not_fetch(self):
        previous = os.environ.get("DORAYAKI_DINOV2_WEIGHTS")
        os.environ["DORAYAKI_DINOV2_WEIGHTS"] = os.path.join(
            "does-not-exist", "dinov2_vits14_lvd142m.safetensors"
        )
        try:
            with self.assertRaises(FileNotFoundError) as raised:
                load_dinov2(224)
        finally:
            if previous is None:
                os.environ.pop("DORAYAKI_DINOV2_WEIGHTS", None)
            else:
                os.environ["DORAYAKI_DINOV2_WEIGHTS"] = previous
        message = str(raised.exception)
        self.assertIn("does-not-exist", message)
        self.assertIn(DINOV2_SHA256, message)
        self.assertIn("does not download", message)

    def test_the_default_path_stays_inside_the_repo(self):
        previous = os.environ.pop("DORAYAKI_DINOV2_WEIGHTS", None)
        try:
            self.assertTrue(dinov2_weights_path().endswith(
                os.path.join("weights", "dinov2_vits14_lvd142m.safetensors")
            ))
        finally:
            if previous is not None:
                os.environ["DORAYAKI_DINOV2_WEIGHTS"] = previous


def _head_payload(class_ids, torch):
    projection = torch.nn.Linear(384, 128)
    prototypes = torch.nn.functional.normalize(torch.randn(len(class_ids), 128), dim=1)
    return {
        "class_ids": list(class_ids),
        "image_size": 518,
        "base_model": DINOV2_MODEL,
        "projection": {key: value.detach().cpu() for key, value in projection.state_dict().items()},
        "prototypes": prototypes.detach().cpu(),
    }


class HeadPinTests(unittest.TestCase):
    def test_a_stale_head_cannot_label_this_catalog(self):
        import tempfile

        import torch

        with tempfile.TemporaryDirectory() as directory:
            with open(os.path.join(directory, "stock.csv"), "w", encoding="utf-8") as handle:
                handle.write("SKU_ID,product_name\nsku_a,A\nsku_b,B\nsku_c,C\n")
            gallery = ProductGallery(directory, ["sku_a", "sku_b", "sku_c"])
            gallery._available = True
            path = gallery._head_path
            subset = ["sku_c", "sku_a"]
            digest = catalog_digest(directory, subset)

            foreign = _head_payload(["other_store", "sku_a"], torch)
            torch.save(foreign, path)
            self.assertFalse(gallery.load_head())
            self.assertIsNone(gallery._head)
            self.assertIn("catalog_mismatch", gallery.head_reason)
            self.assertTrue(gallery.available())

            legacy = _head_payload(subset, torch)
            torch.save(legacy, path)
            self.assertTrue(gallery.load_head())
            self.assertEqual(gallery.head_reason, "legacy_unpinned")
            self.assertEqual(gallery._head["class_ids"], subset)
            self.assertNotIn("other_store", gallery._head["class_ids"])

            stale = _head_payload(subset, torch)
            stale["catalog_digest"] = "0" * 64
            torch.save(stale, path)
            self.assertFalse(gallery.load_head())
            self.assertIsNone(gallery._head)
            self.assertIn("catalog_mismatch", gallery.head_reason)
            self.assertTrue(gallery.available())

            wrong_backbone = _head_payload(subset, torch)
            wrong_backbone["catalog_digest"] = digest
            wrong_backbone["base_model"] = "vit_base_patch14_dinov2.lvd142m"
            torch.save(wrong_backbone, path)
            self.assertFalse(gallery.load_head())
            self.assertIn("backbone_mismatch", gallery.head_reason)

            wrong_size = _head_payload(subset, torch)
            wrong_size["catalog_digest"] = digest
            wrong_size["image_size"] = 224
            torch.save(wrong_size, path)
            self.assertFalse(gallery.load_head())
            self.assertIn("image_size_mismatch", gallery.head_reason)

            torch.save(["not", "a", "head"], path)
            self.assertFalse(gallery.load_head())
            self.assertTrue(gallery.head_reason.startswith("malformed"))
            self.assertTrue(gallery.available())
            self.assertIsNone(gallery._head)

            current = _head_payload(subset, torch)
            current["catalog_digest"] = digest
            torch.save(current, path)
            self.assertTrue(gallery.load_head())
            self.assertEqual(gallery.head_reason, "ok")
            self.assertEqual(gallery._head["class_ids"], subset)
            self.assertIsNone(head_class_set(["sku_a", "sku_a", ""]))


class TrainerPinTests(unittest.TestCase):
    def test_head_training_loads_the_pinned_file_and_keeps_preprocess(self):
        import finetune_identity
        import gallery as gallery_module

        seen = []

        class FakeBackbone:
            def eval(self):
                return self

            def to(self, device):
                return self

        sentinel = FakeBackbone()
        def fake_load(image_size):
            seen.append(image_size)
            return sentinel


        # build_model resolves the loader from gallery at call time, so that is
        # the attribute to replace; finetune_identity has no module-level copy.
        original = gallery_module.load_dinov2
        gallery_module.load_dinov2 = fake_load
        try:
            model, device = finetune_identity.build_model(224)
        finally:
            gallery_module.load_dinov2 = original
        self.assertIs(model, sentinel)
        self.assertIn(device, ("cpu", "mps"))
        self.assertEqual(seen, [224])

        image = np.zeros((10, 20, 3), dtype=np.uint8)
        self.assertEqual(finetune_identity.prepare(image, 32).shape, gallery_module._prepare(image, 32).shape)
        self.assertTrue(np.allclose(
            finetune_identity.prepare(image, 32),
            gallery_module._prepare(image, 32),
        ))

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