"""Tests for the dummy-node cross-batch reference mechanism.

The batching design: when a batch references a target that lives in another
batch, a *dummy* target node is created carrying the real target's id. On the
node side dummies are skipped (``node_generator`` drops ``label == 'dummy'``),
but the *edge* is still emitted with the real endpoint id. BioCypher appends
node/edge CSV part files after every batch, and a single ``neo4j-admin
import`` at the end stitches edges to nodes by id. So as long as the real
target lands in *some* batch, the edge resolves at import time.

These tests model that by emitting nodes/edges per batch and accumulating the
result, mimicking the appended CSVs that the final import consumes.
"""

import networkx as nx

import import_fhir_to_nx_diGraph as pipeline


def _emit_batch(graph):
    """Emit one batch the way run_biocypher does: nodes first, then edges."""
    nodes = {nid: label for (nid, label, _props) in pipeline.node_generator(graph)}
    edges = [
        (eid, src, tgt, label)
        for (eid, src, tgt, label, _attrs) in pipeline.edge_generator(graph)
    ]
    return nodes, edges


def _accumulate(batches):
    """Union of all batches' CSV part files, as the final import would see."""
    all_nodes = {}
    all_edges = []
    for graph in batches:
        nodes, edges = _emit_batch(graph)
        all_nodes.update(nodes)  # neo4j --skip-duplicate-nodes: union by id
        all_edges.extend(edges)
    return all_nodes, all_edges


def _batch_with_dummy_target():
    """Batch 1: real Observation O1 referencing P1, which is elsewhere.

    ``edge_label`` mirrors what ``process_references.py`` always sets on a
    real dummy node (the referenced resource type, e.g. "Patient").
    """
    g = nx.DiGraph()
    g.add_node("obs1", label="observation", unique_id="O1")
    g.add_node("dummy_p1", label="dummy", unique_id="P1", edge_label="Patient")
    g.add_edge("obs1", "dummy_p1", id="e1")
    return g


def _batch_with_real_target():
    """Batch 2: the real Patient node P1 that the dummy stood in for."""
    g = nx.DiGraph()
    g.add_node("pat1", label="patient", unique_id="P1")
    return g


def _unbatched_graph():
    """Everything in one graph: the reference target is the real node."""
    g = nx.DiGraph()
    g.add_node("obs1", label="observation", unique_id="O1")
    g.add_node("pat1", label="patient", unique_id="P1")
    g.add_edge("obs1", "pat1", id="e1")
    return g


def test_dummy_target_skipped_as_node_but_edge_survives():
    nodes, edges = _emit_batch(_batch_with_dummy_target())
    assert set(nodes) == {"O1"}            # dummy P1 not written as a node
    assert len(edges) == 1
    _, src, tgt, _ = edges[0]
    assert (src, tgt) == ("O1", "P1")      # edge keeps the real endpoint id


def test_cross_batch_edge_endpoint_supplied_by_later_batch():
    nodes, edges = _accumulate(
        [_batch_with_dummy_target(), _batch_with_real_target()]
    )
    # Real P1 arrives in batch 2; dummy was skipped. Every edge endpoint now
    # corresponds to a real node, so the import connects them.
    assert set(nodes) == {"O1", "P1"}
    endpoints = {s for _, s, _, _ in edges} | {t for _, _, t, _ in edges}
    assert endpoints <= set(nodes)


def test_batched_connectivity_matches_unbatched():
    batched_nodes, batched_edges = _accumulate(
        [_batch_with_dummy_target(), _batch_with_real_target()]
    )
    full_nodes, full_edges = _emit_batch(_unbatched_graph())

    assert set(batched_nodes) == set(full_nodes)
    # Connectivity (endpoint pairs) is identical, ignoring relationship type.
    assert {(s, t) for _, s, t, _ in batched_edges} == {
        (s, t) for _, s, t, _ in full_edges
    }


def test_within_batch_edge_has_correct_relationship_label():
    # When both endpoints are real in the same batch, the type is correct.
    (_, _, _, label, _), = pipeline.edge_generator(_unbatched_graph())
    assert label == "Observation_to_Patient"


def test_batched_and_unbatched_produce_identical_edges():
    # The full equivalence check: batching must not change the OUTCOME, and
    # that includes relationship *types*, not just connectivity. We compare
    # the complete edge set (endpoints + relationship label) of a batched run
    # against the unbatched run. Since the dummy node carries the real
    # target's resource type via `edge_label`, the relationship label is
    # derived correctly even while the target is still a dummy, so no
    # cross-batch '..._to_Dummy' relationship should ever appear.
    _, batched_edges = _accumulate(
        [_batch_with_dummy_target(), _batch_with_real_target()]
    )
    _, full_edges = _emit_batch(_unbatched_graph())

    def relationships(edges):
        return {(src, tgt, label) for _eid, src, tgt, label in edges}

    # No dummy-typed relationship should be produced at all.
    assert not any(
        label.endswith("_to_Dummy") for _e, _s, _t, label in batched_edges
    )
    # And the batched edge set must match the unbatched one exactly.
    assert relationships(batched_edges) == relationships(full_edges)