"""Tests for RAG vector stores."""

import pytest

from logos_adk.rag import Chunk
from logos_adk.rag.vector_chroma import ChromaVectorStore
from logos_adk.rag.vector_postgresql import PostgreSQLVectorStore
from logos_adk.rag.vector_sqlite import SQLiteVectorStore
from logos_adk.rag.vector_store import InMemoryVectorStore


def _chunk(content="text", doc_id="d1", embedding=None, **meta):
    return Chunk(content=content, metadata=meta, document_id=doc_id, embedding=embedding)


class TestInMemoryVectorStore:
    @pytest.mark.asyncio
    async def test_add_and_search(self):
        store = InMemoryVectorStore()
        c = _chunk("hello", embedding=[1.0, 0.0, 0.0])
        await store.add([c])
        results = await store.search([1.0, 0.0, 0.0], limit=1)
        assert len(results) == 1
        assert results[0].content == "hello"

    @pytest.mark.asyncio
    async def test_search_empty(self):
        store = InMemoryVectorStore()
        results = await store.search([1.0, 0.0], limit=5)
        assert results == []

    @pytest.mark.asyncio
    async def test_search_returns_most_similar(self):
        store = InMemoryVectorStore()
        c1 = _chunk("close", doc_id="d1", embedding=[1.0, 0.0, 0.0])
        c2 = _chunk("far", doc_id="d2", embedding=[0.0, 1.0, 0.0])
        await store.add([c1, c2])
        results = await store.search([0.9, 0.1, 0.0], limit=1)
        assert results[0].content == "close"

    @pytest.mark.asyncio
    async def test_search_limit(self):
        store = InMemoryVectorStore()
        chunks = [_chunk(f"c{i}", doc_id=f"d{i}", embedding=[float(i), 0.0]) for i in range(5)]
        await store.add(chunks)
        results = await store.search([4.0, 0.0], limit=2)
        assert len(results) == 2

    @pytest.mark.asyncio
    async def test_search_with_filter(self):
        store = InMemoryVectorStore()
        c1 = _chunk("a", doc_id="d1", embedding=[1.0, 0.0], topic="science")
        c2 = _chunk("b", doc_id="d2", embedding=[1.0, 0.0], topic="art")
        await store.add([c1, c2])
        results = await store.search([1.0, 0.0], limit=5, filter={"topic": "science"})
        assert len(results) == 1
        assert results[0].content == "a"

    @pytest.mark.asyncio
    async def test_delete_by_document_id(self):
        store = InMemoryVectorStore()
        c1 = _chunk("a", doc_id="d1", embedding=[1.0, 0.0])
        c2 = _chunk("b", doc_id="d1", embedding=[0.0, 1.0])
        c3 = _chunk("c", doc_id="d2", embedding=[1.0, 1.0])
        await store.add([c1, c2, c3])
        await store.delete("d1")
        results = await store.search([1.0, 1.0], limit=10)
        assert len(results) == 1
        assert results[0].document_id == "d2"

    @pytest.mark.asyncio
    async def test_clear(self):
        store = InMemoryVectorStore()
        await store.add([_chunk("a", embedding=[1.0])])
        await store.clear()
        results = await store.search([1.0], limit=10)
        assert results == []

    @pytest.mark.asyncio
    async def test_skips_chunks_without_embedding(self):
        store = InMemoryVectorStore()
        c = _chunk("no vec", embedding=None)
        await store.add([c])
        results = await store.search([1.0, 0.0], limit=10)
        assert results == []


class TestSQLiteVectorStore:
    @pytest.mark.asyncio
    async def test_add_and_search(self, tmp_path):
        store = SQLiteVectorStore(str(tmp_path / "vec.db"))
        c = _chunk("hello", embedding=[1.0, 0.0, 0.0])
        await store.add([c])
        results = await store.search([1.0, 0.0, 0.0], limit=1)
        assert len(results) == 1
        assert results[0].content == "hello"

    @pytest.mark.asyncio
    async def test_persistence(self, tmp_path):
        path = str(tmp_path / "vec.db")
        store1 = SQLiteVectorStore(path)
        await store1.add([_chunk("persist", embedding=[1.0, 0.0])])

        store2 = SQLiteVectorStore(path)
        results = await store2.search([1.0, 0.0], limit=1)
        assert len(results) == 1
        assert results[0].content == "persist"

    @pytest.mark.asyncio
    async def test_delete(self, tmp_path):
        store = SQLiteVectorStore(str(tmp_path / "vec.db"))
        await store.add(
            [
                _chunk("a", doc_id="d1", embedding=[1.0, 0.0]),
                _chunk("b", doc_id="d2", embedding=[0.0, 1.0]),
            ]
        )
        await store.delete("d1")
        results = await store.search([1.0, 0.0], limit=10)
        assert len(results) == 1
        assert results[0].document_id == "d2"

    @pytest.mark.asyncio
    async def test_clear(self, tmp_path):
        store = SQLiteVectorStore(str(tmp_path / "vec.db"))
        await store.add([_chunk("a", embedding=[1.0])])
        await store.clear()
        results = await store.search([1.0], limit=10)
        assert results == []

    @pytest.mark.asyncio
    async def test_filter(self, tmp_path):
        store = SQLiteVectorStore(str(tmp_path / "vec.db"))
        c1 = _chunk("a", doc_id="d1", embedding=[1.0, 0.0], topic="science")
        c2 = _chunk("b", doc_id="d2", embedding=[1.0, 0.0], topic="art")
        await store.add([c1, c2])
        results = await store.search([1.0, 0.0], limit=5, filter={"topic": "science"})
        assert len(results) == 1
        assert results[0].metadata["topic"] == "science"


class _FakePostgresCursor:
    def __init__(self, rows):
        self._rows = rows
        self._fetchall_result = []

    def execute(self, query, params=None):
        normalized = " ".join(query.lower().split())
        if normalized.startswith("insert into"):
            chunk_id, document_id, content, metadata_json, embedding_json = params
            self._rows[chunk_id] = (
                chunk_id,
                document_id,
                content,
                metadata_json,
                embedding_json,
            )
            return
        if normalized.startswith("select id, document_id, content, metadata, embedding"):
            self._fetchall_result = list(self._rows.values())
            return
        if "where document_id = %s" in normalized:
            document_id = params[0]
            for chunk_id, row in list(self._rows.items()):
                if row[1] == document_id:
                    self._rows.pop(chunk_id, None)
            return
        if normalized.startswith("delete from"):
            self._rows.clear()
            return

    def fetchall(self):
        return self._fetchall_result

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


class _FakePostgresConnection:
    def __init__(self, rows):
        self._rows = rows
        self.commits = 0

    def cursor(self):
        return _FakePostgresCursor(self._rows)

    def commit(self):
        self.commits += 1

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


class TestPostgreSQLVectorStore:
    def test_constructor_resolves_env_dsn(self, monkeypatch):
        monkeypatch.setenv("DATABASE_URL", "postgresql://env/rag")
        store = PostgreSQLVectorStore()
        assert store.dsn == "postgresql://env/rag"

    @pytest.mark.asyncio
    async def test_add_search_delete_clear_with_fake_backend(self, monkeypatch):
        rows = {}
        connection = _FakePostgresConnection(rows)
        monkeypatch.setattr(
            "logos_adk.rag.vector_postgresql.pooled_connect",
            lambda dsn: connection,
        )

        store = PostgreSQLVectorStore("postgresql://user:pass@localhost/rag")
        await store.add(
            [
                _chunk("a", doc_id="d1", embedding=[1.0, 0.0], topic="science"),
                _chunk("b", doc_id="d2", embedding=[0.0, 1.0], topic="art"),
            ]
        )

        filtered = await store.search([1.0, 0.0], limit=5, filter={"topic": "science"})
        assert len(filtered) == 1
        assert filtered[0].content == "a"

        await store.delete("d1")
        remaining = await store.search([0.0, 1.0], limit=5)
        assert len(remaining) == 1
        assert remaining[0].document_id == "d2"

        await store.clear()
        assert await store.search([1.0, 0.0], limit=5) == []


class TestChromaVectorStore:
    def test_constructor_stores_collection_name(self):
        store = ChromaVectorStore(collection_name="test_docs")
        assert store._collection_name == "test_docs"
        assert store._client is None
