"""
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_SHA256,
    GalleryEntry,
    ProductGallery,
    dinov2_weights_path,
    load_dinov2,
)


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 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_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 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

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