#include "ipv4.h"
#include "icmp.h"
#include "igmp.h"
#include "net.h"
#include "udp.h"
#include "parse_buffer.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(const ip4_addr& dst, const ip4_addr& our_ip) {
    return dst == our_ip || dst == IP_BROADCAST_ALL || dst == SUBNET_BROADCAST || igmp::is_member(dst);
}

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

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

    switch (ip->protocol) {
    case 1:
        if (!ip_match(ip->dst, ns.ip))
            return;
        icmp::handle(frame, tx, ns.mac, ns.ip);
        break;
    case 2:
        igmp::handle(frame, tx, ns.mac, ns.ip);
        break;
    case 17:
        if (!ip_match(ip->dst, ns.ip))
            return;
        udp::handle(frame, tx);
        break;
    }
}

__attribute__((constructor))
static void register_ethertype() {
    net_register_ethertype(eth::ETH_IPV4, handle);
}

} // namespace ipv4