module ProtocolZoo

using QuantumSavory
import QuantumSavory: get_time_tracker, Tag, isolderthan, onchange, QueryOnRegResult
import QuantumSavory: _message_type
using QuantumSavory: Wildcard, alwaystrue, compactstr, @domain
using QuantumSavory.CircuitZoo: EntanglementSwap, LocalEntanglementSwap

using DocStringExtensions

using Distributions: Geometric
using ConcurrentSim: Simulation, @yield, timeout, @process, now
import ConcurrentSim: Process
import ResumableFunctions
using ResumableFunctions: @resumable
import SumTypes
using PrettyTables: PrettyTables, pretty_table

export
    # protocols
    EntanglerProt, SwapperProt, EntanglementTracker, EntanglementConsumer, CutoffProt,
    protocol_log_context,
    # tags
    EntanglementCounterpart, EntanglementHistory, EntanglementUpdateX, EntanglementUpdateZ,
    EntanglementID, NO_ENTANGLEMENT_ID, fresh_entanglement_id, combine_entanglement_ids,
    # from Switches
    SimpleSwitchDiscreteProt, SwitchRequest,
    # from QTCP
    QDatagram, Flow, LinkLevelRequest,
    QTCPPairBegin, QTCPPairEnd,
    LinkLevelReply, LinkLevelReplyAtHop, LinkLevelReplyAtSource,
    NetworkNodeController, EndNodeController, LinkController

public AbstractProtocol, available_protocol_types, protocol_catalog_metadata,
    permits_virtual_edge

abstract type AbstractProtocol end

"""
Check whether a protocol type or instance permits virtual edges between nodes.

Virtual edges refer to protocol connections between two nodes that do not correspond
to actual network edges/links. Some protocols like [`EntanglementConsumer`](@ref) can operate
between any two nodes in the network regardless of physical connectivity.

The capability is defined on protocol types; instance queries delegate to their type.
"""
permits_virtual_edge(::Type{<:AbstractProtocol}) = false
permits_virtual_edge(prot::AbstractProtocol) = permits_virtual_edge(typeof(prot))

get_time_tracker(prot::AbstractProtocol) = prot.sim::Simulation

Process(prot::AbstractProtocol, args...; kwargs...) = Process((e,a...;k...)->prot(a...;k...), get_time_tracker(prot), args...; kwargs...)

_protocol_nodes(::AbstractProtocol) = ()

"""
    protocol_log_context(prot::AbstractProtocol)

Return the structured logging context for `prot`. The fields are the simulation
context followed by the concrete protocol type name and an immutable, ordered
snapshot of the participating node identifiers.

Custom protocols should overload this function when their node layout is not
represented by a built-in protocol family:

```julia
import QuantumSavory.ProtocolZoo: protocol_log_context

protocol_log_context(prot::MyProtocol) = (
    simulation_log_context(prot.sim)...,
    protocol = :MyProtocol,
    nodes = (prot.node,),
)
```
"""
function protocol_log_context(prot::AbstractProtocol)
    return (;
        simulation_log_context(get_time_tracker(prot))...,
        protocol=nameof(typeof(prot)),
        nodes=_protocol_nodes(prot),
    )
end

"""
    protocol_catalog_metadata(::Type{<:AbstractProtocol})

Opt a protocol type into [`available_protocol_types`](@ref), to make it available to tools like the GUI WebQuantumSavory.

Independent packages extend this method for their own public protocol types and return
a named tuple with exactly these fields:

- `attachment`: one of `:network`, `:node`, or `:edge`;
- `attachment_fields`: respectively `NamedTuple()`, `(node=:field,)`, or
  `(node_a=:field_a, node_b=:field_b)`; and
- `required_fields`: a tuple of configurable constructor-field names which callers
  must supply.

The protocol must have documented `sim` and `net` fields. Every other non-private
field must also be documented; attachment fields are supplied by topology and the
remaining fields are configurable parameters. Underscore-prefixed fields are private.
This generic intentionally has no fallback method so that dispatch is the opt-in.
"""
function protocol_catalog_metadata end

"""Return metadata for available protocol types.

Used to make a protocol available to tools like the GUI WebQuantumSavory.

The result is sorted by qualified type name. Each entry contains `type`, `doc`,
`nodeargs`, `attachment`, `attachment_fields`, `parameters`, and
`permits_virtual_edge`. A parameter entry contains `field`, `type`, `doc`, and
`required`. The `InteractiveUtils` and `REPL` standard libraries must be loaded to
activate this optional method."""
function available_protocol_types end

const QueryArgs = Union{Int,Function,Wildcard}

include("entanglement_ids.jl")


"""
$TYPEDEF

Indicates the current entanglement status with a remote node's slot. Added when a new entanglement is generated through [`EntanglerProt`](@ref) or when a swap happens and
 the [`EntanglementTracker`](@ref) receives an [`EntanglementUpdate`] message.

$TYPEDFIELDS
"""
@kwdef struct EntanglementCounterpart <: AbstractTag
    "the id of the remote node to which we are entangled"
    remote_node::Int
    "the slot in the remote node containing the qubit we are entangled to"
    remote_slot::Int
    "the identifier of the entangled pair"
    pair_id::EntanglementID
end
EntanglementCounterpart(remote_node::Int, remote_slot::Int) = EntanglementCounterpart(remote_node, remote_slot, NO_ENTANGLEMENT_ID)
Base.show(io::IO, tag::EntanglementCounterpart) = print(io, "Entangled to $(tag.remote_node).$(tag.remote_slot) with id $(tag.pair_id)")
Tag(tag::EntanglementCounterpart) = Tag(EntanglementCounterpart, tag.remote_node, tag.remote_slot, tag.pair_id)

function _tag_entanglement_counterpart!(
    slot, remote_node, remote_slot, pair_id, prot::AbstractProtocol
)
    existing = query(slot, EntanglementCounterpart, ❓, ❓, ❓)
    if !isnothing(existing)
        @error(
            "Counterpart tag conflicts with an existing tag",
            _group=LOG_GROUPS.protocol,
            event=:counterpart_tag_conflict,
            protocol_log_context(prot)...,
            slot=slot.idx,
            dst_node=remote_node,
            dst_slot=remote_slot,
            pair_id=pair_id,
        )
    end
    tag!(slot, EntanglementCounterpart, remote_node, remote_slot, pair_id)
end

"""
$TYPEDEF

This tag is used to store the outdated entanglement information after a
swap. It helps to direct incoming entanglement update messages to the right node after a swap.
It helps in situations when locally we have performed a swap, but we are now receiving a message
from a distant node that does not know yet that the swap has occurred (thus the distant node might
have outdated information about who is entangled to whom and we need to update that information).

$TYPEDFIELDS
"""
@kwdef struct EntanglementHistory <: AbstractTag
    "the id of the remote node we used to be entangled to"
    remote_node::Int
    "the slot of the remote node we used to be entangled to"
    remote_slot::Int
    "the id of remote node to which we are entangled after the swap"
    swap_remote_node::Int
    "the slot of the remote node to which we are entangled after the swap"
    swap_remote_slot::Int
    "the slot in this register with whom we performed a swap"
    swapped_local::Int
    "the pair-id chunk corresponding to `remote_node.remote_slot`"
    local_chunk_id::EntanglementID
    "the pair-id chunk corresponding to `swap_remote_node.swap_remote_slot`"
    swapped_chunk_id::EntanglementID
end
EntanglementHistory(remote_node::Int, remote_slot::Int, swap_remote_node::Int, swap_remote_slot::Int, swapped_local::Int) = EntanglementHistory(remote_node, remote_slot, swap_remote_node, swap_remote_slot, swapped_local, NO_ENTANGLEMENT_ID, NO_ENTANGLEMENT_ID)
Base.show(io::IO, tag::EntanglementHistory) = print(io, "Was entangled to $(tag.remote_node).$(tag.remote_slot) with chunk id $(tag.local_chunk_id), but swapped with .$(tag.swapped_local) which was entangled to $(tag.swap_remote_node).$(tag.swap_remote_slot) with chunk id $(tag.swapped_chunk_id)")
Tag(tag::EntanglementHistory) = Tag(EntanglementHistory, tag.remote_node, tag.remote_slot, tag.swap_remote_node, tag.swap_remote_slot, tag.swapped_local, tag.local_chunk_id, tag.swapped_chunk_id)

"""
$TYPEDEF

This tag arrives as a message from a remote node to which the current node was entangled to update the
entanglement information and apply an `X` correction after the remote node performs an entanglement swap.

$TYPEDFIELDS
"""
@kwdef struct EntanglementUpdateX <: AbstractTag
    "the pair id currently known by the receiver for the target slot"
    target_pair_id::EntanglementID
    "the pair-id chunk to combine into the target pair"
    other_pair_id::EntanglementID
    "the id of the node to which you were entangled before the swap"
    past_local_node::Int
    "the slot of the node to which you were entangled before the swap"
    past_local_slot::Int
    "the slot of your node that we were entangled to"
    past_remote_slot::Int
    "the id of the node to which you are now entangled after the swap"
    new_remote_node::Int
    "the slot of the node to which you are now entangled after the swap"
    new_remote_slot::Int
    "what Pauli correction you need to perform"
    correction::Int
end
EntanglementUpdateX(past_local_node::Int, past_local_slot::Int, past_remote_slot::Int, new_remote_node::Int, new_remote_slot::Int, correction::Int) = EntanglementUpdateX(NO_ENTANGLEMENT_ID, NO_ENTANGLEMENT_ID, past_local_node, past_local_slot, past_remote_slot, new_remote_node, new_remote_slot, correction)
Base.show(io::IO, tag::EntanglementUpdateX) = print(io, "Update pair $(tag.target_pair_id) at slot .$(tag.past_remote_slot) which used to be entangled to $(tag.past_local_node).$(tag.past_local_slot) to be entangled to $(tag.new_remote_node).$(tag.new_remote_slot), combine with $(tag.other_pair_id), and apply correction Z$(tag.correction)")
Tag(tag::EntanglementUpdateX) = Tag(EntanglementUpdateX, tag.target_pair_id, tag.other_pair_id, tag.past_local_node, tag.past_local_slot, tag.past_remote_slot, tag.new_remote_node, tag.new_remote_slot, tag.correction)

"""
$TYPEDEF

This tag arrives as a message from a remote node to which the current node was entangled to update the
entanglement information and apply a `Z` correction after the remote node performs an entanglement swap.

$TYPEDFIELDS
"""
@kwdef struct EntanglementUpdateZ <: AbstractTag
    "the pair id currently known by the receiver for the target slot"
    target_pair_id::EntanglementID
    "the pair-id chunk to combine into the target pair"
    other_pair_id::EntanglementID
    "the id of the node to which you were entangled before the swap"
    past_local_node::Int
    "the slot of the node to which you were entangled before the swap"
    past_local_slot::Int
    "the slot of your node that we were entangled to"
    past_remote_slot::Int
    "the id of the node to which you are now entangled after the swap"
    new_remote_node::Int
    "the slot of the node to which you are now entangled after the swap"
    new_remote_slot::Int
    "what Pauli correction you need to perform"
    correction::Int
end
EntanglementUpdateZ(past_local_node::Int, past_local_slot::Int, past_remote_slot::Int, new_remote_node::Int, new_remote_slot::Int, correction::Int) = EntanglementUpdateZ(NO_ENTANGLEMENT_ID, NO_ENTANGLEMENT_ID, past_local_node, past_local_slot, past_remote_slot, new_remote_node, new_remote_slot, correction)
Base.show(io::IO, tag::EntanglementUpdateZ) = print(io, "Update pair $(tag.target_pair_id) at slot .$(tag.past_remote_slot) which used to be entangled to $(tag.past_local_node).$(tag.past_local_slot) to be entangled to $(tag.new_remote_node).$(tag.new_remote_slot), combine with $(tag.other_pair_id), and apply correction X$(tag.correction)")
Tag(tag::EntanglementUpdateZ) = Tag(EntanglementUpdateZ, tag.target_pair_id, tag.other_pair_id, tag.past_local_node, tag.past_local_slot, tag.past_remote_slot, tag.new_remote_node, tag.new_remote_slot, tag.correction)

"""
$TYPEDEF

This tag arrives as a message from a remote node's Cutoff Protocol to which the current node was entangled,
to update the classical metadata of the entangled slot and empty it.
It is also stored at a node to handle incoming `EntanglementUpdate` and `EntanglementDelete` messages.

$TYPEDFIELDS

See also: [`CutoffProt`](@ref)
"""
@kwdef struct EntanglementDelete <: AbstractTag
    "the pair id targeted by this deletion"
    target_pair_id::EntanglementID
    "the node that sent the deletion announcement message after they delete their local qubit"
    send_node::Int
    "the sender's slot containing the decohered qubit"
    send_slot::Int
    "the node receiving the message for qubit deletion"
    rec_node::Int
    "the slot containing decohered qubit"
    rec_slot::Int
end
EntanglementDelete(send_node::Int, send_slot::Int, rec_node::Int, rec_slot::Int) = EntanglementDelete(NO_ENTANGLEMENT_ID, send_node, send_slot, rec_node, rec_slot)
Base.show(io::IO, tag::EntanglementDelete) = print(io, "Deleted pair $(tag.target_pair_id) at $(tag.send_node).$(tag.send_slot) which was entangled to $(tag.rec_node).$(tag.rec_slot)")
Tag(tag::EntanglementDelete) = Tag(EntanglementDelete, tag.target_pair_id, tag.send_node, tag.send_slot, tag.rec_node, tag.rec_slot)

"""
$TYPEDEF

A protocol that generates entanglement between two nodes.
Whenever a pair of empty slots is available, the protocol locks them
and starts probabilistic attempts to establish entanglement.

$TYPEDFIELDS
"""
@kwdef struct EntanglerProt <: AbstractProtocol
    """time-and-schedule-tracking instance from `ConcurrentSim`"""
    sim::Simulation # TODO check that
    """a network graph of registers"""
    net::RegisterNet
    """the vertex index of node A"""
    nodeA::Int
    """the vertex index of node B"""
    nodeB::Int
    """the state being generated (supports symbolic, numeric, noisy, and pure)"""
    pairstate::SymQObj = StabilizerState("ZZ XX")
    """success probability of one attempt of entanglement generation"""
    success_prob::Float64 = 0.001
    """duration of single entanglement attempt"""
    attempt_time::Float64 = 0.001
    """fixed "busy time" duration immediately before starting entanglement generation attempts"""
    local_busy_time_pre::Float64 = 0.0
    """fixed "busy time" duration immediately after the a successful entanglement generation attempt"""
    local_busy_time_post::Float64 = 0.0
    """how long to wait before retrying to lock qubits if no qubits are available (`nothing` for queuing up)"""
    retry_lock_time::Union{Float64,Nothing} = 0.1
    """how many rounds of this protocol to run (`-1` for infinite)"""
    rounds::Int = -1
    """maximum number of attempts to make per round (`-1` for infinite)"""
    attempts::Int = -1
    """function `Int->Bool` or an integer slot number, specifying the slot to take among available free slots in node A"""
    chooseslotA::Union{Int,Function} = alwaystrue
    """function `Int->Bool` or an integer slot number, specifying the slot to take among available free slots in node B"""
    chooseslotB::Union{Int,Function} = alwaystrue
    """whether the protocol should find the first available free slots in the nodes to be entangled or check for free slots randomly from the available slots"""
    randomize::Bool = false
    """whether the protocol should look for unlocked slots to entangle and lock them during the protocol"""
    uselock::Bool = true
    """Repeated rounds of this protocol may lead to monopolizing all slots of a pair of registers, starving or deadlocking other protocols. This field can be used to always leave a minimum number of slots free if there already exists entanglement between the current pair of nodes."""
    margin::Int = 0
    """Like `margin`, but it is enforced even when no entanglement has been established yet. Usually smaller than `margin`."""
    hardmargin::Int = 0
    """concrete `AbstractTag` subtype to add to the entangled qubits, or `nothing` to add no tag. `EntanglementCounterpart` tags include a pair ID; custom tags keep the legacy `tag(remote_node, remote_slot)` shape."""
    tag::Union{Type{<:AbstractTag},Nothing} = EntanglementCounterpart

    function EntanglerProt(sim, net, nodeA, nodeB, pairstate, success_prob, attempt_time, local_busy_time_pre, local_busy_time_post, retry_lock_time, rounds, attempts, chooseslotA, chooseslotB, randomize, uselock, margin, hardmargin, tag)
        @domain 0 < success_prob ≤ 1
        @domain attempt_time ≥ 0
        @domain local_busy_time_pre ≥ 0
        @domain local_busy_time_post ≥ 0
        @domain isnothing(retry_lock_time) || retry_lock_time > 0
        @domain rounds == -1 || rounds ≥ 0
        @domain attempts == -1 || attempts ≥ 0
        @domain margin ≥ 0
        @domain hardmargin ≥ 0
        return new(sim, net, nodeA, nodeB, pairstate, success_prob, attempt_time, local_busy_time_pre, local_busy_time_post, retry_lock_time, rounds, attempts, chooseslotA, chooseslotB, randomize, uselock, margin, hardmargin, tag)
    end
end

protocol_catalog_metadata(::Type{EntanglerProt}) = (
    attachment = :edge,
    attachment_fields = (node_a=:nodeA, node_b=:nodeB),
    required_fields = (),
)

"""Convenience constructor for specifying `rate` of generation instead of success probability and time"""
function EntanglerProt(sim::Simulation, net::RegisterNet, nodeA::Int, nodeB::Int; rate::Union{Nothing,Float64}=nothing, kwargs...)
    if isnothing(rate)
        return EntanglerProt(;sim, net, nodeA, nodeB, kwargs...)
    else
        return EntanglerProt(;sim, net, nodeA, nodeB, kwargs..., success_prob=0.001, attempt_time=0.001/rate)
    end
end

EntanglerProt(net::RegisterNet, nodeA::Int, nodeB::Int; kwargs...) = EntanglerProt(get_time_tracker(net), net, nodeA, nodeB; kwargs...)

#TODO """Convenience constructor for specifying `fidelity` of generation instead of success probability and time"""

@resumable function (prot::EntanglerProt)()
    rounds = prot.rounds
    round = 1
    last_a, last_b = nothing, nothing
    regA = prot.net[prot.nodeA]
    regB = prot.net[prot.nodeB]
    while rounds != 0
        tagtype = prot.tag
        isentangled = !isnothing(tagtype) && (
            tagtype === EntanglementCounterpart ?
            !isnothing(query(regA, tagtype::DataType, prot.nodeB, ❓, ❓; assigned=true)) :
            !isnothing(query(regA, tagtype::DataType, prot.nodeB, ❓; assigned=true))
        )
        margin = isentangled ? prot.margin : prot.hardmargin
        (; chooseslotA, chooseslotB, randomize, uselock) = prot
        a_ = findfreeslot(regA; chooseslot=chooseslotA, randomize=randomize, locked=!uselock, margin=margin)
        b_ = findfreeslot(regB; chooseslot=chooseslotB, randomize=randomize, locked=!uselock, margin=margin)

        if isnothing(a_) || isnothing(b_)
            if isnothing(prot.retry_lock_time)
                @debug(
                    "Free slots unavailable",
                    _group=LOG_GROUPS.protocol,
                    event=:free_slots_unavailable,
                    protocol_log_context(prot)...,
                    round=round,
                    wait_mode=:tag_change,
                    retry_after_s=nothing,
                )
                @yield onchange(regA, Tag) | onchange(regB, Tag)
            else
                @debug(
                    "Free slots unavailable",
                    _group=LOG_GROUPS.protocol,
                    event=:free_slots_unavailable,
                    protocol_log_context(prot)...,
                    round=round,
                    wait_mode=:fixed_delay,
                    retry_after_s=prot.retry_lock_time,
                )
                @yield timeout(prot.sim, prot.retry_lock_time::Float64)
            end
            continue
        end
        # we are now certain that a_ and b_ are not nothing. The compiler is not smart enough to figure this out
        # on its own, so we need to tell it explicitly. A new variable name is needed due to @resumable.
        a = a_::RegRef
        b = b_::RegRef

        if uselock
            @yield lock(a) & lock(b) # this yield is expected to return immediately
        end

        @yield timeout(prot.sim, prot.local_busy_time_pre)
        attempts = if isone(prot.success_prob)
            1
        else
            rand(Geometric(prot.success_prob))+1
        end
        if prot.attempts == -1 || prot.attempts >= attempts
            @yield timeout(prot.sim, attempts * prot.attempt_time)
            initialize!((a,b), prot.pairstate; time=now(prot.sim))
            @yield timeout(prot.sim, prot.local_busy_time_post)

            # generate a random identifier for the new pair
            pair_id = fresh_entanglement_id()
            # tag local node a with EntanglementCounterpart remote_node_idx_b remote_slot_idx_b pair_id
            if tagtype === EntanglementCounterpart
                _tag_entanglement_counterpart!(a, prot.nodeB, b.idx, pair_id, prot)
            elseif !isnothing(tagtype)
                tag!(a, tagtype::DataType, prot.nodeB, b.idx)
            end
            last_a = a.idx
            # tag local node b with EntanglementCounterpart remote_node_idx_a remote_slot_idx_a pair_id
            if tagtype === EntanglementCounterpart
                _tag_entanglement_counterpart!(b, prot.nodeA, a.idx, pair_id, prot)
            elseif !isnothing(tagtype)
                tag!(b, tagtype::DataType, prot.nodeA, a.idx)
            end
            last_b = b.idx

            @debug(
                "Entangled a pair",
                _group=LOG_GROUPS.protocol,
                event=:pair_entangled,
                protocol_log_context(prot)...,
                round=round,
                slots=(a.idx, b.idx),
                pair_id=pair_id,
                attempts=attempts,
            )
        else
            @yield timeout(prot.sim, prot.attempts * prot.attempt_time)
            @debug(
                "Exhausted entanglement attempts",
                _group=LOG_GROUPS.protocol,
                event=:attempts_exhausted,
                protocol_log_context(prot)...,
                round=round,
                attempts=prot.attempts,
            )
        end
        if uselock
            unlock(a)
            unlock(b)
        end
        rounds==-1 || (rounds -= 1)
        round += 1
    end
    return prot.nodeA, last_a, prot.nodeB, last_b
end


"""
$TYPEDEF

A protocol, running at a given node, listening for messages that indicate something has happened to a remote qubit entangled with one of the local qubits.

$TYPEDFIELDS
"""
@kwdef struct EntanglementTracker <: AbstractProtocol
    """time-and-schedule-tracking instance from `ConcurrentSim`"""
    sim::Simulation
    """a network graph of registers"""
    net::RegisterNet
    """the vertex of the node where the tracker is working"""
    node::Int
end

protocol_catalog_metadata(::Type{EntanglementTracker}) = (
    attachment = :node,
    attachment_fields = (node=:node,),
    required_fields = (),
)

EntanglementTracker(net::RegisterNet, node::Int) = EntanglementTracker(get_time_tracker(net), net, node)

@resumable function (prot::EntanglementTracker)()
    nodereg = prot.net[prot.node]
    mb = messagebuffer(prot.net, prot.node)
    while true
        workwasdone = true # waiting is not enough because we might have multiple rounds of work to do
        while workwasdone
            workwasdone = false
            for (updatetagsymbol, updategate) in ((EntanglementUpdateX, Z), (EntanglementUpdateZ, X), (EntanglementDelete, nothing)) # TODO this is getting ugly. Refactor EntanglementUpdateX and EntanglementUpdateZ to be the same parameterized tag
                # look for EntanglementUpdate? or EntanglementDelete message sent to us
                if !isnothing(updategate) # EntanglementUpdate
                    msg = querydelete!(mb, updatetagsymbol, ❓, ❓, ❓, ❓, ❓, ❓, ❓, ❓)
                    isnothing(msg) && continue
                    (src, (_, target_pair_id, other_pair_id, pastremotenode, pastremoteslotid, localslotid, newremotenode, newremoteslotid, correction)) = msg
                    target_pair_id = target_pair_id::EntanglementID
                    other_pair_id = other_pair_id::EntanglementID
                else # EntanglementDelete
                    msg = querydelete!(mb, updatetagsymbol, ❓, ❓, ❓, ❓, ❓)
                    isnothing(msg) && continue
                    (src, (_, target_pair_id, pastremotenode, pastremoteslotid, _, localslotid)) = msg
                    target_pair_id = target_pair_id::EntanglementID
                    other_pair_id = NO_ENTANGLEMENT_ID
                    newremotenode = -1
                    newremoteslotid = -1
                    correction = 0
                end

                @debug(
                    "Received an entanglement update",
                    _group=LOG_GROUPS.protocol,
                    event=:message_received,
                    protocol_log_context(prot)...,
                    src_node=src,
                    src_slot=pastremoteslotid,
                    dst_slot=localslotid,
                    message_type=_message_type(msg.tag),
                    pair_id=target_pair_id,
                )
                workwasdone = true
                localslot = nodereg[localslotid]

                new_pair_id = isnothing(updategate) ? target_pair_id : combine_entanglement_ids(target_pair_id::EntanglementID, other_pair_id::EntanglementID)

                # Check if the local slot is still present and believed to be entangled.
                # The physical slot lock is only needed when we may mutate the live qubit.
                # If no matching counterpart tag is present, fall through to the history
                # metadata path without waiting behind unrelated reuse of an empty slot.
                counterpart = query(localslot, EntanglementCounterpart, pastremotenode, pastremoteslotid, target_pair_id)
                if !isnothing(counterpart)
                    @debug(
                        "Requested a slot lock",
                        _group=LOG_GROUPS.protocol,
                        event=:slot_lock_requested,
                        protocol_log_context(prot)...,
                        slot=localslotid,
                        pair_id=target_pair_id,
                    )
                    @yield lock(localslot)
                    @debug(
                        "Acquired a slot lock",
                        _group=LOG_GROUPS.protocol,
                        event=:slot_lock_acquired,
                        protocol_log_context(prot)...,
                        slot=localslotid,
                        pair_id=target_pair_id,
                    )
                    counterpart = querydelete!(localslot, EntanglementCounterpart, pastremotenode, pastremoteslotid, target_pair_id)
                    if !isnothing(counterpart)
                        if !isassigned(localslot)
                            unlock(localslot)
                            error("There was an error in the entanglement tracking protocol `EntanglementTracker`. We were attempting to forward a classical message from a node that performed a swap to the remote entangled node. However, on reception of that message it was found that the remote node has lost track of its part of the entangled state although it still keeps a `Tag` as a record of it being present.") # TODO make it configurable whether an error is thrown and plug it into the logging module
                        end
                        if !isnothing(updategate) # EntanglementUpdate
                            # Pauli frame correction gate
                            if correction==2
                                apply!(localslot, updategate)
                            end
                            if newremotenode != -1 #TODO: this is a bit hacky
                                # tag local with updated EntanglementCounterpart new_remote_node new_remote_slot_idx
                                _tag_entanglement_counterpart!(
                                    localslot, newremotenode, newremoteslotid,
                                    new_pair_id, prot
                                )
                            else
                                _tag_entanglement_counterpart!(
                                    localslot, pastremotenode, pastremoteslotid,
                                    target_pair_id, prot
                                )
                            end
                        else # EntanglementDelete
                            traceout!(localslot)
                        end
                        unlock(localslot)
                        continue
                    end
                    unlock(localslot)
                end

                # If there is nothing still stored locally, check if we have a record of the entanglement being swapped to a different remote node,
                # and forward the message to that node.
                history = querydelete!(localslot, EntanglementHistory,
                                    pastremotenode, pastremoteslotid, # who we were entangled to (node, slot)
                                    ❓, ❓,                             # who we swapped with (node, slot)
                                    ❓,                                 # which local slot used to be entangled with whom we swapped with
                                    target_pair_id, ❓)                # pair-id chunks for this side and the swapped side
                if !isnothing(history)
                    _, _, _, whoweswappedwith_node, whoweswappedwith_slotidx, swappedlocal_slotidx, local_chunk_id, swapped_chunk_id = history.tag
                    local_chunk_id = local_chunk_id::EntanglementID
                    swapped_chunk_id = swapped_chunk_id::EntanglementID
                    forwarded_target_pair_id = combine_entanglement_ids(local_chunk_id::EntanglementID, swapped_chunk_id::EntanglementID)
                    if !isnothing(updategate) # EntanglementUpdate
                        # A history tag stores the two chunks joined by the
                        # local swap. An update from this side advances only
                        # this side's chunk; the end-to-end ID is recomputed
                        # when notifying the opposite side.
                        updated_local_chunk_id = combine_entanglement_ids(local_chunk_id::EntanglementID, other_pair_id::EntanglementID)
                        tag!(localslot, EntanglementHistory, newremotenode, newremoteslotid, whoweswappedwith_node, whoweswappedwith_slotidx, swappedlocal_slotidx, updated_local_chunk_id, swapped_chunk_id)
                        @debug(
                            "Forwarded an entanglement update",
                            _group=LOG_GROUPS.protocol,
                            event=:message_forwarded,
                            protocol_log_context(prot)...,
                            src_slot=localslotid,
                            dst_node=whoweswappedwith_node,
                            dst_slot=whoweswappedwith_slotidx,
                            message_type=_message_type(msg.tag),
                            pair_id=forwarded_target_pair_id,
                        )
                        msghist = Tag(updatetagsymbol, forwarded_target_pair_id, other_pair_id, pastremotenode, pastremoteslotid, whoweswappedwith_slotidx, newremotenode, newremoteslotid, correction)
                        put!(channel(prot.net, prot.node=>whoweswappedwith_node; permit_forward=true), msghist)
                    else # EntanglementDelete
                        # We have a delete message but the qubit was swapped so add a tag and forward to swapped node
                        @debug(
                            "Forwarded an entanglement deletion",
                            _group=LOG_GROUPS.protocol,
                            event=:message_forwarded,
                            protocol_log_context(prot)...,
                            src_slot=localslotid,
                            dst_node=whoweswappedwith_node,
                            dst_slot=whoweswappedwith_slotidx,
                            message_type=_message_type(msg.tag),
                            pair_id=forwarded_target_pair_id,
                        )
                        msghist = Tag(updatetagsymbol, forwarded_target_pair_id, pastremotenode, pastremoteslotid, whoweswappedwith_node, whoweswappedwith_slotidx)
                        tag!(localslot, updatetagsymbol, target_pair_id, prot.node, localslot.idx, pastremotenode, pastremoteslotid)
                        put!(channel(prot.net, prot.node=>whoweswappedwith_node; permit_forward=true), msghist)
                    end
                    continue
                end

                # Finally, if there the history of a swap is not present in the log anymore,
                # it must be because a delete message was received, and forwarded,
                # and the entanglement history was deleted, and replaced with an entanglement delete tag.
                if !isnothing(querydelete!(localslot, EntanglementDelete, target_pair_id, prot.node, localslot.idx, pastremotenode, pastremoteslotid)) #deletion from both sides of the swap, deletion msg when both qubits of a pair are deleted, or when EU arrives after ED at swap node with two simultaneous swaps and deletion on one side
                    if !(isnothing(updategate)) # EntanglementUpdate
                        # to handle a possible delete-swap-swap case, we need to update the EntanglementDelete tag
                        tag!(localslot, EntanglementDelete, combine_entanglement_ids(target_pair_id::EntanglementID, other_pair_id::EntanglementID), prot.node, localslot.idx, newremotenode, newremoteslotid)
                        @debug(
                            "Applied an update to a deleted qubit marker",
                            _group=LOG_GROUPS.protocol,
                            event=:deleted_qubit_update_applied,
                            protocol_log_context(prot)...,
                            slot=localslotid,
                            message_type=_message_type(msg.tag),
                            pair_id=target_pair_id,
                        )
                    else # EntanglementDelete
                        # when the message is EntanglementDelete and the slot history also has an EntanglementDelete tag (both qubits were deleted), do nothing
                        @debug(
                            "Dropped a message for a deleted qubit",
                            _group=LOG_GROUPS.protocol,
                            event=:message_dropped,
                            protocol_log_context(prot)...,
                            slot=localslotid,
                            message_type=_message_type(msg.tag),
                            pair_id=target_pair_id,
                        )
                    end
                    continue
                end

                # With pair IDs, unmatched updates/deletes should be rare: a live
                # counterpart, a history entry, or a delete marker should normally
                # carry the target ID until the message is handled. The expected
                # benign cause is a bounded history log that discarded the needed
                # entry before a delayed message arrived.
                @warn(
                    "Dropped a stale entanglement message",
                    _group=LOG_GROUPS.protocol,
                    event=:stale_message_dropped,
                    protocol_log_context(prot)...,
                    slot=localslotid,
                    message_type=_message_type(msg.tag),
                    pair_id=target_pair_id,
                )
                continue
            end
        end
        @debug(
            "Started waiting for a message",
            _group=LOG_GROUPS.protocol,
            event=:message_wait_started,
            protocol_log_context(prot)...,
        )
        @yield onchange(mb)
        @debug(
            "Finished waiting for a message",
            _group=LOG_GROUPS.protocol,
            event=:message_wait_finished,
            protocol_log_context(prot)...,
        )
    end
end

"""
$TYPEDEF

A protocol running between two nodes, checking periodically for any entangled pairs between the two nodes and consuming/emptying the qubit slots.

This protocol permits virtual edges, meaning it can operate between any two nodes in the network regardless of whether they are physically connected by an edge.

$FIELDS
"""
@kwdef struct EntanglementConsumer <: AbstractProtocol
    """time-and-schedule-tracking instance from `ConcurrentSim`"""
    sim::Simulation
    """a network graph of registers"""
    net::RegisterNet
    """the vertex index of node A"""
    nodeA::Int
    """the vertex index of node B"""
    nodeB::Int
    """time period between successive queries on the nodes (`nothing` for queuing up and waiting for available pairs)"""
    period::Union{Float64,Nothing} = 0.1
    """concrete `AbstractTag` subtype which the consumer is looking for; defaults to `EntanglementCounterpart`, where reciprocal tags must also agree on pair ID"""
    tag::Type{<:AbstractTag} = EntanglementCounterpart
    """stores the time and resulting observable from querying nodeA and nodeB for `EntanglementCounterpart`; the storage type is not part of the public API and may change in future versions"""
    _log::Vector{@NamedTuple{t::Float64, obs1::Float64, obs2::Float64}} = @NamedTuple{t::Float64, obs1::Float64, obs2::Float64}[]

    function EntanglementConsumer(sim, net, nodeA, nodeB, period, tag, _log)
        @domain isnothing(period) || period > 0
        return new(sim, net, nodeA, nodeB, period, tag, _log)
    end
end

protocol_catalog_metadata(::Type{EntanglementConsumer}) = (
    attachment = :edge,
    attachment_fields = (node_a=:nodeA, node_b=:nodeB),
    required_fields = (),
)

function EntanglementConsumer(sim::Simulation, net::RegisterNet, nodeA::Int, nodeB::Int; kwargs...)
    return EntanglementConsumer(;sim, net, nodeA, nodeB, kwargs...)
end
function EntanglementConsumer(net::RegisterNet, nodeA::Int, nodeB::Int; kwargs...)
    return EntanglementConsumer(get_time_tracker(net), net, nodeA, nodeB; kwargs...)
end

permits_virtual_edge(::Type{EntanglementConsumer}) = true

@resumable function (prot::EntanglementConsumer)()
    regA = prot.net[prot.nodeA]
    regB = prot.net[prot.nodeB]
    while true
        use_pair_id = prot.tag === EntanglementCounterpart
        query1 = if use_pair_id
            query(regA, prot.tag, prot.nodeB, ❓, ❓; locked=false, assigned=true)
        else
            query(regA, prot.tag, prot.nodeB, ❓; locked=false, assigned=true)
        end # TODO Need a `querydelete!` dispatch on `Register` rather than using `query` here followed by `untag!` below
        if isnothing(query1)
            if isnothing(prot.period)
                @debug(
                    "Entanglement unavailable",
                    _group=LOG_GROUPS.protocol,
                    event=:entanglement_unavailable,
                    protocol_log_context(prot)...,
                    src_node=prot.nodeA,
                    wait_mode=:tag_change,
                    retry_after_s=nothing,
                )
                @yield onchange(regA, Tag)
            else
                @debug(
                    "Entanglement unavailable",
                    _group=LOG_GROUPS.protocol,
                    event=:entanglement_unavailable,
                    protocol_log_context(prot)...,
                    src_node=prot.nodeA,
                    wait_mode=:fixed_delay,
                    retry_after_s=prot.period,
                )
                @yield timeout(prot.sim, prot.period::Float64)
            end
            continue
        else
            pair_id = use_pair_id ? query1.tag[4] : nothing
            query2 = if use_pair_id
                query(regB, prot.tag, prot.nodeA, query1.slot.idx, pair_id; locked=false, assigned=true)
            else
                query(regB, prot.tag, prot.nodeA, query1.slot.idx; locked=false, assigned=true)
            end
            if isnothing(query2) # in case EntanglementUpdate hasn't reached the second node yet, but the first node has the EntanglementCounterpart
                if isnothing(prot.period)
                    @debug(
                        "Entanglement unavailable",
                        _group=LOG_GROUPS.protocol,
                        event=:entanglement_unavailable,
                        protocol_log_context(prot)...,
                        dst_node=prot.nodeB,
                        wait_mode=:tag_change,
                        retry_after_s=nothing,
                    )
                    @yield onchange(regB, Tag)
                else
                    @debug(
                        "Entanglement unavailable",
                        _group=LOG_GROUPS.protocol,
                        event=:entanglement_unavailable,
                        protocol_log_context(prot)...,
                        dst_node=prot.nodeB,
                        wait_mode=:fixed_delay,
                        retry_after_s=prot.period,
                    )
                    @yield timeout(prot.sim, prot.period::Float64)
                end
                continue
            end
        end

        q1 = query1.slot
        q2 = query2.slot
        @yield lock(q1) & lock(q2)
        # Re-query under lock and require reciprocal pair IDs. The slot tuple is
        # routing metadata; `pair_id` is the actual entangled-pair identity.
        query1 = if use_pair_id
            query(q1, prot.tag, prot.nodeB, q2.idx, pair_id; locked=true, assigned=true)
        else
            query(q1, prot.tag, prot.nodeB, q2.idx; locked=true, assigned=true)
        end
        query2 = if use_pair_id
            query(q2, prot.tag, prot.nodeA, q1.idx, pair_id; locked=true, assigned=true)
        else
            query(q2, prot.tag, prot.nodeA, q1.idx; locked=true, assigned=true)
        end
        if isnothing(query1) || isnothing(query2)
            @debug(
                "Entanglement query was invalidated",
                _group=LOG_GROUPS.protocol,
                event=:query_invalidated,
                protocol_log_context(prot)...,
                slots=(q1.idx, q2.idx),
                pair_id=pair_id,
            )
            unlock(q1)
            unlock(q2)
            continue
        end

        @debug(
            "Consumed entanglement",
            _group=LOG_GROUPS.protocol,
            event=:entanglement_consumed,
            protocol_log_context(prot)...,
            slots=(q1.idx, q2.idx),
            pair_id=pair_id,
        )
        untag!(q1, query1.id)
        untag!(q2, query2.id)
        # TODO do we need to add EntanglementHistory or EntanglementDelete and should that be a different EntanglementHistory since the current one is specifically for Swapper
        # TODO currently when calculating the observable we assume that EntanglerProt.pairstate is always (|00⟩ + |11⟩)/√2, make it more general for other states
        ob1 = observable((q1, q2), Z⊗Z)
        ob2 = observable((q1, q2), X⊗X)
        if isnothing(ob1) || isnothing(ob2)
            @error(
                "Dropped a stale entangled pair",
                _group=LOG_GROUPS.protocol,
                event=:stale_pair_dropped,
                protocol_log_context(prot)...,
                slots=(q1.idx, q2.idx),
                pair_id=pair_id,
            )
            traceout!(regA[q1.idx], regB[q2.idx])
            unlock(q1)
            unlock(q2)
            continue
        end
        ob1 = real(ob1)
        ob2 = real(ob2)

        traceout!(regA[q1.idx], regB[q2.idx])
        push!(prot._log, (now(prot.sim), ob1, ob2))
        unlock(q1)
        unlock(q2)
        if !isnothing(prot.period)
            @yield timeout(prot.sim, prot.period::Float64)
        end
    end
end


include("cutoff.jl")
include("swapping.jl")
include("switches.jl")
using .Switches
include("qtcp.jl")
using .QTCP
include("mbqc.jl")
using .MBQCEntanglementDistillation

_protocol_nodes(prot::EntanglementTracker) = (prot.node,)
_protocol_nodes(prot::SwapperProt) = (prot.node,)
_protocol_nodes(prot::CutoffProt) = (prot.node,)
_protocol_nodes(prot::EntanglerProt) = (prot.nodeA, prot.nodeB)
_protocol_nodes(prot::EntanglementConsumer) = (prot.nodeA, prot.nodeB)

_protocol_nodes(prot::Switches.SimpleSwitchDiscreteProt) =
    (prot.switchnode, Tuple(prot.clientnodes)...)
protocol_log_context(deleter::Switches._SwitchSynchronizedDelete) =
    protocol_log_context(deleter.prot)

_protocol_nodes(prot::QTCP.EndNodeController) = (prot.node,)
_protocol_nodes(prot::QTCP.NetworkNodeController) = (prot.node,)
_protocol_nodes(prot::QTCP.LinkController) = (prot.nodeA, prot.nodeB)

_protocol_nodes(prot::MBQCEntanglementDistillation.GraphStateConstructor) =
    Tuple(prot.nodes)
_protocol_nodes(prot::MBQCEntanglementDistillation.GraphToResource) =
    Tuple(prot.nodes)
_protocol_nodes(prot::MBQCEntanglementDistillation.PurifierBellMeasurements) =
    Tuple(prot.nodes)
_protocol_nodes(prot::MBQCEntanglementDistillation.MBQCPurificationTracker) =
    Tuple(prot.nodes)

include("show.jl")

end # module
