import math
from genesis.rag.store import RagStore


def _vec(*ones):
    v = [0.0] * 8
    for i in ones:
        v[i] = 1.0
    n = math.sqrt(sum(x * x for x in v))
    return [x / n for x in v]


def test_search_returns_most_similar_chunk():
    s = RagStore(":memory:")
    s.reset_scope("p1")
    s.add("p1", ["订单模块处理创建", "用户认证登录"], [_vec(0, 1), _vec(4, 5)])
    res = s.search("p1", _vec(0, 1), k=1)
    assert res == ["订单模块处理创建"]
    # k<=0 健壮性
    assert s.search("p1", _vec(0, 1), k=0) == []
    s.close()


def test_reset_scope_clears():
    s = RagStore(":memory:")
    s.reset_scope("p1")
    s.add("p1", ["a"], [_vec(0)])
    s.reset_scope("p1")
    assert s.search("p1", _vec(0), k=3) == []
    # chunks 与 embeddings 长度不一致应抛 ValueError
    try:
        import pytest
        with pytest.raises(ValueError):
            s.add("p1", ["a", "b"], [_vec(0)])
    finally:
        s.close()


def test_concurrent_add_and_search_no_crash():
    import threading
    s = RagStore(":memory:")
    s.reset_scope("p1")

    def worker(i):
        # 每个 worker 线程内交替执行 add 与 search，让读写真正并发跨线程运行
        for j in range(5):
            s.add("p1", [f"chunk-{i}-{j}"], [_vec(i % 8)])
            res = s.search("p1", _vec(i % 8), k=3)
            assert all(isinstance(r, str) for r in res)

    try:
        threads = [threading.Thread(target=worker, args=(i,)) for i in range(8)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        res = s.search("p1", _vec(0), k=3)
        assert len(res) <= 3
        assert all(isinstance(r, str) for r in res)
    finally:
        s.close()


def test_count_empty_scope_returns_zero():
    s = RagStore(":memory:")
    assert s.count("missing") == 0


def test_count_after_add_returns_chunk_count():
    s = RagStore(":memory:")
    s.add("p1", ["a", "b", "c"], [_vec(0), _vec(1), _vec(2)])
    assert s.count("p1") == 3


def test_count_scopes_are_isolated():
    s = RagStore(":memory:")
    s.add("p1", ["a"], [_vec(0)])
    s.add("p2", ["b", "c"], [_vec(1), _vec(2)])
    assert s.count("p1") == 1
    assert s.count("p2") == 2


def test_count_after_reset_returns_zero():
    s = RagStore(":memory:")
    s.add("p1", ["a", "b"], [_vec(0), _vec(1)])
    assert s.count("p1") == 2
    s.reset_scope("p1")
    assert s.count("p1") == 0