#include "ipv4.h"
#include "icmp.h"

namespace ipv4 {

static constexpr ip4_addr IP_BROADCAST_ALL = {255, 255, 255, 255};

uint16_t checksum(const void* data, size_t len) {
    auto p = static_cast<const uint8_t*>(data);
    uint32_t sum = 0;
    for (size_t i = 0; i < len - 1; i += 2)
        sum += (p[i] << 8) | p[i + 1];
    if (len & 1)
        sum += p[len - 1] << 8;
    while (sum >> 16)
        sum = (sum & 0xFFFF) + (sum >> 16);
    return __builtin_bswap16(~sum);
}

static bool ip_match_or_broadcast(const ip4_addr& dst, const ip4_addr& our_ip, const ip4_addr& subnet_broadcast) {
    return dst == our_ip || dst == IP_BROADCAST_ALL || dst == subnet_broadcast;
}

void handle(std::span<const uint8_t> frame, span_writer& tx,
            eth::mac_addr our_mac, ip4_addr our_ip, ip4_addr subnet_broadcast,
            std::function<void(std::span<const uint8_t>)> send_raw,
            std::function<void(std::span<const uint8_t>, span_writer&)> handle_udp) {
    if (frame.size() < sizeof(header)) return;
    auto& ip = *reinterpret_cast<const header*>(frame.data());
    if ((ip.ver_ihl >> 4) != 4) return;

    switch (ip.protocol) {
    case 1:
        if (!ip_match_or_broadcast(ip.dst, our_ip, subnet_broadcast))
            return;
        icmp::handle(frame, tx, our_mac, our_ip, send_raw);
        break;
    case 17:
        if (!ip_match_or_broadcast(ip.dst, our_ip, subnet_broadcast))
            return;
        handle_udp(frame, tx);
        break;
    }
}

} // namespace ipv4