"""Tests for model discovery: the Hub search client and its measurements.

Nothing here touches the network.  Every function that would make a request is
monkeypatched, and the safetensors fixtures are packed in the test — the same
approach ``tests/test_semantic.py`` uses for the projection tests.

The dtype cases matter: the code-tuned sibling this project settled on,
``minishlab/potion-code-16M-v2``, is stored as **F16** (63 457 x 256), while the
general ``potion-base-8M`` is F32.  Reading an F16 file with the F32 path would
silently return garbage rows, so the reader is pinned here.
"""

import json
import os
import struct
import tempfile
import unittest
import urllib.parse
from unittest import mock

from hdc import hub
from hdc.semantic import StaticEmbeddings

RAW_ONE = {
    "id": "minishlab/potion-base-8M",
    "downloads": 447_769,
    "likes": 12,
    "library_name": "model2vec",
    "pipeline_tag": "feature-extraction",
    "gated": False,
    "tags": ["model2vec", "license:mit", "feature-extraction"],
    "siblings": [{"rfilename": "model.safetensors"}, {"rfilename": "tokenizer.json"}],
    "sha": "abc123",
}

RAW_TWO = {
    "id": "jinaai/jina-embeddings-v2-base-code",
    "downloads": 156_469,
    "likes": 40,
    "library_name": "sentence-transformers",
    "pipeline_tag": "feature-extraction",
    "gated": True,
    "tags": ["license:apache-2.0", "sentence-transformers"],
    "siblings": [{"rfilename": "pytorch_model.bin"}],
    "sha": "def456",
}


DTYPE_BYTES = {"F32": 4, "F16": 2, "BF16": 2, "I8": 1}


def header(tensors: dict) -> dict:
    """A parsed safetensors header keyed the way ``hub.safetensors`` returns it.

    ``bytes`` is derived from shape x dtype, because that is exactly what the
    real header's ``data_offsets`` gives without reading the tensor buffer.
    """
    return {"repo": "x", "tensors": tensors,
            "bytes": sum(hub._shape_product(t["shape"]) * DTYPE_BYTES.get(t["dtype"], 4)
                         for t in tensors.values()),
            "params": sum(hub._shape_product(t["shape"]) for t in tensors.values())}


def write_static(directory: str, dtype: str, rows: list[list[float]],
                 vocab: dict | None = None) -> str:
    """Pack a tiny safetensors file in ``dtype`` plus a tokenizer, no libraries."""
    flat = [value for row in rows for value in row]
    if dtype == "F16":
        payload = struct.pack("<%de" % len(flat), *flat)
    elif dtype == "BF16":
        # bf16 is the top half of the f32 word
        halves = [struct.unpack("<I", struct.pack("<f", value))[0] >> 16 for value in flat]
        payload = struct.pack("<%dH" % len(flat), *halves)
    else:
        payload = struct.pack("<%df" % len(flat), *flat)
    meta = json.dumps({
        "embeddings": {"dtype": dtype, "shape": [len(rows), len(rows[0])],
                       "data_offsets": [0, len(payload)]}
    }).encode("utf-8")
    meta += b" " * ((8 - len(meta) % 8) % 8)
    with open(os.path.join(directory, "model.safetensors"), "wb") as fh:
        fh.write(struct.pack("<Q", len(meta)))
        fh.write(meta)
        fh.write(payload)
    with open(os.path.join(directory, "tokenizer.json"), "w", encoding="utf-8") as fh:
        json.dump({"model": {"vocab": vocab or {"a": 0}, "type": "WordPiece",
                             "continuing_subword_prefix": "##"}}, fh)
    return directory


class TestNormalise(unittest.TestCase):
    def test_record_is_reduced_and_licence_is_unpacked(self):
        record = hub.normalise(RAW_ONE)
        self.assertEqual(record["id"], "minishlab/potion-base-8M")
        self.assertEqual(record["downloads"], 447_769)
        self.assertEqual(record["likes"], 12)
        self.assertEqual(record["library"], "model2vec")
        self.assertEqual(record["license"], "mit")
        self.assertEqual(record["files"],
                         ["model.safetensors", "tokenizer.json"])
        self.assertFalse(record["gated"])
        self.assertTrue(hub.normalise(RAW_TWO)["gated"])

    def test_missing_fields_do_not_raise(self):
        record = hub.normalise({})
        self.assertEqual(record["id"], "")
        self.assertEqual(record["downloads"], 0)
        self.assertEqual(record["license"], "")


class TestStaticTable(unittest.TestCase):
    def test_single_rank_two_tensor_is_static(self):
        is_static, rows, dim = hub.static_table(header({"embeddings": {"dtype": "F16",
                                                                      "shape": [63457, 256]}}))
        self.assertTrue(is_static)
        self.assertEqual((rows, dim), (63457, 256))

    def test_many_tensors_are_a_transformer(self):
        many = {name: {"dtype": "F32", "shape": [8, 8]}
                for name in ("q_proj", "k_proj", "v_proj")}
        self.assertEqual(hub.static_table(header(many)), (False, 0, 0))

    def test_rank_one_and_rank_three_are_rejected(self):
        rank_one = header({"embeddings": {"dtype": "F32", "shape": [256]}})
        rank_three = header({"embeddings": {"dtype": "F32", "shape": [2, 4, 8]}})
        self.assertEqual(hub.static_table(rank_one), (False, 0, 0))
        self.assertEqual(hub.static_table(rank_three), (False, 0, 0))

    def test_no_header_is_not_static(self):
        self.assertEqual(hub.static_table(None), (False, 0, 0))


class TestCodeCoverage(unittest.TestCase):
    def test_full_vocabulary_covers_everything_whole(self):
        vocab = {token: i for i, token in enumerate(hub.MACHINE_CODE)}
        coverage = hub.code_coverage(vocab)
        self.assertEqual(coverage["coverage"], 1.0)
        self.assertEqual(coverage["whole_fraction"], 1.0)
        self.assertEqual(coverage["any_fraction"], 1.0)
        self.assertEqual(coverage["missing"], [])

    def test_fragments_are_worth_less_than_whole_tokens(self):
        # every token is reachable, but only as first-char + "##" remainder
        vocab = {}
        for token in hub.MACHINE_CODE:
            vocab[token[0]] = len(vocab)
            vocab["##" + token[1:]] = len(vocab)
        coverage = hub.code_coverage(vocab)
        self.assertEqual(coverage["whole_fraction"], 0.0)
        self.assertEqual(coverage["any_fraction"], 1.0)
        # the 0.7 / 0.3 weighting is visible in the number itself
        self.assertAlmostEqual(coverage["coverage"], 0.3 * coverage["any_fraction"], places=6)
        self.assertLess(coverage["coverage"], 0.5)

    def test_token_with_no_pieces_is_missing(self):
        coverage = hub.code_coverage({}, tokens=["qqq", "mov"])
        self.assertEqual(coverage["missing"], ["qqq", "mov"])
        self.assertEqual(coverage["coverage"], 0.0)


class TestScore(unittest.TestCase):
    def setUp(self):
        self.coverage = hub.code_coverage({token: i for i, token in enumerate(hub.MACHINE_CODE)})

    def test_components_match_the_weights(self):
        graded = hub.score(hub.normalise(RAW_ONE),
                           header({"embeddings": {"dtype": "F32", "shape": [29528, 256]}}),
                           self.coverage)
        self.assertEqual(set(graded["components"]), set(hub.WEIGHTS))

    def test_weighted_sum_is_the_total(self):
        graded = hub.score(hub.normalise(RAW_ONE),
                           header({"embeddings": {"dtype": "F32", "shape": [29528, 256]}}),
                           self.coverage)
        expected = sum(hub.WEIGHTS[name] * value
                       for name, value in graded["components"].items())
        self.assertAlmostEqual(graded["total"], expected, places=2)

    def test_transformer_scores_zero_on_runtime_free(self):
        graded = hub.score(hub.normalise(RAW_TWO), None, self.coverage)
        self.assertEqual(graded["components"]["runtime_free"], 0.0)
        self.assertFalse(graded["static"])

    def test_small_static_beats_large_transformer_at_equal_coverage(self):
        small = hub.score(hub.normalise(RAW_ONE),
                          header({"embeddings": {"dtype": "F32", "shape": [29528, 256]}}),
                          self.coverage)
        large = hub.score(hub.normalise(RAW_TWO),
                          header({"model.embed_tokens": {"dtype": "F16", "shape": [250000, 1024]},
                                  "layers.0.q_proj": {"dtype": "F16", "shape": [1024, 1024]}}),
                          self.coverage)
        self.assertGreater(large["bytes"], 400_000_000)
        self.assertGreater(small["total"], large["total"])


class TestClamps(unittest.TestCase):
    def test_size_score_is_clamped_and_monotonic(self):
        self.assertEqual(hub._size_score(30_000_000), 1.0)
        self.assertEqual(hub._size_score(600_000_000), 0.0)
        self.assertEqual(hub._size_score(0), 0.0)
        self.assertGreater(hub._size_score(60_000_000), hub._size_score(200_000_000))

    def test_dim_score_prefers_256_and_clamps(self):
        self.assertEqual(hub._dim_score(256), 1.0)
        self.assertEqual(hub._dim_score(0), 0.0)
        self.assertLess(hub._dim_score(128), hub._dim_score(256))
        self.assertGreaterEqual(hub._dim_score(2048), 0.0)


class TestSearch(unittest.TestCase):
    def setUp(self):
        self.calls: list[str] = []

    def fake_get_json(self, url, timeout=45):
        self.calls.append(url)
        return [RAW_ONE, RAW_TWO]

    def query(self) -> dict:
        parsed = urllib.parse.urlparse(self.calls[-1])
        return urllib.parse.parse_qs(parsed.query)

    def test_text_query_keeps_the_task_filter(self):
        with mock.patch.object(hub, "_get_json", self.fake_get_json):
            records = hub.search("code embeddings")
        self.assertEqual(len(records), 2)
        self.assertEqual(records[0]["id"], RAW_ONE["id"])
        query = self.query()
        self.assertEqual(query["search"], ["code embeddings"])
        self.assertEqual(query["filter"], ["feature-extraction"])

    def test_author_query_drops_the_filter(self):
        # measured API quirk: author + filter returns zero rows, silently
        with mock.patch.object(hub, "_get_json", self.fake_get_json):
            hub.search(author="minishlab")
        query = self.query()
        self.assertNotIn("filter", query)
        self.assertEqual(query["author"], ["minishlab"])

    def test_non_list_payload_is_an_empty_result(self):
        with mock.patch.object(hub, "_get_json", lambda url, timeout=45: {"error": "nope"}):
            self.assertEqual(hub.search("anything"), [])


class TestCandidatePool(unittest.TestCase):
    def test_queries_are_de_duplicated_and_ordered_by_downloads(self):
        shared = hub.normalise(dict(RAW_ONE, id="a/shared", downloads=5))
        more = hub.normalise(dict(RAW_ONE, id="b/bigger", downloads=9))
        least = hub.normalise(dict(RAW_ONE, id="c/small", downloads=1))

        def fake_search(text=None, **kwargs):
            return [shared, more, least, more]      # the same repo from two queries

        with mock.patch.object(hub, "search", fake_search):
            pool = hub.candidate_pool(limit=10)
        self.assertEqual([row["id"] for row in pool], ["b/bigger", "a/shared", "c/small"])

    def test_limit_is_applied(self):
        rows = [hub.normalise(dict(RAW_ONE, id="r/%d" % i, downloads=i)) for i in range(5)]
        with mock.patch.object(hub, "search", lambda *a, **k: rows):
            self.assertEqual(len(hub.candidate_pool(limit=2)), 2)


class TestBakeoff(unittest.TestCase):
    def test_offline_uses_the_cache_and_touches_no_network(self):
        entry = {"id": "minishlab/potion-code-16M-v2", "sha": "s", "downloads": 7,
                 "components": {"runtime_free": 1.0}, "total": 1.0,
                 "static": True, "rows": 63457, "dim": 256, "bytes": 32_489_984,
                 "coverage": None, "license": "mit"}
        with mock.patch.object(hub, "load_cache", lambda: {"minishlab/potion-code-16M-v2@s": entry}), \
             mock.patch.object(hub, "save_cache", lambda cache: None), \
             mock.patch.object(hub, "candidate_pool", lambda **kwargs: []):
            result = hub.bakeoff(offline=True)
        self.assertEqual(result["pool"], 1)
        self.assertEqual(result["graded"][0]["id"], "minishlab/potion-code-16M-v2")
        self.assertTrue(result["graded"][0]["cached"])
        self.assertEqual(hub.best(result)["total"], 1.0)


class TestHeaderRange(unittest.TestCase):
    def test_two_range_requests_give_shape_size_and_params(self):
        meta = json.dumps({
            "embeddings": {"dtype": "F16", "shape": [2000, 256],
                           "data_offsets": [0, 1_024_000]},
        }).encode("utf-8")
        seen: list[dict] = []

        def fake_request(url, headers=None, timeout=45):
            seen.append(headers or {})
            if (headers or {}).get("Range") == "bytes=0-7":
                return struct.pack("<Q", len(meta))
            return meta

        with mock.patch.object(hub, "_request", fake_request):
            parsed = hub.safetensors("minishlab/potion-code-16M-v2")
        self.assertEqual(len(seen), 2)
        self.assertEqual(parsed["bytes"], 1_024_000)
        self.assertEqual(parsed["params"], 2000 * 256)
        self.assertEqual(parsed["tensors"]["embeddings"]["dtype"], "F16")
        self.assertEqual(hub.static_table(parsed), (True, 2000, 256))


class TestDtypes(unittest.TestCase):
    """F16 is what the code model actually ships; F32 is what the reader used to assume."""

    ROWS = [
        [1.0, 0.5, -1.5, 0.25],          # every value exact in f16 and bf16
        [2.0, 0.0, 1.875, -0.25],
    ]

    def build(self, dtype: str) -> StaticEmbeddings:
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        return StaticEmbeddings(write_static(tmp.name, dtype, self.ROWS))

    def test_float16_rows_decode_exactly(self):
        model = self.build("F16")
        self.assertEqual(model.dtype, "F16")
        self.assertEqual(model.itemsize, 2)
        self.assertEqual(model.row(0), (1.0, 0.5, -1.5, 0.25))
        self.assertEqual(model.row(1), (2.0, 0.0, 1.875, -0.25))

    def test_bfloat16_rows_decode_exactly(self):
        model = self.build("BF16")
        self.assertEqual(model.dtype, "BF16")
        self.assertEqual(model.row(0), (1.0, 0.5, -1.5, 0.25))
        self.assertEqual(model.row(1), (2.0, 0.0, 1.875, -0.25))

    def test_rows_do_not_overlap(self):
        # the offset is rows * dim * itemsize; a wrong stride returns row 0 twice
        model = self.build("F16")
        self.assertNotEqual(model.row(0), model.row(1))

    def test_unsupported_dtype_is_rejected(self):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        write_static(tmp.name, "I8", self.ROWS)
        with self.assertRaises(ValueError):
            StaticEmbeddings(tmp.name)


class TestCacheLifetime(unittest.TestCase):
    """The score cache outlives the code that wrote it.

    An older run of ``python -m hdc hub --offline`` raised ``KeyError`` on a
    coverage entry that predated a new field, so ``load_cache`` has to drop what
    it cannot understand instead of handing it to the report.
    """

    def write(self, payload):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        path = os.path.join(tmp.name, "scores.json")
        with open(path, "w", encoding="utf-8") as fh:
            json.dump(payload, fh)
        return path

    def test_current_entries_are_kept(self):
        entry = {"id": "a/b", "coverage": {"any_fraction": 1.0, "whole_fraction": 0.5}}
        path = self.write({"a/b@sha": entry})
        with mock.patch.object(hub, "cache_path", return_value=path):
            self.assertIn("a/b@sha", hub.load_cache())

    def test_stale_entries_are_dropped(self):
        stale = {"id": "a/b", "coverage": {"coverage": 1.0, "missing": []}}
        fresh = {"id": "c/d", "coverage": {"any_fraction": 1.0, "whole_fraction": 0.4}}
        path = self.write({"a/b@old": stale, "c/d@new": fresh})
        with mock.patch.object(hub, "cache_path", return_value=path):
            cache = hub.load_cache()
        self.assertEqual(sorted(cache), ["c/d@new"])

    def test_missing_or_corrupt_cache_is_empty(self):
        with mock.patch.object(hub, "cache_path",
                               return_value=os.path.join("does", "not", "exist.json")):
            self.assertEqual(hub.load_cache(), {})

    def test_report_survives_a_coverage_without_any_fraction(self):
        # older cached rows reach report() when offline; it must not raise
        entry = {"id": "a/b", "total": 0.9, "static": True, "dim": 256,
                 "rows": 10, "bytes": 30_000_000, "license": "mit",
                 "components": {"code": 0.5},
                 "coverage": {"coverage": 1.0, "missing": [], "tokens": 42}}
        with mock.patch.object(hub, "bakeoff", return_value={"pool": 1, "graded": [entry]}):
            self.assertIn("a/b", hub.report())


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