"""Tests for rbsp.indexer.vector_store -- Unified Vector Store."""

from __future__ import annotations

import math
import os
import random
import sys
import threading
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

from rbsp.indexer.vector_store import VectorStore
from rbsp.core.hnsw import HNSWConfig


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _random_vector(dim, rng):
    """Generate a random unit-ish vector."""
    v = [rng.gauss(0, 1) for _ in range(dim)]
    norm = math.sqrt(sum(x * x for x in v))
    if norm > 0:
        v = [x / norm for x in v]
    return v


# ---------------------------------------------------------------------------
# Basic operations
# ---------------------------------------------------------------------------


class TestVectorStoreBasics(unittest.TestCase):
    """Basic insert, search, remove, properties."""

    def test_insert_and_size(self):
        vs = VectorStore(dim=4)
        vs.insert("a", [1.0, 0.0, 0.0, 0.0])
        vs.insert("b", [0.0, 1.0, 0.0, 0.0])
        self.assertEqual(vs.size, 2)
        self.assertEqual(len(vs), 2)

    def test_contains(self):
        vs = VectorStore(dim=3)
        vs.insert("x", [1.0, 0.0, 0.0])
        self.assertIn("x", vs)
        self.assertNotIn("y", vs)

    def test_dim_inferred(self):
        vs = VectorStore()
        self.assertIsNone(vs.dim)
        vs.insert("a", [1.0, 2.0, 3.0])
        self.assertEqual(vs.dim, 3)

    def test_remove(self):
        vs = VectorStore(dim=2)
        vs.insert("a", [1.0, 0.0])
        self.assertTrue(vs.remove("a"))
        self.assertEqual(vs.size, 0)
        self.assertNotIn("a", vs)

    def test_remove_nonexistent(self):
        vs = VectorStore(dim=2)
        self.assertFalse(vs.remove("ghost"))

    def test_get_vector(self):
        vs = VectorStore(dim=3)
        vs.insert("a", [1.0, 2.0, 3.0])
        vec = vs.get_vector("a")
        self.assertEqual(vec, [1.0, 2.0, 3.0])

    def test_get_vector_nonexistent(self):
        vs = VectorStore(dim=2)
        self.assertIsNone(vs.get_vector("nope"))

    def test_repr(self):
        vs = VectorStore(dim=4)
        r = repr(vs)
        self.assertIn("VectorStore", r)
        self.assertIn("dim=4", r)

    def test_update_overwrites(self):
        vs = VectorStore(dim=2)
        vs.insert("a", [1.0, 0.0])
        vs.insert("a", [0.0, 1.0])
        self.assertEqual(vs.size, 1)
        self.assertEqual(vs.get_vector("a"), [0.0, 1.0])


# ---------------------------------------------------------------------------
# Batch insert
# ---------------------------------------------------------------------------


class TestBatchInsert(unittest.TestCase):
    """batch_insert for bulk loading."""

    def test_basic(self):
        vs = VectorStore(dim=3)
        ids = ["a", "b", "c"]
        vecs = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
        vs.batch_insert(ids, vecs)
        self.assertEqual(vs.size, 3)

    def test_length_mismatch_raises(self):
        vs = VectorStore(dim=2)
        with self.assertRaises(ValueError):
            vs.batch_insert(["a"], [[1, 0], [0, 1]])

    def test_large_batch(self):
        rng = random.Random(42)
        vs = VectorStore(dim=16)
        ids = [f"v_{i}" for i in range(100)]
        vecs = [_random_vector(16, rng) for _ in range(100)]
        vs.batch_insert(ids, vecs)
        self.assertEqual(vs.size, 100)


# ---------------------------------------------------------------------------
# Exact search
# ---------------------------------------------------------------------------


class TestExactSearch(unittest.TestCase):
    """Brute-force exact nearest neighbour search."""

    def test_finds_nearest(self):
        vs = VectorStore(dim=3, distance="l2")
        vs.insert("a", [1.0, 0.0, 0.0])
        vs.insert("b", [0.0, 1.0, 0.0])
        vs.insert("c", [0.0, 0.0, 1.0])

        results = vs.search([1.0, 0.0, 0.0], k=1, mode="exact")
        self.assertEqual(results[0][0], "a")
        self.assertAlmostEqual(results[0][1], 0.0)

    def test_returns_k(self):
        vs = VectorStore(dim=2, distance="l2")
        for i in range(10):
            vs.insert(f"v{i}", [float(i), 0.0])
        results = vs.search([0.0, 0.0], k=3, mode="exact")
        self.assertEqual(len(results), 3)

    def test_sorted_ascending(self):
        vs = VectorStore(dim=2, distance="l2")
        vs.insert("a", [0.0, 0.0])
        vs.insert("b", [1.0, 0.0])
        vs.insert("c", [3.0, 0.0])
        results = vs.search([0.0, 0.0], k=3, mode="exact")
        dists = [d for _, d in results]
        self.assertEqual(dists, sorted(dists))

    def test_empty_store(self):
        vs = VectorStore(dim=2)
        results = vs.search([1.0, 0.0], k=5, mode="exact")
        self.assertEqual(results, [])


# ---------------------------------------------------------------------------
# HNSW search
# ---------------------------------------------------------------------------


class TestHNSWSearch(unittest.TestCase):
    """HNSW approximate nearest neighbour search."""

    def test_finds_nearest(self):
        cfg = HNSWConfig(M=4, ef_construction=50, ef_search=20, distance="l2")
        vs = VectorStore(dim=4, distance="l2", hnsw_config=cfg)
        rng = random.Random(42)

        target = [1.0, 0.0, 0.0, 0.0]
        vs.insert("target", target)

        for i in range(50):
            vs.insert(f"noise_{i}", _random_vector(4, rng))

        results = vs.search(target, k=1, mode="hnsw")
        self.assertEqual(results[0][0], "target")

    def test_returns_k_results(self):
        cfg = HNSWConfig(M=4, ef_construction=50, ef_search=20, distance="l2")
        vs = VectorStore(dim=4, distance="l2", hnsw_config=cfg)
        rng = random.Random(42)
        for i in range(20):
            vs.insert(f"v_{i}", _random_vector(4, rng))
        results = vs.search([1.0, 0.0, 0.0, 0.0], k=5, mode="hnsw")
        self.assertLessEqual(len(results), 5)

    def test_default_mode_is_hnsw(self):
        vs = VectorStore(dim=2)
        vs.insert("a", [1.0, 0.0])
        # Default mode should work (hnsw)
        results = vs.search([1.0, 0.0], k=1)
        self.assertEqual(len(results), 1)


# ---------------------------------------------------------------------------
# Binary search
# ---------------------------------------------------------------------------


class TestBinarySearch(unittest.TestCase):
    """Binary quantization Hamming distance search."""

    def test_finds_nearest(self):
        vs = VectorStore(dim=8, distance="l2")
        # Insert orthogonal-ish vectors
        vs.insert("a", [1, 1, 1, 1, 0, 0, 0, 0])
        vs.insert("b", [0, 0, 0, 0, 1, 1, 1, 1])
        vs.insert("c", [1, 1, 1, 0, 0, 0, 0, 0])

        results = vs.search([1, 1, 1, 1, 0, 0, 0, 0], k=1, mode="binary")
        self.assertEqual(results[0][0], "a")
        self.assertAlmostEqual(results[0][1], 0.0)

    def test_returns_hamming_distance(self):
        vs = VectorStore(dim=4)
        vs.insert("a", [1.0, 1.0, 0.0, 0.0])
        vs.insert("b", [0.0, 0.0, 1.0, 1.0])
        results = vs.search([1.0, 1.0, 0.0, 0.0], k=2, mode="binary")
        # "a" should have distance 0
        self.assertEqual(results[0][0], "a")
        self.assertEqual(results[0][1], 0.0)


# ---------------------------------------------------------------------------
# Hybrid search
# ---------------------------------------------------------------------------


class TestHybridSearch(unittest.TestCase):
    """Hybrid: binary pre-filter -> exact rerank."""

    def test_basic_hybrid(self):
        vs = VectorStore(dim=16, distance="l2")
        rng = random.Random(42)

        target = _random_vector(16, rng)
        vs.insert("target", target)
        for i in range(100):
            vs.insert(f"noise_{i}", _random_vector(16, rng))

        results = vs.search(target, k=5, mode="hybrid")
        self.assertGreater(len(results), 0)
        # Target should be in top results
        top_ids = [id_ for id_, _ in results]
        self.assertIn("target", top_ids)

    def test_hybrid_recall(self):
        """Hybrid mode should have decent recall on a moderate dataset."""
        rng = random.Random(42)
        dim = 32
        vs_exact = VectorStore(dim=dim, distance="l2")
        vs_hybrid = VectorStore(dim=dim, distance="l2")

        vecs = [_random_vector(dim, rng) for _ in range(200)]
        ids = [f"v_{i}" for i in range(200)]

        for id_, vec in zip(ids, vecs):
            vs_exact.insert(id_, vec)
            vs_hybrid.insert(id_, vec)

        query = _random_vector(dim, rng)
        exact_results = vs_exact.search(query, k=10, mode="exact")
        hybrid_results = vs_hybrid.search(query, k=10, mode="hybrid")

        exact_ids = {id_ for id_, _ in exact_results}
        hybrid_ids = {id_ for id_, _ in hybrid_results}
        recall = len(exact_ids & hybrid_ids) / len(exact_ids) if exact_ids else 1.0

        # Hybrid should have >70% recall
        self.assertGreaterEqual(recall, 0.7, f"Hybrid recall {recall:.2f} < 0.7")

    def test_empty_store(self):
        vs = VectorStore(dim=4)
        results = vs.search([1, 0, 0, 0], k=5, mode="hybrid")
        self.assertEqual(results, [])


# ---------------------------------------------------------------------------
# Invalid mode
# ---------------------------------------------------------------------------


class TestInvalidMode(unittest.TestCase):
    """Invalid search mode raises ValueError."""

    def test_invalid_mode(self):
        vs = VectorStore(dim=2)
        vs.insert("a", [1.0, 0.0])
        with self.assertRaises(ValueError):
            vs.search([1.0, 0.0], mode="invalid")


# ---------------------------------------------------------------------------
# Quantizer rebuild
# ---------------------------------------------------------------------------


class TestRebuildQuantizer(unittest.TestCase):
    """rebuild_quantizer refits binary thresholds."""

    def test_rebuild(self):
        vs = VectorStore(dim=4)
        rng = random.Random(42)
        for i in range(50):
            vs.insert(f"v_{i}", _random_vector(4, rng))

        vs.rebuild_quantizer()
        # After rebuild, thresholds should be set
        self.assertTrue(vs._quantizer_fitted)
        self.assertIsNotNone(vs._quantizer.thresholds)

    def test_rebuild_with_sample_size(self):
        vs = VectorStore(dim=4)
        rng = random.Random(42)
        for i in range(100):
            vs.insert(f"v_{i}", _random_vector(4, rng))

        vs.rebuild_quantizer(sample_size=20)
        self.assertTrue(vs._quantizer_fitted)

    def test_rebuild_empty_no_error(self):
        vs = VectorStore(dim=4)
        vs.rebuild_quantizer()  # Should not raise


# ---------------------------------------------------------------------------
# Thread safety
# ---------------------------------------------------------------------------


class TestThreadSafety(unittest.TestCase):
    """Concurrent insert and search should not crash."""

    def test_concurrent_inserts(self):
        vs = VectorStore(dim=8, distance="l2")
        rng = random.Random(42)
        errors = []

        def inserter(start):
            try:
                for i in range(50):
                    vs.insert(f"t_{start}_{i}", _random_vector(8, rng))
            except Exception as e:
                errors.append(e)

        threads = [threading.Thread(target=inserter, args=(t,)) for t in range(4)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        self.assertEqual(len(errors), 0, f"Errors: {errors}")
        self.assertEqual(vs.size, 200)

    def test_concurrent_search(self):
        vs = VectorStore(dim=4, distance="l2")
        rng = random.Random(42)
        for i in range(50):
            vs.insert(f"v_{i}", _random_vector(4, rng))

        errors = []

        def searcher():
            try:
                for _ in range(20):
                    vs.search(_random_vector(4, rng), k=5, mode="exact")
            except Exception as e:
                errors.append(e)

        threads = [threading.Thread(target=searcher) for _ in range(4)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        self.assertEqual(len(errors), 0, f"Errors: {errors}")


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