#include "net/udp.hpp"

#include <cstdint>

#include <memory>
#include <utility>

using namespace net;

UDPSocket::UDPSocket(UDPProvider &backend)
    : backend(backend), handler(nullptr), listening(false)
{
}

void
UDPSocket::handleUDPMessage(span<std::uint8_t> msg)
{
    if (handler != nullptr) {
        handler->handleUDPMessage(*this, msg);
    }
}

void
UDPSocket::send(span<const std::uint8_t> buf)
{
    backend.send(*this, buf);
}

void
UDPSocket::disconnect()
{
    backend.disconnect(*this);
}

UDPProvider::UDPProvider(IPProvider &backend)
    : IPHandler(backend, 0x11), numSockets(0), freePort(1024)
{
}

bool
UDPProvider::onIPReceived(NetOrder<std::uint32_t> srcIP,
                          NetOrder<std::uint32_t> dstIP,
                          span<std::uint8_t> msg)
{
    if (msg.size() < UDPHeader::size()) {
        return false;
    }

    UDPHeader udpHdr(msg.data());

    UDPSocket *socket = nullptr;
    for (std::uint16_t i = 0U; i < numSockets && socket == nullptr; ++i) {
        if (sockets[i]->localPort == udpHdr.dstPort &&
            sockets[i]->localIP == dstIP &&
            sockets[i]->listening) {
            socket = sockets[i].get();
            socket->listening = false;
            socket->remotePort = udpHdr.srcPort;
            socket->remoteIP = srcIP;
        } else if (sockets[i]->localPort == udpHdr.dstPort &&
                   sockets[i]->localIP == dstIP &&
                   sockets[i]->remotePort == udpHdr.srcPort &&
                   sockets[i]->remoteIP == srcIP) {
            socket = sockets[i].get();
        }
    }

    if (socket != nullptr) {
        socket->handleUDPMessage(msg.after(udpHdr.size()));
    }

    return false;
}

UDPSocket &
UDPProvider::connect(NetOrder<std::uint32_t> ip, std::uint16_t port)
{
    std::unique_ptr<UDPSocket> socket(new UDPSocket(*this));

    socket->remotePort = toNetOrder(port);
    socket->remoteIP = ip;
    socket->localPort = toNetOrder(freePort++);
    socket->localIP = backend.getIPAddress();

    sockets[numSockets++] = std::move(socket);
    return *sockets[numSockets - 1U];
}

UDPSocket &
UDPProvider::listen(std::uint16_t port)
{
    std::unique_ptr<UDPSocket> socket(new UDPSocket(*this));

    socket->listening = true;
    socket->localPort = toNetOrder(port);
    socket->localIP = backend.getIPAddress();

    sockets[numSockets++] = std::move(socket);
    return *sockets[numSockets - 1U];
}

void
UDPProvider::disconnect(UDPSocket &socket)
{
    for (std::uint16_t i = 0; i < numSockets; ++i) {
        if (sockets[i].get() == &socket) {
            sockets[i] = std::move(sockets[--numSockets]);
            break;
        }
    }
}

void
UDPProvider::send(UDPSocket &socket, span<const std::uint8_t> buf)
{
    std::uint16_t totalLength = buf.size() + UDPHeader::size();
    std::unique_ptr<std::uint8_t[]> buffer(new std::uint8_t[totalLength]);

    UDPHeader msg(buffer.get());

    msg.srcPort = socket.localPort;
    msg.dstPort = socket.remotePort;
    msg.length = toNetOrder(totalLength);

    std::uint8_t *payload = buffer.get() + UDPHeader::size();
    for (std::uint16_t i = 0U; i < buf.size(); ++i) {
        payload[i] = buf[i];
    }

    msg.checksum = toNetOrder<std::uint16_t>(0);
    IPHandler::send(socket.remoteIP, { buffer.get(), totalLength });
}

void
UDPProvider::bind(UDPSocket &socket, UDPHandler *handler)
{
    socket.handler = handler;
}