import pytest

from deerflow.skill_knowledge.graph_writer import (
    EDGE_QUERY,
    NODE_QUERY,
    execute_plan,
    parameterized_statements,
    plan_from_ir,
)


class RecordingTransaction:
    def __init__(self) -> None:
        self.calls: list[tuple[str, dict]] = []

    async def run(self, query: str, parameters: dict) -> None:
        self.calls.append((query, parameters))


def _plan():
    return plan_from_ir(
        {
            "entities": [
                {"local_id": "a", "type": "person", "name": "张三", "confidence": 1},
                {"local_id": "b", "type": "system", "name": "知识库", "confidence": 1},
            ],
            "relations": [
                {
                    "source_local_id": "a",
                    "target_local_id": "b",
                    "predicate": "uses",
                    "confidence": 0.9,
                    "evidence": [{"file": "SKILL.md", "line": 2}],
                }
            ],
        },
        target_type="weknora",
        target_id="kb-1",
        binding_id="binding-1",
        snapshot_id="snapshot-1",
        job_id="job-1",
    )


def test_direct_graph_plan_uses_only_fixed_parameterized_templates() -> None:
    statements = parameterized_statements(_plan())
    assert len(statements) == 3
    assert {row["statement"] for row in statements} == {NODE_QUERY, EDGE_QUERY}
    assert all("parameters" in row for row in statements)
    assert all("张三" not in row["statement"] for row in statements)


@pytest.mark.asyncio
async def test_direct_graph_writer_obeys_explicit_enable_switch() -> None:
    transaction = RecordingTransaction()
    disabled = await execute_plan(transaction, _plan(), enabled=False)
    assert disabled == {"nodes": 0, "edges": 0}
    assert transaction.calls == []

    enabled = await execute_plan(transaction, _plan(), enabled=True)
    assert enabled == {"nodes": 2, "edges": 1}
    assert len(transaction.calls) == 3