#include "udp.h"
#include <array>
#include "eth.h"
#include "ipv4.h"
#include "parse_buffer.h"
#include "debug_log.h"

namespace udp {

struct port_entry {
    uint16_t port_be;
    port_handler fn;
};
static std::array<port_entry, 8> port_handlers;
static size_t port_handler_count = 0;

void register_port(uint16_t port_be, port_handler fn) {
    if (port_handler_count >= port_handlers.size()) {
        dlogf("udp::register_port overflow: port=%u dropped", __builtin_bswap16(port_be));
        return;
    }
    port_handlers[port_handler_count++] = {port_be, fn};
}

void handle(std::span<const uint8_t> frame, span_writer& tx) {
    parse_buffer pb(frame);
    auto* eth_hdr = pb.consume<eth::header>();
    auto* ip = pb.consume<ipv4::header>();
    if (!ip) return;
    if (!ipv4::addressed_to_us(ip->dst)) return;

    size_t options_len = ip->header_len() - sizeof(ipv4::header);
    if (options_len > 0 && !pb.skip(options_len)) return;

    auto* uhdr = pb.consume<header>();
    if (!uhdr) return;

    size_t udp_len = __builtin_bswap16(uhdr->length);
    if (udp_len < sizeof(header)) return;
    size_t payload_len = udp_len - sizeof(header);
    if (pb.remaining_size() < payload_len) return;

    for (size_t i = 0; i < port_handler_count; i++) {
        if (port_handlers[i].port_be == uhdr->dst_port) {
            address from{eth_hdr->src, ip->src, uhdr->src_port};
            port_handlers[i].fn(pb.remaining().subspan(0, payload_len), from);
            break;
        }
    }
}

void init() {
    ipv4::register_protocol(17, handle);
}

} // namespace udp