#ifndef TOS__NET__UDP_HPP__
#define TOS__NET__UDP_HPP__

#include <cstdint>

#include <memory>

#include "net/ipv4.hpp"
#include "utils/span.hpp"

namespace net {

struct UDPHeader
{
    NetField<+0, 2, std::uint16_t> srcPort;
    NetField<+2, 2, std::uint16_t> dstPort;
    NetField<+4, 2, std::uint16_t> length;
    NetField<+6, 2, std::uint16_t> checksum;

    explicit UDPHeader(std::uint8_t data[])
        : srcPort(data), dstPort(data), length(data), checksum(data)
    {
    }

    static constexpr std::ptrdiff_t size()
    {
        using lastField = decltype(checksum);
        return lastField::offset + lastField::size;
    }
};

class UDPSocket;
class UDPProvider;

class UDPHandler
{
public:
    virtual ~UDPHandler() = default;

public:
    virtual void handleUDPMessage(UDPSocket &socket,
                                  span<std::uint8_t> msg) = 0;
};

class UDPSocket
{
    friend class UDPProvider;

public:
    UDPSocket(UDPProvider &backend);
    virtual ~UDPSocket() = default;

public:
    virtual void handleUDPMessage(span<std::uint8_t> msg);
    virtual void send(span<const std::uint8_t> buf);
    virtual void disconnect();

private:
    NetOrder<std::uint16_t> remotePort;
    NetOrder<std::uint32_t> remoteIP;
    NetOrder<std::uint16_t> localPort;
    NetOrder<std::uint32_t> localIP;
    UDPProvider &backend;
    UDPHandler *handler;
    bool listening;
};

class UDPProvider : IPHandler
{
protected:
    std::unique_ptr<UDPSocket> sockets[65536];
    std::uint16_t numSockets;
    std::uint16_t freePort;

public:
    UDPProvider(IPProvider &backend);

public:
    virtual bool onIPReceived(NetOrder<std::uint32_t> srcIP,
                              NetOrder<std::uint32_t> dstIP,
                              span<std::uint8_t> msg) override;

    virtual UDPSocket & connect(NetOrder<std::uint32_t> ip, std::uint16_t port);
    virtual UDPSocket & listen(std::uint16_t port);
    virtual void disconnect(UDPSocket &socket);
    virtual void send(UDPSocket &socket, span<const std::uint8_t> buf);

    virtual void bind(UDPSocket &socket, UDPHandler* handler);
};

}

#endif // TOS__NET__UDP_HPP__