import os
import tempfile

import numpy as np
import pytest

from neurovault.search.hnsw import HNSWIndex


def test_hnsw_add_and_search():
    idx = HNSWIndex(dim=8, M=4)
    rng = np.random.default_rng(42)
    vectors = rng.standard_normal((120, 8), dtype=np.float32)

    for i in range(120):
        idx.add(vectors[i].tolist(), f"mem_{i}".encode())

    assert idx.size == 120

    hits = idx.search(vectors[0].tolist(), k=5)
    assert len(hits) > 0
    assert any(mid == b"mem_0" for mid, _ in hits)


def test_hnsw_empty_index_search():
    idx = HNSWIndex(dim=8)
    assert idx.search([0.0] * 8, k=5) == []


def test_hnsw_save_and_load_roundtrip():
    idx = HNSWIndex(dim=4, M=4)
    rng = np.random.default_rng(7)
    for i in range(20):
        vec = rng.standard_normal(4, dtype=np.float32)
        idx.add(vec.tolist(), f"m{i}".encode())

    with tempfile.NamedTemporaryFile(suffix=".idx", delete=False) as handle:
        path = handle.name
    try:
        idx.save(path)
        loaded = HNSWIndex.load(path)
        assert loaded.size == 20
        assert loaded.dimension == 4
    finally:
        os.unlink(path)


def test_hnsw_load_rejects_invalid_payload(tmp_path):
    path = tmp_path / "index.idx"
    path.write_bytes(b"this-is-not-a-valid-hnsw-index")
    with pytest.raises(Exception):
        HNSWIndex.load(path)
