import unittest
from pathlib import Path
from unittest.mock import AsyncMock

from rag.mock_ingest import MOCK_DOCUMENTS, ingest_mock_documents
from rag.milvus_delete import delete_document_vectors


class MockIngestTests(unittest.IsolatedAsyncioTestCase):
    async def test_mock_documents_are_chunked_embedded_and_inserted(self):
        milvus = AsyncMock()
        embedder = AsyncMock()
        embedder.return_value = [[0.0] * 768]

        await ingest_mock_documents(milvus, embedder=embedder)

        self.assertEqual(len(MOCK_DOCUMENTS), 3)
        self.assertEqual(embedder.await_count, len(MOCK_DOCUMENTS))
        self.assertEqual(milvus.insert.await_count, len(MOCK_DOCUMENTS))
        inserted = [call.kwargs for call in milvus.insert.await_args_list]
        self.assertEqual(
            {item["collection_name"] for item in inserted},
            {item.collection for item in MOCK_DOCUMENTS},
        )
        for call_args in inserted:
            rows = call_args["data"]
            self.assertTrue(rows)
            self.assertTrue({"doc_id", "title", "section_title", "text", "strategy", "vector"} <= rows[0].keys())
            self.assertEqual(len(rows[0]["vector"]), 768)

    def test_mock_markdown_files_are_present(self):
        for item in MOCK_DOCUMENTS:
            self.assertTrue(Path(item.path).is_file())
            self.assertEqual(Path(item.path).suffix, ".md")
        faq = next(item for item in MOCK_DOCUMENTS if item.collection == "fin_faq")
        self.assertEqual(
            sum(line.startswith("Q:") for line in Path(faq.path).read_text(encoding="utf-8").splitlines()),
            40,
        )

    async def test_deletes_document_rows_by_doc_id_from_all_knowledge_collections(self):
        milvus = AsyncMock()

        await delete_document_vectors("doc-123", milvus_client=milvus)

        self.assertEqual(milvus.delete.await_count, 3)
        for call_args in milvus.delete.await_args_list:
            self.assertEqual(call_args.kwargs["filter"], 'doc_id == "doc-123"')


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