"""Unit tests for PGGraphStorage chunk-level transaction batch paths.

Covers the PR-2 chunk-level fallback: ``upsert_nodes_batch`` /
``upsert_edges_batch`` group per-row Cypher into payload/record-bounded chunks
run in a single transaction; ``remove_nodes`` runs all chunks in ONE
transaction (all-or-nothing); ``remove_edges`` runs one transaction per chunk.
"""

import pytest
from unittest.mock import AsyncMock

from lightrag.kg.postgres_impl import PGGraphStorage


# ---------------------------------------------------------------------------
# Capture harness
# ---------------------------------------------------------------------------


class _Capture:
    def __init__(self) -> None:
        self.calls: list[dict] = []  # every connection.execute(sql, *args)
        self.tx_count = 0  # connection.transaction() opens
        self.run_count = 0  # _run_with_retry invocations (= chunks issued)


class _FakeTransaction:
    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc, tb):
        return False


class _FakeConnection:
    def __init__(self, capture: _Capture):
        self._capture = capture

    def transaction(self):
        self._capture.tx_count += 1
        return _FakeTransaction()

    async def execute(self, sql, *args):
        self._capture.calls.append({"sql": sql, "args": args})
        return ""


def make_graph_storage(
    *,
    max_upsert_records: int | None = None,
    max_upsert_payload: int | None = None,
    max_delete_records: int | None = None,
) -> tuple[PGGraphStorage, _Capture]:
    storage = PGGraphStorage.__new__(PGGraphStorage)
    storage.workspace = "test_ws"
    storage.namespace = "test_graph"
    storage.graph_name = "test_graph"
    storage.__post_init__()  # resolves the chunk-level batch limits
    if max_upsert_records is not None:
        storage._max_upsert_records_per_batch = max_upsert_records
    if max_upsert_payload is not None:
        storage._max_upsert_payload_bytes = max_upsert_payload
    if max_delete_records is not None:
        storage._max_delete_records_per_batch = max_delete_records

    capture = _Capture()
    conn = _FakeConnection(capture)

    async def fake_run_with_retry(operation, **_kwargs):
        capture.run_count += 1
        return await operation(conn)

    storage.db = AsyncMock()
    storage.db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
    return storage, capture


def _sql_count(capture: _Capture, needle: str) -> int:
    return sum(1 for c in capture.calls if needle in c["sql"])


# ---------------------------------------------------------------------------
# upsert_nodes_batch
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_upsert_nodes_batch_splits_into_chunk_transactions():
    storage, cap = make_graph_storage(max_upsert_records=2)
    nodes = [(f"n{i}", {"entity_id": f"n{i}", "v": str(i)}) for i in range(5)]

    await storage.upsert_nodes_batch(nodes)

    # 5 nodes / cap 2 => 3 chunks, each its own _run_with_retry + transaction.
    assert cap.run_count == 3
    assert cap.tx_count == 3
    # One MERGE per node.
    assert _sql_count(cap, "MERGE (n:base") == 5


@pytest.mark.asyncio
async def test_upsert_nodes_batch_dedupes_last_write_wins():
    storage, cap = make_graph_storage()
    await storage.upsert_nodes_batch(
        [
            ("A", {"entity_id": "A", "weight": "first"}),
            ("A", {"entity_id": "A", "weight": "second"}),
        ]
    )

    merge_calls = [c for c in cap.calls if "MERGE (n:base" in c["sql"]]
    assert len(merge_calls) == 1
    assert '"second"' in merge_calls[0]["sql"]
    assert '"first"' not in merge_calls[0]["sql"]


@pytest.mark.asyncio
async def test_upsert_nodes_batch_splits_by_payload_bytes():
    # Disable the record cap, force a small byte budget so each big node lands
    # in its own chunk.
    storage, cap = make_graph_storage(max_upsert_records=0, max_upsert_payload=200)
    nodes = [(f"n{i}", {"entity_id": f"n{i}", "blob": "X" * 150}) for i in range(3)]

    await storage.upsert_nodes_batch(nodes)

    assert cap.run_count == 3


@pytest.mark.asyncio
async def test_upsert_nodes_batch_missing_entity_id_raises_before_tx():
    storage, cap = make_graph_storage()
    with pytest.raises(ValueError, match="entity_id"):
        await storage.upsert_nodes_batch([("A", {"foo": "bar"})])
    # The builder runs before the transaction opens, so nothing was issued.
    assert cap.run_count == 0


@pytest.mark.asyncio
async def test_upsert_nodes_batch_empty_noop():
    storage, cap = make_graph_storage()
    await storage.upsert_nodes_batch([])
    assert cap.run_count == 0


# ---------------------------------------------------------------------------
# upsert_edges_batch
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_upsert_edges_batch_splits_into_chunk_transactions():
    storage, cap = make_graph_storage(max_upsert_records=2)
    edges = [(f"s{i}", f"t{i}", {"w": str(i)}) for i in range(5)]

    await storage.upsert_edges_batch(edges)

    # 5 edges / cap 2 => 3 chunks; each chunk = one transaction.
    assert cap.run_count == 3
    assert cap.tx_count == 3
    # One CREATE cypher per edge, and exactly one graph-wide lock per chunk.
    assert _sql_count(cap, "CREATE (source)-[r:DIRECTED") == 5
    assert _sql_count(cap, "pg_advisory_xact_lock") == 3


@pytest.mark.asyncio
async def test_upsert_edges_batch_takes_one_graph_lock_per_chunk():
    """The batch chunk takes a single graph-wide advisory lock (keyed on
    graph_name only), not one lock per edge -- bounded regardless of edge count."""
    storage, cap = make_graph_storage()
    await storage.upsert_edges_batch(
        [(f"s{i}", f"t{i}", {"w": str(i)}) for i in range(4)]
    )

    lock_calls = [c for c in cap.calls if "pg_advisory_xact_lock" in c["sql"]]
    assert len(lock_calls) == 1  # one chunk, one lock (not 4)
    assert lock_calls[0]["args"] == ("test_graph",)  # graph-keyed, no endpoints
    assert _sql_count(cap, "CREATE (source)-[r:DIRECTED") == 4
    # The lock is acquired before any edge cypher in the chunk.
    assert "pg_advisory_xact_lock" in cap.calls[0]["sql"]


@pytest.mark.asyncio
async def test_upsert_edges_batch_empty_noop():
    storage, cap = make_graph_storage()
    await storage.upsert_edges_batch([])
    assert cap.run_count == 0


# ---------------------------------------------------------------------------
# remove_nodes — all chunks in ONE transaction (all-or-nothing)
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_remove_nodes_chunks_share_one_transaction():
    storage, cap = make_graph_storage(max_delete_records=2)
    await storage.remove_nodes([f"n{i}" for i in range(5)])

    # One transaction wraps all chunks (preserves single-statement atomicity).
    assert cap.run_count == 1
    assert cap.tx_count == 1
    # 5 ids / cap 2 => 3 bounded IN [...] DETACH DELETE statements.
    assert _sql_count(cap, "DETACH DELETE n") == 3


@pytest.mark.asyncio
async def test_remove_nodes_empty_noop():
    storage, cap = make_graph_storage()
    await storage.remove_nodes([])
    assert cap.run_count == 0


# ---------------------------------------------------------------------------
# remove_edges — one transaction per chunk
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_remove_edges_one_transaction_per_chunk():
    storage, cap = make_graph_storage(max_delete_records=2)
    edges = [(f"s{i}", f"t{i}") for i in range(5)]

    await storage.remove_edges(edges)

    # 5 edges / cap 2 => 3 chunks, each its own transaction.
    assert cap.run_count == 3
    assert cap.tx_count == 3
    # One DELETE r statement per edge.
    assert _sql_count(cap, "DELETE r") == 5


@pytest.mark.asyncio
async def test_remove_edges_empty_noop():
    storage, cap = make_graph_storage()
    await storage.remove_edges([])
    assert cap.run_count == 0