module QTCP

using QuantumSavory
import QuantumSavory: Tag, get_time_tracker
using QuantumSavory.ProtocolZoo
using QuantumSavory.ProtocolZoo: AbstractProtocol
import QuantumSavory.ProtocolZoo: protocol_catalog_metadata
using QuantumSavory.CircuitZoo: LocalEntanglementSwap
import ConcurrentSim
using ConcurrentSim: Simulation, @yield, timeout, @process, now, Process, Resource
import ResumableFunctions
using ResumableFunctions: @resumable
using DocStringExtensions
import Graphs

export QDatagram, Flow,
    LinkLevelRequest, LinkLevelReply, LinkLevelReplyAtHop, LinkLevelReplyAtSource,
    QTCPPairBegin, QTCPPairEnd,
    NetworkNodeController, EndNodeController, LinkController

###
# Message types
###

"""
$TYPEDEF

$TYPEDFIELDS
"""
@kwdef struct Flow <: AbstractTag
    "who initiates the request and also initiates the qdatagrams"
    src::Int
    "the destination node"
    dst::Int
    "the number of requested pairs"
    npairs::Int
    "the unique id of the flow"
    uuid::Int
end
Base.show(io::IO, tag::Flow) = print(io, "Flow `$(tag.uuid)` | $(tag.src) → $(tag.dst) | $(tag.npairs) pairs")
Tag(tag::Flow) = Tag(Flow, tag.src, tag.dst, tag.npairs, tag.uuid) # TODO automate this


"""
$TYPEDEF

$TYPEDFIELDS
"""
@kwdef struct QTCPPairBegin <: AbstractTag
    "the uuid of the flow we are generated for"
    flow_uuid::Int
    "who initiates the flow request and also initiates the qdatagrams"
    flow_src::Int
    "the destination node for the flow"
    flow_dst::Int
    "sequence number of the qdataframe in the given flow"
    seq_num::Int
    "the memory slot of the entangled qubit"
    memory_slot::Int
    "original start time of the qdatagram"
    start_time::Float64
end
Base.show(io::IO, tag::QTCPPairBegin) = print(io, "QTCPPairBegin `$(tag.flow_uuid).$(tag.seq_num)` $(tag.flow_src) → $(tag.flow_dst) | started at $(tag.start_time) | memory slot $(tag.memory_slot)")
Tag(tag::QTCPPairBegin) = Tag(QTCPPairBegin, tag.flow_uuid, tag.flow_src, tag.flow_dst, tag.seq_num, tag.memory_slot, tag.start_time) # TODO automate this


"""
$TYPEDEF

$TYPEDFIELDS
"""
@kwdef struct QTCPPairEnd <: AbstractTag
    "the uuid of the flow we are generated for"
    flow_uuid::Int
    "who initiates the flow request and also initiates the qdatagrams"
    flow_src::Int
    "the destination node for the flow"
    flow_dst::Int
    "sequence number of the qdataframe in the given flow"
    seq_num::Int
    "the memory slot of the entangled qubit"
    memory_slot::Int
    "original start time of the qdatagram"
    start_time::Float64
end
Base.show(io::IO, tag::QTCPPairEnd) = print(io, "QTCPPairEnd `$(tag.flow_uuid).$(tag.seq_num)` $(tag.flow_src) → $(tag.flow_dst) | started at $(tag.start_time) | memory slot $(tag.memory_slot)")
Tag(tag::QTCPPairEnd) = Tag(QTCPPairEnd, tag.flow_uuid, tag.flow_src, tag.flow_dst, tag.seq_num, tag.memory_slot, tag.start_time) # TODO automate this


"""
$TYPEDEF

$TYPEDFIELDS
"""
@kwdef struct QDatagram <: AbstractTag
    "the uuid of the flow we are generated for"
    flow_uuid::Int
    "who initiates the flow request and also initiates the qdatagrams"
    flow_src::Int
    "the destination node for the flow"
    flow_dst::Int
    "the Pauli frame correction"
    correction::Int
    "sequence number of the qdataframe in the given flow"
    seq_num::Int
    "original start time of the qdatagram"
    start_time::Float64
end
Base.show(io::IO, tag::QDatagram) = print(io, "QDatagram `$(tag.flow_uuid).$(tag.seq_num)` $(tag.flow_src) → $(tag.flow_dst) | started at $(tag.start_time) | correction $(tag.correction)")
Tag(tag::QDatagram) = Tag(QDatagram, tag.flow_uuid, tag.flow_src, tag.flow_dst, tag.correction, tag.seq_num, tag.start_time) # TODO automate this

"""
$TYPEDEF

$TYPEDFIELDS
"""
@kwdef struct QDatagramSuccess <: AbstractTag
    "the uuid of the flow we are generated for"
    flow_uuid::Int
    "sequence number of the qdataframe in the given flow"
    seq_num::Int
    "time at which the qdatagram was generated by the flow source"
    start_time::Float64
end
Base.show(io::IO, tag::QDatagramSuccess) = print(io, "QDatagramSuccess `$(tag.flow_uuid).$(tag.seq_num)` | started at $(tag.start_time)")
Tag(tag::QDatagramSuccess) = Tag(QDatagramSuccess, tag.flow_uuid, tag.seq_num, tag.start_time) # TODO automate this

"""
$TYPEDEF

$TYPEDFIELDS
"""
@kwdef struct LinkLevelRequest <: AbstractTag
    "the uuid of the flow we are providing entanglement for"
    flow_uuid::Int
    "sequence number of the qdataframe we are providing entanglement for"
    seq_num::Int
    "the remote node with which we want to establish entanglement"
    remote_node::Int
end
Base.show(io::IO, tag::LinkLevelRequest) = print(io, "LinkLevelRequest for flow $(tag.flow_uuid), sequence $(tag.seq_num), remote node $(tag.remote_node)")
Tag(tag::LinkLevelRequest) = Tag(LinkLevelRequest, tag.flow_uuid, tag.seq_num, tag.remote_node) # TODO automate this

"""
$TYPEDEF

$TYPEDFIELDS
"""
@kwdef struct LinkLevelReply <: AbstractTag
    "the uuid of the flow we are providing entanglement for"
    flow_uuid::Int
    "sequence number of the qdataframe we are providing entanglement for"
    seq_num::Int
    "where is the entangled qubit stored"
    memory_slot::Int
end
Base.show(io::IO, tag::LinkLevelReply) = print(io, "LinkLevelReply for flow $(tag.flow_uuid), sequence $(tag.seq_num), memory slot $(tag.memory_slot)")
Tag(tag::LinkLevelReply) = Tag(LinkLevelReply, tag.flow_uuid, tag.seq_num, tag.memory_slot) # TODO automate this

"""
$TYPEDEF

$TYPEDFIELDS
"""
@kwdef struct LinkLevelReplyAtSource <: AbstractTag
    "the uuid of the flow we are providing entanglement for"
    flow_uuid::Int
    "sequence number of the qdataframe we are providing entanglement for"
    seq_num::Int
    "where is the entangled qubit stored"
    memory_slot::Int
end
Base.show(io::IO, tag::LinkLevelReplyAtSource) = print(io, "LinkLevelReplyAtSource for flow $(tag.flow_uuid), sequence $(tag.seq_num), memory slot $(tag.memory_slot)")
Tag(tag::LinkLevelReplyAtSource) = Tag(LinkLevelReplyAtSource, tag.flow_uuid, tag.seq_num, tag.memory_slot) # TODO automate this

"""
$TYPEDEF

$TYPEDFIELDS
"""
@kwdef struct LinkLevelReplyAtHop <: AbstractTag
    "the uuid of the flow we are providing entanglement for"
    flow_uuid::Int
    "sequence number of the qdataframe we are providing entanglement for"
    seq_num::Int
    "where is the entangled qubit stored"
    memory_slot::Int
end
Base.show(io::IO, tag::LinkLevelReplyAtHop) = print(io, "LinkLevelReplyAtHop for flow $(tag.flow_uuid), sequence $(tag.seq_num), memory slot $(tag.memory_slot)")
Tag(tag::LinkLevelReplyAtHop) = Tag(LinkLevelReplyAtHop, tag.flow_uuid, tag.seq_num, tag.memory_slot) # TODO automate this


###
# Protocol declaration
###

"""
$TYPEDEF

$FIELDS

Managing the following transformations of classical control signals:

- `Flow` @ start node → sequence of `QDatagram` @ start node
- `QDatagramSuccess` received from destination node → continuing the flow's sequence, reevaluating window
    - and converting `LinkLevelReplyAtSource` to `QTCPPairBegin`
- `QDatagram` @ destination node → `QDatagramSuccess` sent to start node
    - and converting `LinkLevelReplyAtHop` to `QTCPPairEnd`
"""
@kwdef struct EndNodeController <: AbstractProtocol
    """time-and-schedule-tracking instance from `ConcurrentSim`"""
    sim::Simulation
    """a network graph of registers"""
    net::RegisterNet
    """the vertex index of where the protocol is located"""
    node::Int
end

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

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

"""
$TYPEDEF

$FIELDS

Managing the following transformations of classical control signals:

- `QDatagram` → `LinkLevelRequest` to link with next hop (according to a pathfinding algorithm)
- `LinkLevelReply` → `QDatagram` forwarded to next hop and
    - either a swap with the pre-existing `LinkLevelReplyAtHop`
    - or just tagging `LinkLevelReplyAtSource` if we are at the source node of the flow
- this is the entity that actually:
    - performs pathfinding
    - performs entanglement swapping
"""
@kwdef struct NetworkNodeController <: AbstractProtocol
    """time-and-schedule-tracking instance from `ConcurrentSim`"""
    sim::Simulation
    """a network graph of registers"""
    net::RegisterNet
    """the vertex index of where the protocol is located"""
    node::Int
end

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

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

"""
$TYPEDEF

$FIELDS

Managing the following transformations of classical control signals:

- `LinkLevelRequest` → `LinkLevelReply` and `LinkLevelReplyAtHop`
- this is the entity that actually establishes the link-level entanglement
"""
@kwdef struct LinkController <: AbstractProtocol
    """time-and-schedule-tracking instance from `ConcurrentSim`"""
    sim::Simulation
    """a network graph of registers"""
    net::RegisterNet
    """the vertex index of one of the nodes in the link (Alice)"""
    nodeA::Int
    """the vertex index of one of the nodes in the link (Bob)"""
    nodeB::Int
end

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

LinkController(net::RegisterNet, nodeA::Int, nodeB::Int) = LinkController(get_time_tracker(net), net, nodeA, nodeB)


###
# Protocol Implementation
###

@resumable function (prot::EndNodeController)()
    (;sim, net, node) = prot
    mb = messagebuffer(net, node)

    current_flows = Set{Int}() # the uuids of flows currently being processed
    # these are keyed by uuid:
    qdatagrams_in_flight   = Dict{Int,Int}() # number of datagrams currently believed to be in flight # TODO turn value to ordered list
    qdatagrams_sent        = Dict{Int,Int}() # total number of datagrams already sent
    pairs_left_to_fulfill = Dict{Int,Int}() # total number of pairs still to be established
    destination            = Dict{Int,Int}() # the destination

    WINDOW = 3 # TODO per flow

    while true
        workwasdone = true
        while workwasdone
            workwasdone = false

            # check if there is an approved flow and start it
            flow = querydelete!(mb, Flow, node, ❓, ❓, ❓)
            if !isnothing(flow)
                workwasdone = true
                _, _, dst, npairs, uuid = flow.tag
                push!(current_flows, uuid)
                qdatagrams_in_flight[uuid]   = 0
                qdatagrams_sent[uuid]        = 0
                pairs_left_to_fulfill[uuid] = npairs
                destination[uuid]            = dst
                @debug(
                    "Started a flow",
                    _group=LOG_GROUPS.protocol,
                    event=:flow_started,
                    protocol_log_context(prot)...,
                    flow_id=uuid,
                    src_node=node,
                    dst_node=dst,
                )
            end

            # check if there are datagram acknowledgements
            success = querydelete!(mb, QDatagramSuccess, ❓, ❓, ❓)
            if !isnothing(success)
                workwasdone = true
                # TODO implement drop detection and window modification
                _, flow_uuid, seq_num, start_time = success.tag
                start_time = start_time::Float64
                qdatagrams_in_flight[flow_uuid]   -= 1
                pairs_left_to_fulfill[flow_uuid] -= 1

                # Check if there are any LinkLevelReplyAtSource messages and turn them into QTCPPairBegin messages
                link_reply = querydelete!(mb, LinkLevelReplyAtSource, flow_uuid, seq_num, ❓)
                @assert !isnothing(link_reply) "No LinkLevelReplyAtSource message found after the success of flow $(flow_uuid), sequence $(seq_num) at node $(node)"
                _, _, _, memory_slot = link_reply.tag
                pair_begin = QTCPPairBegin(;
                    flow_uuid,
                    flow_src=node,
                    flow_dst=success.src,
                    seq_num,
                    memory_slot,
                    start_time
                )
                put!(net[node], pair_begin)
                @debug(
                    "Acknowledged a datagram",
                    _group=LOG_GROUPS.protocol,
                    event=:datagram_acknowledged,
                    protocol_log_context(prot)...,
                    flow_id=flow_uuid,
                    sequence_number=seq_num,
                    src_node=success.src,
                    dst_node=node,
                )

                # if we have fulfilled all pairs, remove the flow in every data structure
                if pairs_left_to_fulfill[flow_uuid] == 0
                    delete!(current_flows, flow_uuid)
                    delete!(qdatagrams_in_flight, flow_uuid)
                    delete!(qdatagrams_sent, flow_uuid)
                    delete!(pairs_left_to_fulfill, flow_uuid)
                    delete!(destination, flow_uuid)
                    @debug(
                        "Completed a flow",
                        _group=LOG_GROUPS.protocol,
                        event=:flow_completed,
                        protocol_log_context(prot)...,
                        flow_id=flow_uuid,
                    )
                end
            end

            # check if we just received a qdatagram for which we are the flow destination
            qdatagram = querydelete!(mb, QDatagram, ❓, ❓, node, ❓, ❓, ❓)
            if !isnothing(qdatagram)
                workwasdone = true
                # We need to generate a QDatagramSuccess message and send it back to the flow source
                _, flow_uuid, flow_src, flow_dst, corrections, seq_num, start_time = qdatagram.tag
                start_time = start_time::Float64
                qdatagram_success = QDatagramSuccess(flow_uuid, seq_num, start_time)
                put!(channel(net, node=>flow_src; permit_forward=true), qdatagram_success)
                # TODO implement Pauli corrections

                # Check if there are any LinkLevelReplyAtHop messages and turn them into QTCPPairEnd messages
                link_reply = querydelete!(mb, LinkLevelReplyAtHop, flow_uuid, seq_num, ❓)
                @assert !isnothing(link_reply) "No LinkLevelReplyAtHop message found after the success of flow $(flow_uuid), sequence $(seq_num) at node $(node)"
                _, _, _, memory_slot = link_reply.tag
                pair_end = QTCPPairEnd(;
                    flow_uuid,
                    flow_src=flow_src,
                    flow_dst=node,
                    seq_num,
                    memory_slot,
                    start_time
                )
                put!(net[node], pair_end)
                @debug(
                    "Delivered a datagram",
                    _group=LOG_GROUPS.protocol,
                    event=:datagram_delivered,
                    protocol_log_context(prot)...,
                    flow_id=flow_uuid,
                    sequence_number=seq_num,
                    src_node=flow_src,
                    dst_node=node,
                )
            end
        end

        # send qdatagrams
        for uuid in current_flows
            while qdatagrams_in_flight[uuid] < WINDOW && qdatagrams_in_flight[uuid] < pairs_left_to_fulfill[uuid]
                qdatagrams_in_flight[uuid] += 1
                dst        = destination[uuid]
                seq_num    = qdatagrams_sent[uuid] += 1
                start_time = now(sim)::Float64
                corrections = 0 # TODO implement Pauli corrections
                qdatagram = QDatagram(uuid, node, dst, corrections, seq_num, start_time)
                put!(net[node], qdatagram)
            end
        end

        # wait until we have received a message
        @yield onchange(mb)
    end
end


@resumable function (prot::NetworkNodeController)()
    (;sim, net, node) = prot
    mb = messagebuffer(net, node)
    datagrams_in_waiting = Dict{Tuple{Int,Int},Tuple{Tag,Int}}() # keyed by flow_uuid, seq_num; storing datagram tag and next hop
    while true
        workwasdone = true
        while workwasdone
            workwasdone = false

            incoming_qdatagram = querydelete!(mb, QDatagram, ❓, ❓, !=(node), ❓, ❓, ❓)
            if !isnothing(incoming_qdatagram)
                workwasdone = true
                _, flow_uuid, flow_src, flow_dst, corrections, seq_num, start_time = incoming_qdatagram.tag
                nexthop = first(Graphs.a_star(net.graph, node, flow_dst::Int)).dst
                request = LinkLevelRequest(flow_uuid, seq_num, nexthop)
                datagrams_in_waiting[(flow_uuid, seq_num)] = (incoming_qdatagram.tag, nexthop)
                put!(mb, request)
            end

            # Check for LinkLevelReply messages
            # TODO have timeouts on how long to wait for a reply
            llreply = querydelete!(mb, LinkLevelReply, ❓, ❓, ❓)
            if !isnothing(llreply)
                workwasdone = true
                _, flow_uuid, seq_num, memory_slot = llreply.tag
                # Find the corresponding QDatagram that matches this reply
                queued_tag, nexthop = pop!(datagrams_in_waiting, (flow_uuid, seq_num))
                # Process the entanglement and forward the datagram
                _, flow_uuid, flow_src, flow_dst, corrections, seq_num, start_time = queued_tag

                # Perform entanglement swapping
                if node == flow_src
                    put!(net[node], LinkLevelReplyAtSource(flow_uuid, seq_num, memory_slot))
                else
                    # Find the corresponding LinkLevelReplyAtHop message from the previous hop (when the current node was the destination node)
                    llreply_at_destination = querydelete!(mb, LinkLevelReplyAtHop, flow_uuid, seq_num, ❓)
                    @assert !isnothing(llreply_at_destination) "No LinkLevelReplyAtHop message found for flow $(flow_uuid), sequence $(seq_num) at node $(node)"
                    _, _, _, memory_slot_at_destination = llreply_at_destination.tag
                    swapcircuit = LocalEntanglementSwap()
                    xmeas, zmeas = swapcircuit(net[node,memory_slot], net[node,memory_slot_at_destination])
                    # TODO: use xmeas and zmeas to add a correction to the datagram
                end

                # Forward the datagram to the next node in the path
                new_qdatagram = QDatagram(flow_uuid, flow_src, flow_dst, corrections, seq_num, start_time)
                put!(channel(net, node=>nexthop; permit_forward=false), new_qdatagram)
            end
        end

        # Wait until we have received a message
        @yield onchange(mb)
    end
end


@resumable function _link_handle_request(sim, net, nodeA, nodeB, originator_node, destination_node, flow_uuid, seq_num, link_resource)
    @yield lock(link_resource)
    entangler = EntanglerProt(;
        sim, net,
        nodeA, nodeB,
        tag=nothing,
        rounds=1, attempts=-1, success_prob=1.0, attempt_time=1.0 # TODO parameterize the link time and quality
    )
    # TODO have a timeout on how long to wait for an entangler to complete
    proc = @process entangler()
    _, slotA, _, slotB = @yield proc
    unlock(link_resource)
    # slotA is at nodeA, slotB is at nodeB.
    # Map to originator/destination so each node gets its own local slot.
    originator_slot, destination_slot = if originator_node == nodeA
        slotA, slotB
    else
        slotB, slotA
    end
    put!(net[originator_node], LinkLevelReply(flow_uuid=flow_uuid, seq_num=seq_num, memory_slot=originator_slot))
    put!(net[destination_node], LinkLevelReplyAtHop(flow_uuid=flow_uuid, seq_num=seq_num, memory_slot=destination_slot))
end

@resumable function (prot::LinkController)()
    (;sim, net, nodeA, nodeB) = prot
    mbA = messagebuffer(net, nodeA)
    mbB = messagebuffer(net, nodeB)
    link_resource = Resource(sim, 1)

    while true
        workwasdone = true
        while workwasdone
            workwasdone = false

            llrequestA = querydelete!(mbA, LinkLevelRequest, ❓, ❓, nodeB)
            llrequest, originator_node, destination_node = if isnothing(llrequestA)
                querydelete!(mbB, LinkLevelRequest, ❓, ❓, nodeA), nodeB, nodeA
            else
                llrequestA, nodeA, nodeB
            end
            if !isnothing(llrequest)
                workwasdone = true
                @assert originator_node != destination_node "LinkController $(nodeA) $(nodeB) has a link request with originator node $(originator_node) equal to the destination node $(destination_node)"
                _, flow_uuid, seq_num, remote_node = llrequest.tag
                @process _link_handle_request(sim, net, nodeA, nodeB, originator_node, destination_node, flow_uuid, seq_num, link_resource)
            end
        end
        # wait until we have received a message
        @yield (onchange(mbA) | onchange(mbB))
    end
end

end # module
