"""db.py 单元测试"""

from index1.db import Database, _floats_to_blob, _blob_to_floats


class TestDatabase:
    def _make_db(self, tmp_path, dim=4):
        db = Database(tmp_path / "test.db", embedding_dim=dim)
        db.connect()
        return db

    def test_init_schema_idempotent(self, tmp_path):
        """建表可重复调用"""
        db = self._make_db(tmp_path)
        db._init_schema()  # 再调一次
        db.close()

    def test_vec_loaded(self, tmp_path):
        db = self._make_db(tmp_path)
        assert db._vec_loaded is True
        db.close()

    def test_upsert_document(self, tmp_path):
        db = self._make_db(tmp_path)
        doc_id = db.upsert_document("docs/api.md", "abc123", "default", "markdown")
        assert doc_id > 0

        # 同路径再 upsert → 同 doc_id
        doc_id2 = db.upsert_document("docs/api.md", "def456", "default", "markdown")
        assert doc_id2 == doc_id

        # hash 已更新
        assert db.get_document_hash("docs/api.md") == "def456"
        db.close()

    def test_get_document_hash_missing(self, tmp_path):
        db = self._make_db(tmp_path)
        assert db.get_document_hash("nonexistent.md") is None
        db.close()

    def test_delete_document(self, tmp_path):
        db = self._make_db(tmp_path)
        doc_id = db.upsert_document("test.md", "hash1", "default", "markdown")
        db.insert_chunks(doc_id, [
            {"content": "hello", "chunk_index": 0, "tags": ["test"]},
        ])
        db.delete_document("test.md")
        assert db.get_document_hash("test.md") is None
        db.close()

    def test_insert_and_get_chunk(self, tmp_path):
        db = self._make_db(tmp_path)
        doc_id = db.upsert_document("api.md", "h1", "default", "markdown")
        db.insert_chunks(doc_id, [
            {
                "content": "Liquidation triggers when health < 1.2",
                "section": "## Liquidation",
                "parent_section": "# Risk",
                "chunk_index": 0,
                "tags": ["liquidation", "risk"],
                "token_count": 8,
                "summary": "Liquidation triggers...",
            },
        ])

        chunk = db.get_chunk(1)
        assert chunk is not None
        assert chunk.content == "Liquidation triggers when health < 1.2"
        assert chunk.section == "## Liquidation"
        assert chunk.parent_section == "# Risk"
        assert chunk.source == "api.md"
        assert chunk.tags == "liquidation, risk"
        db.close()

    def test_insert_chunk_with_embedding(self, tmp_path):
        db = self._make_db(tmp_path)
        doc_id = db.upsert_document("test.md", "h1", "default", "markdown")
        db.insert_chunks(doc_id, [
            {
                "content": "test content",
                "chunk_index": 0,
                "tags": [],
                "embedding": [1.0, 2.0, 3.0, 4.0],
            },
        ])
        # 向量搜索应能找到
        results = db.vector_search([1.0, 2.0, 3.0, 4.0], None, 5)
        assert len(results) == 1
        assert results[0][0] == 1  # chunk_id
        assert results[0][1] == 0.0  # distance = 0（完全匹配）
        db.close()

    def test_delete_chunks_by_doc(self, tmp_path):
        db = self._make_db(tmp_path)
        doc_id = db.upsert_document("test.md", "h1", "default", "markdown")
        db.insert_chunks(doc_id, [
            {"content": "chunk 1", "chunk_index": 0, "tags": ["a"]},
            {"content": "chunk 2", "chunk_index": 1, "tags": ["b"]},
        ])
        assert db.get_chunk(1) is not None
        assert db.get_chunk(2) is not None

        db.delete_chunks_by_doc(doc_id)
        assert db.get_chunk(1) is None
        assert db.get_chunk(2) is None
        db.close()

    def test_fts5_search(self, tmp_path):
        db = self._make_db(tmp_path)
        doc_id = db.upsert_document("api.md", "h1", "default", "markdown")
        db.insert_chunks(doc_id, [
            {"content": "liquidation triggers when health factor drops", "chunk_index": 0, "tags": ["risk"]},
            {"content": "margin call notification system", "chunk_index": 1, "tags": ["margin"]},
        ])

        results = db.fts5_search("liquidation", None, 5)
        assert len(results) >= 1
        assert results[0][1] > 0  # BM25 分数

    def test_fts5_search_with_collection(self, tmp_path):
        db = self._make_db(tmp_path)
        db.upsert_document("a.md", "h1", "project-a", "markdown")
        db.insert_chunks(1, [{"content": "alpha content", "chunk_index": 0, "tags": []}])

        db.upsert_document("b.md", "h2", "project-b", "markdown")
        db.insert_chunks(2, [{"content": "alpha different", "chunk_index": 0, "tags": []}])

        results_a = db.fts5_search("alpha", "project-a", 5)
        results_b = db.fts5_search("alpha", "project-b", 5)

        assert len(results_a) == 1
        assert len(results_b) == 1

    def test_vector_search_no_vec(self, tmp_path):
        """vec 未加载时返回空"""
        db = self._make_db(tmp_path)
        db._vec_loaded = False
        results = db.vector_search([1.0, 2.0, 3.0, 4.0], None, 5)
        assert results == []
        db.close()

    def test_get_chunk_context(self, tmp_path):
        db = self._make_db(tmp_path)
        doc_id = db.upsert_document("test.md", "h1", "default", "markdown")
        db.insert_chunks(doc_id, [
            {"content": "intro", "section": "# Intro", "chunk_index": 0, "tags": []},
            {"content": "middle", "section": "## Middle", "chunk_index": 1, "tags": []},
            {"content": "end", "section": "## End", "chunk_index": 2, "tags": []},
        ])

        ctx = db.get_chunk_context(2)  # middle chunk
        assert ctx["prev_section"] == "# Intro"
        assert ctx["next_section"] == "## End"
        db.close()

    def test_get_stats(self, tmp_path):
        db = self._make_db(tmp_path)
        doc_id = db.upsert_document("test.md", "h1", "default", "markdown")
        db.insert_chunks(doc_id, [
            {"content": "chunk", "chunk_index": 0, "tags": []},
        ])

        stats = db.get_stats()
        assert stats["doc_count"] == 1
        assert stats["chunk_count"] == 1
        assert stats["collections"] == ["default"]
        assert stats["vec_loaded"] is True
        db.close()


class TestBlobConversion:
    def test_roundtrip(self):
        original = [1.0, 2.5, 3.7, 4.0]
        blob = _floats_to_blob(original)
        restored = _blob_to_floats(blob, 4)
        for a, b in zip(original, restored):
            assert abs(a - b) < 1e-6
