"""graph.py 单元测试 — PageRank + 图扩展 + repo-map"""

import pytest

from index1.db import Database
from index1.graph import (
    compute_and_store_pagerank,
    expand_context,
    generate_repo_map,
    personalized_pagerank,
)


class TestPersonalizedPageRank:
    def test_empty_graph(self):
        result = personalized_pagerank({})
        assert result == {}

    def test_single_node_no_edges(self):
        result = personalized_pagerank({1: {}})
        assert len(result) == 1
        assert abs(result[1] - 1.0) < 0.01

    def test_two_nodes(self):
        adj = {1: {2: 1.0}, 2: {1: 1.0}}
        result = personalized_pagerank(adj)
        assert len(result) == 2
        assert abs(result[1] - result[2]) < 0.01  # 对称图 → 等rank
        assert abs(sum(result.values()) - 1.0) < 0.01  # 归一化

    def test_star_graph(self):
        """星型图: hub 节点 rank 最高"""
        adj = {
            1: {2: 1.0, 3: 1.0, 4: 1.0},  # hub
            2: {1: 1.0},
            3: {1: 1.0},
            4: {1: 1.0},
        }
        result = personalized_pagerank(adj)
        assert result[1] > result[2]
        assert result[1] > result[3]

    def test_personalization(self):
        """personalization 偏置应影响 rank 分布"""
        adj = {1: {2: 1.0}, 2: {3: 1.0}, 3: {1: 1.0}}
        # 无偏置
        uniform = personalized_pagerank(adj)
        # 偏置 node 1
        biased = personalized_pagerank(adj, personalization={1: 1.0})
        # node 1 在偏置模式下应有更高 rank
        assert biased[1] >= uniform[1] - 0.01

    def test_disconnected_graph(self):
        """断开图应正常处理"""
        adj = {1: {2: 1.0}, 3: {4: 1.0}}
        result = personalized_pagerank(adj)
        assert len(result) == 4
        assert abs(sum(result.values()) - 1.0) < 0.01

    def test_convergence(self):
        """PageRank 应在有限次迭代内收敛"""
        adj = {i: {(i + 1) % 10: 1.0} for i in range(10)}
        result = personalized_pagerank(adj, max_iter=100, tol=1e-8)
        # 环形图 → 所有节点等 rank
        values = list(result.values())
        assert max(values) - min(values) < 0.01

    def test_weighted_edges(self):
        """加权边应影响 rank 分布"""
        adj = {
            1: {2: 10.0, 3: 1.0},  # node 2 权重更大
            2: {1: 1.0},
            3: {1: 1.0},
        }
        result = personalized_pagerank(adj)
        assert result[2] > result[3]  # 2 应该比 3 rank 高


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

    def test_empty_graph(self, tmp_path):
        db = self._make_db(tmp_path)
        count = compute_and_store_pagerank(db, "test")
        assert count == 0
        db.close()

    def test_with_data(self, tmp_path):
        db = self._make_db(tmp_path)

        # 创建文档
        doc1 = db.upsert_document("a.py", "h1", "test", "python")
        doc2 = db.upsert_document("b.py", "h2", "test", "python")

        # 创建符号
        db.insert_symbols(doc1, [{"name": "foo", "kind": "function", "line_start": 0}])
        db.insert_symbols(doc2, [{"name": "bar", "kind": "function", "line_start": 0}])

        # 创建引用（a.py → b.py）
        db.insert_refs([{
            "from_doc_id": doc1,
            "to_doc_id": doc2,
            "kind": "call",
            "confidence": 0.9,
        }])

        count = compute_and_store_pagerank(db, "test")
        assert count >= 1
        db.close()


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

    def test_empty_seeds(self, tmp_path):
        db = self._make_db(tmp_path)
        result = expand_context(db, [])
        assert result == []
        db.close()

    def test_no_symbols(self, tmp_path):
        db = self._make_db(tmp_path)
        result = expand_context(db, [999])  # nonexistent chunk
        assert result == []
        db.close()


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

    def test_empty_collection(self, tmp_path):
        db = self._make_db(tmp_path)
        result = generate_repo_map(db, "nonexistent")
        assert result == ""
        db.close()

    def test_with_symbols(self, tmp_path):
        db = self._make_db(tmp_path)

        doc_id = db.upsert_document("src/math.py", "h1", "test", "python")
        db.insert_symbols(doc_id, [
            {"name": "add", "kind": "function", "line_start": 1},
            {"name": "subtract", "kind": "function", "line_start": 5},
        ])

        result = generate_repo_map(db, "test", budget=2048)
        assert "src/math.py" in result
        assert "add" in result
        assert "subtract" in result
        db.close()

    def test_call_annotations(self, tmp_path):
        """repo-map 应标注跨文件调用关系 [→ target]"""
        db = self._make_db(tmp_path)

        doc1 = db.upsert_document("src/search.py", "h1", "test", "python")
        doc2 = db.upsert_document("src/db.py", "h2", "test", "python")

        # search.py 有一个函数 search
        ids1 = db.insert_symbols(doc1, [
            {"name": "search", "kind": "function", "line_start": 1},
        ])
        # db.py 有一个函数 connect
        ids2 = db.insert_symbols(doc2, [
            {"name": "connect", "kind": "function", "line_start": 1},
        ])

        # search.py:search() 调用 db.py:connect()（跨文件引用）
        db.insert_refs([{
            "from_doc_id": doc1,
            "to_doc_id": doc2,
            "to_symbol_id": ids2[0],
            "ref_name": "connect",
            "kind": "call",
            "from_line": 5,
            "confidence": 0.9,
        }])

        result = generate_repo_map(db, "test", budget=2048)
        # search 函数应该标注 [→ db.connect]（带文件前缀）
        assert "→ db.connect" in result
        db.close()

    def test_call_annotations_dedup(self, tmp_path):
        """同名目标应去重"""
        db = self._make_db(tmp_path)
        doc1 = db.upsert_document("src/main.py", "h1", "test", "python")
        doc2 = db.upsert_document("src/db.py", "h2", "test", "python")

        db.insert_symbols(doc1, [
            {"name": "main", "kind": "function", "line_start": 1},
        ])
        ids2 = db.insert_symbols(doc2, [
            {"name": "connect", "kind": "function", "line_start": 1},
        ])

        # 同一个函数两次调用同名目标（不同行）
        db.insert_refs([
            {"from_doc_id": doc1, "to_doc_id": doc2, "to_symbol_id": ids2[0],
             "ref_name": "connect", "kind": "call", "from_line": 5, "confidence": 0.9},
            {"from_doc_id": doc1, "to_doc_id": doc2, "to_symbol_id": ids2[0],
             "ref_name": "connect", "kind": "call", "from_line": 8, "confidence": 0.9},
        ])

        result = generate_repo_map(db, "test", budget=2048)
        # 应该只出现一次 db.connect（去重）
        assert result.count("db.connect") == 1
        db.close()

    def test_call_annotations_boundary(self, tmp_path):
        """调用应精确归属到所属函数（相邻函数边界）"""
        db = self._make_db(tmp_path)
        doc1 = db.upsert_document("src/app.py", "h1", "test", "python")
        doc2 = db.upsert_document("src/util.py", "h2", "test", "python")

        # app.py 有两个函数: foo(line 1) 和 bar(line 20)
        db.insert_symbols(doc1, [
            {"name": "foo", "kind": "function", "line_start": 1},
            {"name": "bar", "kind": "function", "line_start": 20},
        ])
        ids2 = db.insert_symbols(doc2, [
            {"name": "helper_a", "kind": "function", "line_start": 1},
            {"name": "helper_b", "kind": "function", "line_start": 10},
        ])

        # foo() 内调用 helper_a (from_line=5, < bar 的 line_start=20)
        # bar() 内调用 helper_b (from_line=25, >= bar 的 line_start=20)
        db.insert_refs([
            {"from_doc_id": doc1, "to_doc_id": doc2, "to_symbol_id": ids2[0],
             "ref_name": "helper_a", "kind": "call", "from_line": 5, "confidence": 0.9},
            {"from_doc_id": doc1, "to_doc_id": doc2, "to_symbol_id": ids2[1],
             "ref_name": "helper_b", "kind": "call", "from_line": 25, "confidence": 0.9},
        ])

        result = generate_repo_map(db, "test", budget=2048)
        # foo 应标注 helper_a，bar 应标注 helper_b
        result_lines = result.split("\n")
        foo_line = [l for l in result_lines if "foo" in l][0]
        bar_line = [l for l in result_lines if "bar" in l][0]
        assert "helper_a" in foo_line
        assert "helper_b" not in foo_line
        assert "helper_b" in bar_line
        assert "helper_a" not in bar_line
        db.close()

    def test_budget_limit(self, tmp_path):
        db = self._make_db(tmp_path)

        # 创建多个文档和符号
        for i in range(20):
            doc_id = db.upsert_document(f"src/mod{i}.py", f"h{i}", "test", "python")
            db.insert_symbols(doc_id, [
                {"name": f"func_{i}_{j}", "kind": "function", "line_start": j * 10}
                for j in range(10)
            ])

        # 小预算应截断
        result = generate_repo_map(db, "test", budget=100)
        token_est = len(result) // 4
        assert token_est <= 150  # 允许一定超出（按行截断非精确）
        db.close()


class TestDbSymbolsCrud:
    """db.py symbols/refs CRUD 测试"""

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

    def test_insert_and_get_symbols(self, tmp_path):
        db = self._make_db(tmp_path)
        doc_id = db.upsert_document("a.py", "h1", "test", "python")
        ids = db.insert_symbols(doc_id, [
            {"name": "foo", "kind": "function", "line_start": 1},
            {"name": "Bar", "kind": "class", "line_start": 10},
        ])
        assert len(ids) == 2
        syms = db.get_symbols_by_doc(doc_id)
        assert len(syms) == 2
        db.close()

    def test_get_symbol_by_name(self, tmp_path):
        db = self._make_db(tmp_path)
        doc_id = db.upsert_document("a.py", "h1", "test", "python")
        db.insert_symbols(doc_id, [
            {"name": "unique_func", "kind": "function", "line_start": 1},
        ])
        results = db.get_symbol_by_name("unique_func", "test")
        assert len(results) == 1
        assert results[0]["name"] == "unique_func"
        db.close()

    def test_cascade_delete(self, tmp_path):
        """删除文档应级联删除 symbols 和 refs"""
        db = self._make_db(tmp_path)
        doc_id = db.upsert_document("a.py", "h1", "test", "python")
        db.insert_symbols(doc_id, [
            {"name": "foo", "kind": "function", "line_start": 1},
        ])
        db.insert_refs([{
            "from_doc_id": doc_id,
            "kind": "call",
            "confidence": 0.9,
        }])

        # 删除文档
        db.delete_document("a.py")

        # 验证级联
        syms = db.get_symbols_by_doc(doc_id)
        assert len(syms) == 0
        db.close()

    def test_insert_refs(self, tmp_path):
        db = self._make_db(tmp_path)
        doc1 = db.upsert_document("a.py", "h1", "test", "python")
        doc2 = db.upsert_document("b.py", "h2", "test", "python")
        ids1 = db.insert_symbols(doc1, [
            {"name": "caller", "kind": "function", "line_start": 1},
        ])
        ids2 = db.insert_symbols(doc2, [
            {"name": "callee", "kind": "function", "line_start": 1},
        ])

        count = db.insert_refs([{
            "from_doc_id": doc1,
            "to_doc_id": doc2,
            "from_symbol_id": ids1[0],
            "to_symbol_id": ids2[0],
            "kind": "call",
            "confidence": 0.9,
        }])
        assert count == 1

        refs = db.get_refs_by_symbol(ids2[0], direction="callers")
        assert len(refs) == 1
        assert refs[0]["direction"] == "caller"
        db.close()

    def test_get_file_graph(self, tmp_path):
        db = self._make_db(tmp_path)
        doc1 = db.upsert_document("a.py", "h1", "test", "python")
        doc2 = db.upsert_document("b.py", "h2", "test", "python")
        db.insert_refs([{
            "from_doc_id": doc1,
            "to_doc_id": doc2,
            "kind": "call",
            "confidence": 0.8,
        }])

        edges = db.get_file_graph("test")
        assert len(edges) >= 1
        from_id, to_id, weight = edges[0]
        assert from_id == doc1
        assert to_id == doc2
        assert weight > 0
        db.close()

    def test_symbol_stats(self, tmp_path):
        db = self._make_db(tmp_path)
        doc_id = db.upsert_document("a.py", "h1", "test", "python")
        db.insert_symbols(doc_id, [
            {"name": "foo", "kind": "function", "line_start": 1},
        ])

        stats = db.get_symbol_stats("test")
        assert stats["symbol_count"] >= 1
        db.close()

    def test_refs_with_ref_name(self, tmp_path):
        """refs 表应支持 ref_name 列"""
        db = self._make_db(tmp_path)
        doc_id = db.upsert_document("a.py", "h1", "test", "python")
        count = db.insert_refs([{
            "from_doc_id": doc_id,
            "ref_name": "some_function",
            "kind": "call",
            "confidence": 0.5,
        }])
        assert count == 1

        # 验证 ref_name 存储
        row = db.conn.execute(
            "SELECT ref_name FROM refs WHERE from_doc_id = ?", (doc_id,)
        ).fetchone()
        assert row[0] == "some_function"
        db.close()
