#pragma once
#include <cstdint>
#include <span>
#include "eth.h"
#include "ipv4.h"
#include "span_writer.h"

namespace udp {

struct __attribute__((packed)) header {
    uint16_t src_port;
    uint16_t dst_port;
    uint16_t length;
    uint16_t checksum;
};
static_assert(sizeof(header) == 8);

struct address {
    // mac is carried here until we grow an ARP cache; once we do, a
    // destination mac comes from resolving ip and this field goes away.
    eth::mac_addr mac;
    ipv4::ip4_addr ip;
    uint16_t port;
};

template <typename Buf>
void prepend(Buf& buf, const eth::mac_addr& dst_mac, const eth::mac_addr& src_mac,
             ipv4::ip4_addr src_ip, ipv4::ip4_addr dst_ip,
             uint16_t src_port, uint16_t dst_port,
             size_t payload_len, uint8_t ttl = 64) {
    auto* u = buf.template prepend<header>();
    u->src_port = src_port;
    u->dst_port = dst_port;
    u->length = __builtin_bswap16(sizeof(header) + payload_len);
    u->checksum = 0;
    ipv4::prepend(buf, dst_mac, src_mac, src_ip, dst_ip, 17, sizeof(header) + payload_len, ttl);
}

void handle(std::span<const uint8_t> frame, span_writer& tx);

using port_handler = void (*)(std::span<const uint8_t> payload, const address& from);
void register_port(uint16_t port_be, port_handler fn);

} // namespace udp