import numpy as np
from typing import Mapping

from hieromem_core.persistence.faiss_index import FaissIndex
from hieromem_core.persistence.sharded_vector_store import ShardedFaissIndex


def expect(condition: bool, message: str | None = None) -> None:
    """Fail the test when *condition* is false without using ``assert``."""

    if not condition:
        raise AssertionError(message or "Expectation failed")


def test_sharded_faiss_merge_results() -> None:
    shard_a = FaissIndex(dim=2, normalize=False, index_type="hnsw")
    shard_b = FaissIndex(dim=2, normalize=False, index_type="hnsw")

    shard_a.add_embeddings(["a1"], np.array([[1.0, 0.0]], dtype=np.float32))
    shard_b.add_embeddings(["b1"], np.array([[0.0, 1.0]], dtype=np.float32))

    def shard_selector(meta: Mapping[str, str] | None) -> int:
        return 0 if meta and meta.get("namespace") == "alpha" else 1

    sharded = ShardedFaissIndex(
        {0: shard_a, 1: shard_b}, shard_fn=shard_selector)

    # Query should return both shards merged by score
    results = sharded.search(np.array([1.0, 1.0], dtype=np.float32), k=2)
    ids = {rid for rid, _ in results}
    expect(ids == {"a1", "b1"}, "should return ids from both shards")

    # Verify routing respects provided metadata for inserts
    sharded.add_embeddings(["a2"], np.array(
        [[2.0, 0.0]], dtype=np.float32), meta={"namespace": "alpha"})
    expect("a2" in shard_a.ids, "shard selector should route to shard A")
    expect("a2" not in shard_b.ids, "shard B should not receive alpha namespace")
