######################################################################
# Pia Link - Pia's standard transport protocol.
# -----
# Licensed under Apache 2.0
module LinkProtocol
using ..Futures
using ..Lockables
using ..Pia
using ..Packets
using ..PacketFactories
using ..Applications
using Base.Threads
using ExtraFun
using Sockets


# ===== UDPSocket Wrapper =====
export Socket
struct Socket
    raw::UDPSocket
    endpoint::Endpoint
end
function Socket(endpoint::Endpoint)
    socket = UDPSocket()
    bind(socket, endpoint.host, endpoint.port)
    Socket(socket, endpoint)
end
Socket(host::IPAddr, port::Integer) = Socket(Endpoint(host, port))

Base.write(socket::Socket, remote::Endpoint, data) = send(socket.raw, remote.host, remote.port, data)
Base.read(socket::Socket) = recvfrom(socket.raw)

Base.close(socket::Socket) = close(socket.raw)
Base.isopen(socket::Socket) = isopen(socket.raw)


# ===== Protocol =====
export LinkPeer
mutable struct LinkPeer <: Peer
    endpoint::Endpoint
    deadline::Float64
end


export PiaLink
mutable struct PiaLink <: TransportProtocol
    num_desired_connections::UInt16
    num_max_connections::UInt16
    peer_lifetime::Float64
    
    socket::Socket
    wakeup_socket::Socket
    bootstrap::Set{Endpoint}
    connections::Lockable{Dict{Endpoint, Peer}}
    expectance::Lockable{Dict{Endpoint, Future{Vector{UInt8}}}}
    terminate::Atomic{Bool}
end
function PiaLink(
        host::IPAddr = IPv6(0),
        port::Integer = DEFAULT_PORT;
        num_desired_connections::Integer = Pia.DEFAULT_DESIRED_CONNECTIONS,
        num_max_connections::Integer = Pia.DEFAULT_CONNECTION_HARDCAP,
        peer_lifetime::Real = Pia.DEFAULT_PEER_LIFETIME,
        wakeup_port::Integer = randport(50000)
    )
    PiaLink(
        num_desired_connections,
        num_max_connections,
        peer_lifetime,
        
        Socket(host, port),
        Socket(loopback(host), wakeup_port),
        Set(),
        Lockable(Dict{Endpoint, Peer}()),
        Lockable(Dict{Endpoint, Future{Vector{UInt8}}}()),
        Atomic{Bool}(false)
    )
end


Pia.endpoint(link::PiaLink) = link.socket.endpoint
Pia.endpoint(peer::LinkPeer) = peer.endpoint


Pia.connect(link::PiaLink, app::AbstractApplication, host::IPAddr, port::Integer = Pia.DEFAULT_PORT) = Pia.connect(link, app, Endpoint(host, port))
function Pia.connect(link::PiaLink, app::AbstractApplication, endpoint::Endpoint; timeout::Real = Pia.DEFAULT_TIMEOUT)::Peer
    if haskey(link.connections[], endpoint)
        lock(link.connections) do
            if haskey(link.connections[], endpoint)
                return link.connections[][endpoint]
            end
        end
    end
    
    @with ch = PeerChannel(link, endpoint) begin
        handshake_id = rand(UInt8)
        write(ch, packetbytes(Packets.HandshakeStart, appname(app), appversion(app), handshake_id))
        
        (pktcls, pktid), response_id = packetcontent(read(ch, timeout=timeout))
        if pktcls !== Packets.ConnectivityClass || pktid !== Packets.HandshakeResponse || response_id != handshake_id+0x1
            throw(ConnectionError("invalid handshake response"))
        end
        
        write(ch, packetbytes(Packets.HandshakeConfirm))
    end
    
    peer = LinkPeer(endpoint, time() + Pia.DEFAULT_PEER_LIFETIME)
    lock(link.connections) do
        link.connections[][endpoint] = peer
    end
    return peer
end
function Pia.connect(link::PiaLink, app::AbstractApplication)::Vector{Peer}
    if isempty(link.bootstrap)
        throw(StateError("no known bootstrap peers"))
    end
    
    peers = Peer[]
    errors = []
    for remote ∈ link.bootstrap
        try
            push!(peers, Pia.connect(link, app, remote))
        catch ex
            push!(errors, ex) # TODO: Capture stack trace
        end
    end
    if isempty(peers)
        throw(ConnectionError("failed to connect to bootstrap peers", errors))
    end
    return peers
end

Pia.isconnected(link::PiaLink) = !isempty(link.connections[])

Pia.bootstrap(link::PiaLink, ip::IPAddr, port::Integer = Pia.DEFAULT_PORT) = Pia.bootstrap(link, Endpoint(ip, port))
function Pia.bootstrap(link::PiaLink, endpoint::Endpoint)
    push!(link.bootstrap, endpoint)
    link
end

Base.write(link::PiaLink, ::AbstractApplication, endpoint::Endpoint, data) = write(link, endpoint, packetbytes(Packets.AppData, data))
Base.write(link::PiaLink, endpoint::Endpoint, data) = write(link.socket, endpoint, data)

function Base.read(link::PiaLink, ::AbstractApplication, endpoint::Endpoint; kwargs...)
    (clsid, pktid), bytes = packetcontent(read(link, endpoint; kwargs...))
    @assert clsid === Packets.DataClass
    @assert pktid === Packets.AppData
    return bytes
end
function Base.read(link::PiaLink, endpoint::Endpoint; timeout::Real = Pia.DEFAULT_TIMEOUT)
    future = Future{Vector{UInt8}}()
    lock(link.expectance) do
        link.expectance[][endpoint] = future
    end
    return wait(future, timeout=timeout)
end

function Base.close(link::PiaLink)
    if !atomic_cas!(link.terminate, false, true)
        write(link.wakeup_socket, Endpoint(loopback(link.socket.endpoint.host), link.socket.endpoint.port), packetbytes(Packets.Wakeup))
        
        lock(link.connections) do
            bytes = packetbytes(Packets.Goodbye)
            for peer ∈ values(link.connections[])
                write(link.socket, peer.endpoint, bytes)
            end
        end
        
        close(link.wakeup_socket)
        close(link.socket)
    end
end

function Base.show(io::IO, link::PiaLink)
    endpoint    = link.socket.endpoint
    num_conns   = length(link.connections[])
    num_max     = link.num_max_connections
    num_desired = link.num_desired_connections
    
    print(io, "Pia Link @ $endpoint($num_conns/$num_max peers ($num_desired desired))")
end


# ===== PeerInterface Implementation =====
struct LinkPeerInterface <: PeerInterface
    link::PiaLink
end

Pia.peers(link::PiaLink) = LinkPeerInterface(link)

function Base.getindex(iface::LinkPeerInterface, endpoint::Endpoint)
    lock(iface.link.connections) do
        return iface.link.connections[][endpoint]
    end
end

Base.haskey(iface::LinkPeerInterface, endpoint::Endpoint) = haskey(iface.link.connections[], endpoint)
Base.lock(cb, iface::LinkPeerInterface) = lock(cb, iface.link.connections)
Base.lock(iface::LinkPeerInterface) = lock(iface.link.connections)
Base.unlock(iface::LinkPeerInterface) = unlock(iface.link.connections)

function refresh(iface::LinkPeerInterface, endpoint::Endpoint)
    link = iface.link
    lock(link.connections) do
        if haskey(link.connections[], endpoint)
            link.connections[][endpoint].deadline = time()+link.peer_lifetime
        end
    end
    return iface
end

# ===== Pia Link Loop =====
# --- Connection Multiplexing Solution ---
function Pia.loop(link::PiaLink, app::AbstractApplication)
    while !link.terminate[]
        remote, packet = @await yield() recvfrom(link.socket.raw)
        if isself(link, remote) continue end
        
        try
            _, pktid = peekpacket(packet)
            Applications.process_packet(Ident{pktid}(), app, PeerChannel(link, remote), packet)
        catch ex
            println(stderr, ex)
        end
    end
end
isself(link::PiaLink, other::Endpoint) = isself(link.socket.endpoint, other) || isself(link.wakeup_socket.endpoint, other)
isself(me::Endpoint, other::Endpoint) = me.host == other.host && me.port == other.port


# ===== Packet Integration =====
# --- Expectancy Anti-DOS Solution ---
function Applications.process_packet(::Ident, ::AbstractApplication, link::PiaLink, endpoint::Endpoint, bytes)
    lock(link.expectance) do
        if haskey(link.expectance[], endpoint)
            link.expectance[][endpoint][] = bytes
        end
    end
    nothing
end

# --- Handshake Start Packet ---
function Applications.process_packet(::Ident{Packets.HandshakeStart}, app::AbstractApplication, link::PiaLink, remote::Endpoint, bytes)
    # Validate packet
    remote_appname, remote_appversion, handshake_id = packetcontent(bytes, without_id=true)
    if remote_appname != appname(app)
        return LogConnectivity("invalid app name $remote_appname from peer $(ch.endpoint)")
    end
    if remote_appversion != appversion(app)
        return LogConnectivity("invalid app version $remote_appversion from peer $(ch.endpoint)")
    end
    
    # Send HandshakeResponse
    write(link, remote, packetbytes(Packets.HandshakeResponse, handshake_id+0x1))
    
    # Anticipate HandshakeConfirm (async w/ timeout)
    @async begin
        (rcls, rpkt), _ = packetcontent(read(link, remote; timeout=Pia.DEFAULT_TIMEOUT))
        if rcls === Packets.ConnectivityClass && rpkt === Packets.HandshakeConfirm
            lock(link.connections) do
                link.connections[][remote] = LinkPeer(remote, time()+link.peer_lifetime)
            end
        else
            LogConnectivity("invalid handshake confirm from peer $(ch.endpoint)")
        end
    end
end

# --- Goodbye Packet ---
function Applications.process_packet(::Ident{Packets.Goodbye}, ::AbstractApplication, link::PiaLink, remote::Endpoint, _)
    lock(link.connections) do
        delete!(link.connections[], remote)
    end
end

end # module LinkProtocol
