"""Tests for vector store implementations."""

import pytest

from bifrost_rag.store import ChromaStore, Document, VectorStore


class TestDocument:
    def test_creation(self):
        doc = Document(id="doc1", text="hello world", metadata={"source": "test"})
        assert doc.id == "doc1"
        assert doc.text == "hello world"
        assert doc.metadata["source"] == "test"

    def test_default_metadata(self):
        doc = Document(id="doc1", text="hello")
        assert doc.metadata == {}

    def test_embedding_optional(self):
        doc = Document(id="doc1", text="hello")
        assert doc.embedding is None

    def test_embedding_set(self):
        doc = Document(id="doc1", text="hello", embedding=[0.1, 0.2, 0.3])
        assert doc.embedding == [0.1, 0.2, 0.3]


class TestVectorStore:
    def test_abstract_interface(self):
        with pytest.raises(TypeError):
            VectorStore()


class TestChromaStore:
    def test_creation(self):
        store = ChromaStore(collection_name="test")
        assert store.collection_name == "test"

    def test_add_and_query(self):
        store = ChromaStore(collection_name="test_add_query")
        docs = [
            Document(
                id="doc1",
                text="Python is a programming language",
                embedding=[1.0, 0.0, 0.0],
            ),
            Document(
                id="doc2",
                text="JavaScript runs in browsers",
                embedding=[0.0, 1.0, 0.0],
            ),
        ]
        store.add(docs)
        results = store.query(query_embedding=[1.0, 0.0, 0.0], top_k=1)
        assert len(results) == 1
        assert results[0].id == "doc1"

    def test_query_top_k(self):
        store = ChromaStore(collection_name="test_top_k")
        docs = [
            Document(id=f"doc{i}", text=f"text {i}", embedding=[float(i), 0.0, 0.0])
            for i in range(5)
        ]
        store.add(docs)
        results = store.query(query_embedding=[4.0, 0.0, 0.0], top_k=3)
        assert len(results) == 3

    def test_delete(self):
        store = ChromaStore(collection_name="test_delete")
        docs = [
            Document(id="doc1", text="text 1", embedding=[1.0, 0.0]),
            Document(id="doc2", text="text 2", embedding=[0.0, 1.0]),
        ]
        store.add(docs)
        store.delete(["doc1"])
        results = store.query(query_embedding=[1.0, 0.0], top_k=10)
        assert len(results) == 1
        assert results[0].id == "doc2"

    def test_count(self):
        store = ChromaStore(collection_name="test_count")
        assert store.count() == 0
        docs = [
            Document(id="doc1", text="text 1", embedding=[1.0, 0.0]),
            Document(id="doc2", text="text 2", embedding=[0.0, 1.0]),
        ]
        store.add(docs)
        assert store.count() == 2

    def test_add_empty_list(self):
        store = ChromaStore(collection_name="test_empty_add")
        store.add([])
        assert store.count() == 0
