module QTCP

using QuantumSavory
using QuantumSavory: Tag
using QuantumSavory.ProtocolZoo
using QuantumSavory.ProtocolZoo: AbstractProtocol
using QuantumSavory.CircuitZoo: LocalEntanglementSwap
import ConcurrentSim
using ConcurrentSim: Simulation, @yield, timeout, @process, now, Process
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
    "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
    "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
    "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
    "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
    "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
    "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
    "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
    "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
    "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

"""
$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

"""
$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 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
        # check if there is an approved flow and start it
        flow = querydelete!(mb, Flow, node, ❓, ❓, ❓)
        if !isnothing(flow)
            _, _, 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 "[$(now(sim))]: flow $(uuid) started"
        end

        # check if there are datagram acknowledgements
        success = querydelete!(mb, QDatagramSuccess, ❓, ❓, ❓)
        if !isnothing(success)
            # TODO implement drop detection and window modification
            _, flow_uuid, seq_num, start_time = success.tag
            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, ❓, ❓, ❓)
            @assert !isnothing(link_reply) "No LinkLevelReplyAtSource message found after the success of flow $(flow_uuid), sequence $(seq_num) at node $(node)"
            _, flow_uuid, seq_num, 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 "[$(now(sim))]: datagram success notification from flow $(flow_uuid) pair $(seq_num) returned to start 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)
                current_time = now(sim)
                @debug "[$(current_time)]: flow $(flow_uuid) completed and deallocated"
            end
        end

        # check if we just received a qdatagram for which we are the flow destination
        qdatagram = querydelete!(mb, QDatagram, ❓, ❓, node, ❓, ❓, ❓)
        if !isnothing(qdatagram)
            # 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
            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, ❓, ❓, ❓)
            @assert !isnothing(link_reply) "No LinkLevelReplyAtHop message found after the success of flow $(flow_uuid), sequence $(seq_num) at node $(node)"
            _, flow_uuid, seq_num, memory_slot = link_reply.tag
            pair_end = QTCPPairEnd(;
                flow_uuid,
                flow_src=0, # TODO we do not actually know the source node, it is not something that was kept track of
                flow_dst=node,
                seq_num,
                memory_slot,
                start_time
            )
            put!(net[node], pair_end)
            @debug "[$(now(sim))]: datagram from flow $(flow_uuid) pair $(seq_num) reached final destination"
        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)
                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 wait(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 and next hop
    while true
        qdatagram = querydelete!(mb, QDatagram, ❓, ❓, !=(node), ❓, ❓, ❓)
        if !isnothing(qdatagram)
            _, flow_uuid, flow_src, flow_dst, corrections, seq_num, start_time = qdatagram.tag
            nexthop = first(Graphs.a_star(net.graph, node, flow_dst)).dst
            request = LinkLevelRequest(flow_uuid, seq_num, nexthop)
            datagrams_in_waiting[(flow_uuid, seq_num)] = (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)
            _, flow_uuid, seq_num, memory_slot = llreply.tag
            # Find the corresponding QDatagram that matches this reply
            qdatagram, 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 = qdatagram

            # 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, ❓, ❓, ❓)
                @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

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


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

    while true
        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)
            @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
            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
            # Create the reply with the flow information and memory slot
            reply = LinkLevelReply(
                flow_uuid=flow_uuid,
                seq_num=seq_num,
                memory_slot=slotA
            )
            reply_at_destination = LinkLevelReplyAtHop(
                flow_uuid=flow_uuid,
                seq_num=seq_num,
                memory_slot=slotB
            )

            # Put the reply in the appropriate node's message buffer
            put!(net[originator_node], reply)
            put!(net[destination_node], reply_at_destination)
        end
        # wait until we have received a message
        @yield (wait(mbA) | wait(mbB))
    end
end

end # module
