#include <devif/queue_interface_backend.h>
#include <devif/backends/net/enet_devif.h>
#include <aos/aos.h>
#include <driverkit/driverkit.h>
#include <dev/imx8x/enet_dev.h>

#include <aos/ump_binding.h>
#include <aos/ump_chan.h>
#include <aos/ump_net.h>

#include <netutil/etharp.h>
#include <netutil/ip.h>
#include <netutil/icmp.h>
#include <netutil/udp.h>
#include <netutil/tcp.h>
#include <netutil/checksum.h>
#include <netutil/types.h>

#include "enet.h"

// https://datatracker.ietf.org/doc/html/rfc6335#section-6
#define EPHEMERAL_PORT_START 49152
#define EPHEMERAL_PORT_END 65535

extern struct enet_driver_state *st;

static void write_eth_ip_header (
    struct devq_buf *buf, void *vaddr,
    uint64_t eth_dst, uint32_t ip_dest, uint8_t proto, size_t payload_len
);

static void tx_alloc (struct tx_alloc_queue_entry *entry, tx_alloc_callback_t cb, void *cb_arg) {
    errval_t err = DEVQ_ERR_QUEUE_EMPTY;
    struct devq_buf buf;

    // If the queue is not empty, always append to the queue, so that FIFO order is maintained
    if (st->tx_alloc_queue_head == NULL) {
        err = devq_dequeue((struct devq*) st->txq, &buf.rid, &buf.offset,
            &buf.length, &buf.valid_data, &buf.valid_length,
            &buf.flags);
    }
    if (err_is_fail(err)) {
        if (st->tx_init_i < st->txq->size) {
            buf.rid = st->tx_rid;
            buf.offset = st->tx_init_i * 2048;
            buf.length = 2048;
            buf.valid_data = 0;
            buf.valid_length = 2048;
            buf.flags = 0;
            st->tx_init_i++;
        } else {
            entry->cb = cb;
            entry->cb_arg = cb_arg;
            entry->next = NULL;
            if (st->tx_alloc_queue_tail == NULL) {
                st->tx_alloc_queue_head = entry;
            } else {
                st->tx_alloc_queue_tail->next = entry;
            }
            st->tx_alloc_queue_tail = entry;
            return;
        }
    }

    struct region_entry *region_entry = enet_get_region(st->txq, buf.rid);
    assert(region_entry != NULL);
    void *vaddr = (void*)region_entry->mem.vbase + buf.offset;
    cb(cb_arg, &buf, vaddr);
}

static void tx_send (struct devq_buf *buf) {
    errval_t err;
    err = devq_enqueue((struct devq*) st->txq, buf->rid, buf->offset,
                        buf->length, buf->valid_data, buf->valid_length,
                        buf->flags);
    assert(err_is_ok(err));
}

static void rx_release (struct devq_buf *buf) {
    errval_t err;
    err = devq_enqueue((struct devq*) st->rxq, buf->rid, buf->offset,
                        buf->length, buf->valid_data, buf->valid_length,
                        buf->flags);
    assert(err_is_ok(err));
}

// ARP: https://datatracker.ietf.org/doc/html/rfc826

static void arp_send (
    struct devq_buf *buf, void *vaddr,
    uint64_t eth_dst, uint32_t ip_dest,
    uint16_t op
) {
    buf->valid_data = 0;
    buf->valid_length = ETH_HLEN + sizeof(struct arp_hdr);

    struct eth_hdr *eth_hdr = vaddr;
    eth_hdr->dst = eth_addr_wr(eth_dst);
    eth_hdr->src = eth_addr_wr(st->mac);
    eth_hdr->type = uint16_wr(ETH_TYPE_ARP);

    struct arp_hdr *arp_hdr = vaddr + ETH_HLEN;
    arp_hdr->hwtype = uint16_wr(ARP_HW_TYPE_ETH);
    arp_hdr->proto = uint16_wr(ETH_TYPE_IP);
    arp_hdr->hwlen = ETH_ADDR_LEN;
    arp_hdr->protolen = IP_ADDR_LEN;
    arp_hdr->opcode = uint16_wr(op);
    arp_hdr->eth_src = eth_addr_wr(st->mac);
    arp_hdr->ip_src = uint32_wr(st->my_ip);
    arp_hdr->eth_dst = eth_addr_wr(eth_dst);
    arp_hdr->ip_dst = uint32_wr(ip_dest);

    tx_send(buf);
}

static void arp_reply (void *cb_arg, struct devq_buf *buf, void *vaddr) {
    struct arp_tx_queue_entry *entry = &st->arp.tx_queue[st->arp.tx_queue_tail];
    assert(cb_arg == entry);

    arp_send(buf, vaddr, entry->eth_dst, entry->ip_dst, ARP_OP_REP);

    st->arp.tx_queue_tail = (st->arp.tx_queue_tail + 1) & (ARP_TX_QUEUE_LEN - 1);
    st->arp.tx_queue_len--;
}

static void arp_handle (struct devq_buf *buf, const void *vaddr, const struct eth_hdr *eth_hdr) {
    ENET_DEBUG("Received ARP packet\n");

    if (buf->valid_length < sizeof(struct arp_hdr)) {
        ENET_WARN("Received ARP packet too small\n");
        rx_release(buf);
        return;
    }
    const struct arp_hdr *arp_hdr = vaddr + buf->valid_data;
    uint16_t opcode = uint16_rd(arp_hdr->opcode);
    if (
        uint16_rd(arp_hdr->hwtype) != ARP_HW_TYPE_ETH ||
        uint16_rd(arp_hdr->proto) != ETH_TYPE_IP ||
        arp_hdr->hwlen != ETH_ADDR_LEN ||
        arp_hdr->protolen != IP_ADDR_LEN ||
        (opcode != ARP_OP_REQ && opcode != ARP_OP_REP)
    ) {
        ENET_WARN("Unknown ARP packet type\n");
        rx_release(buf);
        return;
    }

    uint32_t ip_src = uint32_rd(arp_hdr->ip_src);

    if ((st->my_ip & st->net_mask) != (ip_src & st->net_mask)) {
        ENET_WARN("Received ARP packet for address outside my network\n");
        rx_release(buf);
        return;
    }

    uint64_t eth_src = eth_addr_rd(arp_hdr->eth_src);
    uint64_t eth_old = (uintptr_t)collections_hash_find(st->arp.table, ip_src);
    // check if sender protocol address is in translation table, if so update entry
    if (eth_old != 0 && eth_src != eth_old) {
        ENET_WARN("MAC address of %d.%d.%d.%d has changed\n",
            arp_hdr->ip_src.val[0],
            arp_hdr->ip_src.val[1],
            arp_hdr->ip_src.val[2],
            arp_hdr->ip_src.val[3]
        );
        collections_hash_delete(st->arp.table, ip_src);
        collections_hash_insert(st->arp.table, ip_src, (void*)eth_src);
    }

    if (uint32_rd(arp_hdr->ip_dst) == st->my_ip) {
        ENET_DEBUG("Received ARP packet addressed to me\n");

        // if sender protocol address was not in translation table, add entry
        if (eth_old == 0) {
            collections_hash_insert(st->arp.table, ip_src, (void*)eth_src);
        }

        if (opcode == ARP_OP_REQ) {
            if (st->arp.tx_queue_len >= ARP_TX_QUEUE_LEN) {
                ENET_WARN("Too many queued ARP replies, not sending reply\n");
            } else {
                struct arp_tx_queue_entry *entry = &st->arp.tx_queue[(st->arp.tx_queue_tail + st->arp.tx_queue_len) & (ARP_TX_QUEUE_LEN - 1)];
                st->arp.tx_queue_len++;
                entry->eth_dst = eth_src;
                entry->ip_dst = ip_src;
                tx_alloc(&entry->tx_entry, arp_reply, entry);
            }
        }
    }

    rx_release(buf);
}

// ICMP: https://datatracker.ietf.org/doc/html/rfc792

static void icmp_reply (void *cb_arg, struct devq_buf *buf, void *vaddr) {
    struct icmp_echo_reply_meta *entry = cb_arg;

    write_eth_ip_header(
        buf, vaddr,
        eth_addr_rd(entry->eth_hdr->src),
        uint32_rd(entry->ip_hdr->src),
        IP_PROTO_ICMP,
        entry->rx_buf.valid_length
    );

    struct icmp_echo_hdr *icmp_echo_hdr = vaddr + ETH_HLEN + IP_HLEN_MIN;
    memcpy(icmp_echo_hdr, entry->icmp_echo_hdr, entry->rx_buf.valid_length);
    icmp_echo_hdr->type = ICMP_ER;

    // Incremental update: https://datatracker.ietf.org/doc/html/rfc1071
    uint32_t chksum = uint16_rd(icmp_echo_hdr->chksum);
    chksum += (ICMP_ECHO << 8) + ((ICMP_ER << 8) ^ 0x0000ffffUL);
    chksum = (chksum >> 16) + (chksum & 0x0000ffffUL);
    icmp_echo_hdr->chksum = uint16_wr(chksum);

    rx_release(&entry->rx_buf);
    simpleslab_free(&st->rx_meta_slab, entry);
    tx_send(buf);
}

static void icmp_handle (struct devq_buf *buf, const void *vaddr, const struct eth_hdr *eth_hdr, const struct ip_hdr *ip_hdr) {
    if (buf->valid_length < 4) {
        ENET_WARN("Received ICMP packet too small\n");
        rx_release(buf);
        return;
    }

    if (inet_checksum(vaddr + buf->valid_data, buf->valid_length) != 0) {
        ENET_WARN("Received ICMP packet with bad checksum\n");
        rx_release(buf);
        return;
    }

    uint8_t type = *(const uint8_t*)(vaddr + buf->valid_data);
    if (type == ICMP_ECHO) {
        if (buf->valid_length < sizeof(struct icmp_echo_hdr)) {
            ENET_WARN("Received ICMP packet too small\n");
            rx_release(buf);
            return;
        }

        const struct icmp_echo_hdr *icmp_echo_hdr = vaddr + buf->valid_data;
        ENET_DEBUG("Received ICMP echo, seq=%d\n", uint16_rd(icmp_echo_hdr->seqno));

        struct icmp_echo_reply_meta *meta = simpleslab_alloc(&st->rx_meta_slab);
        meta->rx_buf = *buf;
        meta->eth_hdr = eth_hdr;
        meta->ip_hdr = ip_hdr;
        meta->icmp_echo_hdr = icmp_echo_hdr;

        tx_alloc(&meta->tx_entry, icmp_reply, meta);
        return;
    } else {
        ENET_WARN("Received ICMP packet with unknown type %d\n", type);
    }

    rx_release(buf);
}

// UDP/TCP checksum

static uint16_t ip_data_checksum (const struct ip_hdr *ip_hdr, const void *data, uint16_t data_len) {
    uint32_t chksum = inet_checksum(data, data_len) ^ 0x0000ffffUL;
    // add pseudo header
    chksum += (ip_hdr->src.val[0] << 8) | ip_hdr->src.val[1];
    chksum += (ip_hdr->src.val[2] << 8) | ip_hdr->src.val[3];
    chksum += (ip_hdr->dest.val[0] << 8) | ip_hdr->dest.val[1];
    chksum += (ip_hdr->dest.val[2] << 8) | ip_hdr->dest.val[3];
    chksum += ip_hdr->proto;
    chksum += data_len;
    chksum = (chksum >> 16) + (chksum & 0x0000ffffUL);
    chksum = (chksum >> 16) + (chksum & 0x0000ffffUL);
    return ~(uint16_t)chksum;
}

// UDP: https://datatracker.ietf.org/doc/html/rfc768

static void udp_recv_ump_callback (void *arg, struct ump_send_queue_entry *entry) {
    struct udp_recv *meta = arg;
    rx_release(&meta->rx_buf);
    simpleslab_free(&st->rx_meta_slab, meta);
}

static void udp_handle (struct devq_buf *buf, const void *vaddr, const struct eth_hdr *eth_hdr, const struct ip_hdr *ip_hdr) {
    if (buf->valid_length < UDP_HLEN) {
        ENET_WARN("Received UDP packet too small\n");
        rx_release(buf);
        return;
    }
    const struct udp_hdr *udp_hdr = vaddr + buf->valid_data;
    uint16_t udp_len = uint16_rd(udp_hdr->len);
    if (udp_len < UDP_HLEN || buf->valid_length < udp_len) {
        ENET_WARN("Received UDP packet too small\n");
        rx_release(buf);
        return;
    }

    if (
        uint16_rd(udp_hdr->chksum) != 0 && // 0 means no checksum
        ip_data_checksum(ip_hdr, udp_hdr, uint16_rd(udp_hdr->len)) != 0
    ) {
        ENET_WARN("Received UDP packet with bad checksum\n");
        rx_release(buf);
        return;
    }

    uint16_t src_port = uint16_rd(udp_hdr->src);
    uint16_t dest_port = uint16_rd(udp_hdr->dest);
    const void *payload = (const void*)udp_hdr + UDP_HLEN;
    uint16_t payload_len = udp_len - UDP_HLEN;
    ENET_DEBUG("Received UDP packet from %d to %d, len %d\n", src_port, dest_port, payload_len);

    struct ump_client *client = collections_hash_find(st->udp.listen_table, dest_port);
    if (client == NULL) {
        // TODO: send ICMP unreachable
        rx_release(buf);
    } else {
        struct udp_recv *meta = simpleslab_alloc(&st->rx_meta_slab);
        meta->rx_buf = *buf;
        meta->ump_header = (struct ump_net_in_header){
            .op = UMP_NET_EV_UDP_RECV,
            .ev = { .udp_recv = {
                .src_ip = uint32_rd(ip_hdr->src),
                .dest_ip = uint32_rd(ip_hdr->dest),
                .src_port = src_port,
                .dest_port = dest_port,
            } }
        };
        ump_send(client->send_chan, &meta->ump_entry,
            sizeof(struct ump_net_in_header), &meta->ump_header,
            payload_len, payload,
            udp_recv_ump_callback, meta);
        return;
    }
}

// TCP: https://datatracker.ietf.org/doc/html/rfc793

static void tcp_recv_ump_callback (void *arg, struct ump_send_queue_entry *entry) {
    struct tcp_recv *meta = arg;
    rx_release(&meta->rx_buf);
    simpleslab_free(&st->rx_meta_slab, meta);
}

static void tcp_reply (void *cb_arg, struct devq_buf *buf, void *vaddr) {
    struct tcp_reply_buf *reply = cb_arg;

    write_eth_ip_header(
        buf, vaddr,
        reply->eth_dst,
        reply->ip_dst,
        IP_PROTO_TCP,
        TCP_HLEN_MIN
    );
    struct ip_hdr *ip_hdr = vaddr + ETH_HLEN;
    struct tcp_hdr *tcp_hdr = vaddr + ETH_HLEN + IP_HLEN_MIN;
    tcp_hdr->src = uint16_wr(reply->port_src);
    tcp_hdr->dest = uint16_wr(reply->port_dst);
    tcp_hdr->seq = uint32_wr(reply->seq);
    tcp_hdr->ack = uint32_wr(reply->ack);
    tcp_hdr->flags = uint16_wr(reply->flags | (5 << 12));
    tcp_hdr->wnd = uint16_wr(0);
    tcp_hdr->chksum = uint16_wr(0);
    tcp_hdr->up = uint16_wr(0);

    uint16_t checksum = ip_data_checksum(ip_hdr, tcp_hdr, TCP_HLEN_MIN);
    tcp_hdr->chksum = uint16_wr(checksum);

    simpleslab_free(&st->tcp.tcp_reply_slab, reply);
    tx_send(buf);
}

static void tcp_handle (struct devq_buf *buf, const void *vaddr, const struct eth_hdr *eth_hdr, const struct ip_hdr *ip_hdr) {
    if (buf->valid_length < TCP_HLEN_MIN) {
        ENET_WARN("Received TCP packet too small\n");
        rx_release(buf);
        return;
    }
    const struct tcp_hdr *tcp_hdr = vaddr + buf->valid_data;
    uint16_t hlen = TCP_HLEN(tcp_hdr);
    if (hlen < TCP_HLEN_MIN || buf->valid_length < hlen) {
        ENET_WARN("Received TCP packet too small\n");
        rx_release(buf);
        return;
    }

    if (ip_data_checksum(ip_hdr, tcp_hdr, buf->valid_length) != 0) {
        ENET_WARN("Received TCP packet with bad checksum\n");
        rx_release(buf);
        return;
    }

    uint16_t flags = uint16_rd(tcp_hdr->flags);
    uint16_t dest_port = uint16_rd(tcp_hdr->dest);
    const void *payload = (const void*)tcp_hdr + hlen;
    uint16_t payload_len = buf->valid_length - hlen;
    ENET_DEBUG("Received TCP packet from %d to %d, len %d\n", uint16_rd(tcp_hdr->src), dest_port, payload_len);

    struct ump_client *client = collections_hash_find(st->tcp.listen_table, dest_port);
    if (client == NULL) {
        if (!(flags & TCP_RST)) {
            // Send RST
            struct tcp_reply_buf *reply = simpleslab_alloc(&st->tcp.tcp_reply_slab);
            if (reply == NULL) {
                ENET_WARN("Too many queued TCP RST replies, not sending reply\n");
            } else {
                reply->eth_dst = eth_addr_rd(eth_hdr->src);
                reply->ip_dst = uint32_rd(ip_hdr->src);
                reply->port_src = uint16_rd(tcp_hdr->dest);
                reply->port_dst = uint16_rd(tcp_hdr->src);
                if (flags & TCP_ACK) {
                    reply->seq = uint32_rd(tcp_hdr->ack);
                    reply->ack = 0;
                    reply->flags = TCP_RST;
                } else {
                    reply->seq = 0;
                    uint16_t seg_len = payload_len;
                    if (flags & (TCP_SYN | TCP_FIN)) seg_len++;
                    reply->ack = uint32_rd(tcp_hdr->seq) + seg_len;
                    reply->flags = TCP_RST | TCP_ACK;
                }
                tx_alloc(&reply->tx_entry, tcp_reply, reply);
            }
        }
        rx_release(buf);
    } else {
        struct tcp_recv *meta = simpleslab_alloc(&st->rx_meta_slab);
        meta->rx_buf = *buf;
        meta->ump_header = (struct ump_net_in_header){
            .op = UMP_NET_EV_TCP_RECV,
            .ev = { .tcp_recv = {
                .src_ip = uint32_rd(ip_hdr->src),
                .dest_ip = uint32_rd(ip_hdr->dest),
            } }
        };
        memcpy(&meta->ump_header.ev.tcp_recv.tcp_header, tcp_hdr, hlen);
        ump_send(client->send_chan, &meta->ump_entry,
            sizeof(struct ump_net_in_header), &meta->ump_header,
            payload_len, payload,
            tcp_recv_ump_callback, meta);
    }
}

// IP: https://datatracker.ietf.org/doc/html/rfc791#section-3.1

static void ip_handle (struct devq_buf *buf, const void *vaddr, const struct eth_hdr *eth_hdr) {
    if (buf->valid_length < IP_HLEN_MIN) {
        ENET_WARN("Received IP packet too small\n");
        rx_release(buf);
        return;
    }
    const struct ip_hdr *ip_hdr = vaddr + buf->valid_data;
    uint16_t ip_len = uint16_rd(ip_hdr->len);
    uint16_t header_len = IPH_HL(ip_hdr) * 4;
    if (
        IPH_V(ip_hdr) != 4 ||
        header_len < 20 ||
        header_len > ip_len ||
        ip_len > buf->valid_length
    ) {
        ENET_WARN("Received bad IP packet\n");
        rx_release(buf);
        return;
    }

    if (inet_checksum(ip_hdr, header_len) != 0) {
        ENET_WARN("Received IP packet with bad checksum\n");
        rx_release(buf);
        return;
    }

    uint16_t offset_flags = uint16_rd(ip_hdr->offset);
    if (
        (offset_flags & IP_MF) != 0 ||
        (offset_flags & IP_OFFMASK) != 0
    ) {
        ENET_WARN("Received IP fragment, dropping\n");
        rx_release(buf);
        return;
    }

    ENET_DEBUG("Received IP packet from %d.%d.%d.%d to %d.%d.%d.%d, proto %d\n",
        ip_hdr->src.val[0],
        ip_hdr->src.val[1],
        ip_hdr->src.val[2],
        ip_hdr->src.val[3],
        ip_hdr->dest.val[0],
        ip_hdr->dest.val[1],
        ip_hdr->dest.val[2],
        ip_hdr->dest.val[3],
        ip_hdr->proto
    );

    buf->valid_data += header_len;
    buf->valid_length = ip_len - header_len;

    if (uint32_rd(ip_hdr->dest) == st->my_ip) {
        if (ip_hdr->proto == IP_PROTO_ICMP) {
            icmp_handle(buf, vaddr, eth_hdr, ip_hdr);
        } else if (ip_hdr->proto == IP_PROTO_UDP) {
            udp_handle(buf, vaddr, eth_hdr, ip_hdr);
        } else if (ip_hdr->proto == IP_PROTO_TCP) {
            tcp_handle(buf, vaddr, eth_hdr, ip_hdr);
        } else {
            ENET_WARN("Received IP packet with unknown protocol %d\n", ip_hdr->proto);
            rx_release(buf);
        }
    } else {
        ENET_WARN("Received IP packet for someone else\n");
        rx_release(buf);
    }
}

static void write_eth_ip_header (
    struct devq_buf *buf, void *vaddr,
    uint64_t eth_dst, uint32_t ip_dest, uint8_t proto, size_t payload_len
) {
    buf->valid_data = 0;
    buf->valid_length = ETH_HLEN + IP_HLEN_MIN + payload_len;
    assert(buf->valid_length <= ETH_HLEN + 1500);

    struct eth_hdr *eth_hdr = vaddr;
    eth_hdr->dst = eth_addr_wr(eth_dst);
    eth_hdr->src = eth_addr_wr(st->mac);
    eth_hdr->type = uint16_wr(ETH_TYPE_IP);

    struct ip_hdr *ip_hdr = vaddr + ETH_HLEN;
    IPH_VHL_SET(ip_hdr, 4, 5);
    ip_hdr->tos = 0;
    ip_hdr->len = uint16_wr(IP_HLEN_MIN + payload_len);
    ip_hdr->id = uint16_wr(st->ip.next_id++);
    ip_hdr->offset = uint16_wr(IP_DF);
    ip_hdr->ttl = 64;
    ip_hdr->proto = proto;
    ip_hdr->chksum = uint16_wr(0);
    ip_hdr->src = uint32_wr(st->my_ip);
    ip_hdr->dest = uint32_wr(ip_dest);
    ip_hdr->chksum = uint16_wr(inet_checksum(ip_hdr, 20));
}

// Ethernet

static void rx_handle (struct devq_buf *buf) {
    struct region_entry *entry = enet_get_region(st->rxq, buf->rid);
    assert(entry != NULL);
    const void *vaddr = (void*)entry->mem.vbase + buf->offset;
    const void *eth_vaddr = vaddr + buf->valid_data;

    #if defined(ENET_DEBUG_OPTION)
    debug_printf("Received Packet of size %lu:", buf->valid_length);
    for (size_t i = 0; i < buf->valid_length; i++) {
        printf(" %02x", ((uint8_t*)eth_vaddr)[i]);
    }
    printf("\n");
    #endif

    if (buf->valid_length < ETH_HLEN) {
        ENET_WARN("Ethernet packet too small\n");
        rx_release(buf);
        return;
    }

    const struct eth_hdr *eth_hdr = eth_vaddr;
    uint16_t type = uint16_rd(eth_hdr->type);

    buf->valid_data += ETH_HLEN;
    buf->valid_length -= ETH_HLEN;

    if (type == ETH_TYPE_ARP) {
        arp_handle(buf, vaddr, eth_hdr);
    } else if (type == ETH_TYPE_IP) {
        ip_handle(buf, vaddr, eth_hdr);
    } else {
        ENET_WARN("Received packet of unknown type: %04x\n", type);
        rx_release(buf);
    }
}

// UMP

static void netump_header_handle (void *arg, size_t header_size, void *header_raw, size_t payload_size);

static void netump_next (struct ump_client *client) {
    if (client->ump_reply_slab.free != 0) {
        ump_recv_header(client->recv_chan, netump_header_handle, client);
    } else {
        client->blocked_on_reply_slab = true;
    }
}

static void netump_reply_callback (void *arg, struct ump_send_queue_entry *entry) {
    struct ump_reply_buf *reply = arg;
    struct ump_client *client = reply->client;
    simpleslab_free(&client->ump_reply_slab, reply);
    if (client->blocked_on_reply_slab) {
        client->blocked_on_reply_slab = false;
        netump_next(client);
    }
}

static void netump_payload_ignore (void *arg, size_t payload_size, void *payload) {
    netump_next(arg);
}

static void netump_send_reply (struct ump_reply_buf *reply, errval_t err) {
    reply->ump_header.ret.err = err;
    ump_send(reply->client->send_chan, &reply->ump_entry,
        sizeof(struct ump_net_in_header), &reply->ump_header,
        0, NULL,
        netump_reply_callback, reply);
}

static void netump_udp_tx_payload (void *arg, size_t payload_size, void *payload) {
    struct ump_client *client = arg;

    struct udp_hdr *udp_hdr = (void*)client->tx_ip_hdr + IP_HLEN_MIN;
    uint16_t checksum = ip_data_checksum(client->tx_ip_hdr, udp_hdr, uint16_rd(udp_hdr->len));
    if (checksum == 0) checksum = 0xffff;
    udp_hdr->chksum = uint16_wr(checksum);

    tx_send(&client->tx_buf);
    netump_next(client);
}

static void netump_udp_tx_allocated (void *arg, struct devq_buf *buf, void *vaddr) {
    struct ump_reply_buf *reply = arg;
    struct ump_client *client = reply->client;
    struct ump_net_out_header *header = client->recv_chan->header;
    size_t payload_size = client->recv_chan->next_payload_size;

    uint64_t eth_dest = (uintptr_t)collections_hash_find(st->arp.table, header->d.udp_send.dest_ip);
    if (eth_dest == 0) {
        ENET_WARN("IP not in ARP table, dropping packet and sending ARP request\n");
        arp_send(buf, vaddr, 0xffffffffffff,
                 header->d.udp_send.dest_ip, ARP_OP_REQ);
        netump_send_reply(reply, LIB_ERR_NET_ARP_MISS);
        ump_recv_payload(client->recv_chan, NULL, netump_payload_ignore, client);
        return;
    }
    write_eth_ip_header(
        buf, vaddr,
        eth_dest,
        header->d.udp_send.dest_ip,
        IP_PROTO_UDP,
        UDP_HLEN + payload_size
    );
    struct udp_hdr *udp_hdr = vaddr + ETH_HLEN + IP_HLEN_MIN;
    udp_hdr->src = uint16_wr(header->d.udp_send.src_port);
    udp_hdr->dest = uint16_wr(header->d.udp_send.dest_port);
    udp_hdr->len = uint16_wr(UDP_HLEN + payload_size);
    udp_hdr->chksum = uint16_wr(0);
    void *payload = (void*)udp_hdr + UDP_HLEN;
    client->tx_buf = *buf;
    client->tx_ip_hdr = vaddr + ETH_HLEN;

    netump_send_reply(reply, SYS_ERR_OK);
    ump_recv_payload(client->recv_chan, payload, netump_udp_tx_payload, client);
}

static void netump_tcp_tx_payload (void *arg, size_t payload_size, void *payload) {
    struct ump_client *client = arg;

    struct tcp_hdr *tcp_hdr = (void*)client->tx_ip_hdr + IP_HLEN_MIN;
    tcp_hdr->chksum = uint16_wr(0);
    uint16_t checksum = ip_data_checksum(client->tx_ip_hdr, tcp_hdr, client->tx_ip_payload_size);
    tcp_hdr->chksum = uint16_wr(checksum);

    tx_send(&client->tx_buf);
    netump_next(client);
}

static void netump_tcp_tx_allocated (void *arg, struct devq_buf *buf, void *vaddr) {
    struct ump_reply_buf *reply = arg;
    struct ump_client *client = reply->client;
    struct ump_net_out_header *header = client->recv_chan->header;
    size_t payload_size = client->recv_chan->next_payload_size;

    uint64_t eth_dest = (uintptr_t)collections_hash_find(st->arp.table, header->d.tcp_send.dest_ip);
    if (eth_dest == 0) {
        ENET_WARN("IP not in ARP table, dropping packet and sending ARP request\n");
        arp_send(buf, vaddr, 0xffffffffffff,
                 header->d.tcp_send.dest_ip, ARP_OP_REQ);
        netump_send_reply(reply, LIB_ERR_NET_ARP_MISS);
        ump_recv_payload(client->recv_chan, NULL, netump_payload_ignore, client);
        return;
    }
    const struct tcp_hdr *ump_tcp_hdr = (const struct tcp_hdr*)&header->d.tcp_send.tcp_header;
    uint16_t tcp_hlen = TCP_HLEN(ump_tcp_hdr);
    write_eth_ip_header(
        buf, vaddr,
        eth_dest,
        header->d.tcp_send.dest_ip,
        IP_PROTO_TCP,
        tcp_hlen + payload_size
    );
    client->tx_buf = *buf;
    client->tx_ip_hdr = vaddr + ETH_HLEN;
    client->tx_ip_payload_size = tcp_hlen + payload_size;
    void *tcp_header = vaddr + ETH_HLEN + IP_HLEN_MIN;
    memcpy(tcp_header, ump_tcp_hdr, tcp_hlen);
    void *payload = tcp_header + tcp_hlen;

    netump_send_reply(reply, SYS_ERR_OK);
    ump_recv_payload(client->recv_chan, payload, netump_tcp_tx_payload, client);
}

static void netump_header_handle (void *arg, size_t header_size, void *header_raw, size_t payload_size) {
    struct ump_client *client = arg;
    struct ump_net_out_header *header = header_raw;
    errval_t err;

    struct ump_reply_buf *reply = simpleslab_alloc(&client->ump_reply_slab);
    assert(reply != NULL);
    reply->client = client;
    reply->ump_header.op = header->op;

    if (header->op == UMP_NET_OP_UDP_SEND) {
        if (payload_size > 1500 - IP_HLEN_MIN - UDP_HLEN) {
            err = LIB_ERR_NET_PACKET_TOO_BIG;
        } else {
            tx_alloc(&client->tx_entry, netump_udp_tx_allocated, reply);
            return;
        }
    } else if (header->op == UMP_NET_OP_UDP_LISTEN) {
        uint16_t port = header->d.udp_listen.dest_port;
        if (port == 0) {
            // allocate an ephemeral port
            uint16_t first = st->udp.next_ephemeral_port;
            port = first;
            while (true) {
                uint16_t next_port;
                if (port == EPHEMERAL_PORT_END) {
                    next_port = EPHEMERAL_PORT_START;
                } else {
                    next_port = port + 1;
                }
                if (collections_hash_find(st->udp.listen_table, port) == NULL) {
                    st->udp.next_ephemeral_port = next_port;
                    break;
                }
                port = next_port;
                if (port == first) {
                    port = 0;
                    break;
                }
            }
        }
        if (port == 0) {
            err = LIB_ERR_NET_ALLOC_PORT;
        } else if (collections_hash_find(st->udp.listen_table, port) != NULL) {
            err = LIB_ERR_NET_PORT_IN_USE;
        } else {
            ENET_DEBUG("Listening on UDP port %d\n", port);
            reply->ump_header.ret.d.udp_listen.dest_port = port;
            collections_hash_insert(st->udp.listen_table, port, client);
            err = SYS_ERR_OK;
        }
    } else if (header->op == UMP_NET_OP_UDP_LISTEN_STOP) {
        uint16_t port = header->d.udp_listen_stop.dest_port;
        if (collections_hash_find(st->udp.listen_table, port) != client) {
            err = LIB_ERR_NET_NOT_LISTENING;
        } else {
            ENET_DEBUG("Stopped listening on UDP port %d\n", port);
            collections_hash_delete(st->udp.listen_table, port);
            err = SYS_ERR_OK;
        }
    } else if (header->op == UMP_NET_OP_TCP_SEND) {
        if (payload_size > 1500 - IP_HLEN_MIN - TCP_HLEN((struct tcp_hdr*)&header->d.tcp_send.tcp_header)) {
            err = LIB_ERR_NET_PACKET_TOO_BIG;
        } else {
            tx_alloc(&client->tx_entry, netump_tcp_tx_allocated, reply);
            return;
        }
    } else if (header->op == UMP_NET_OP_TCP_LISTEN) {
        uint16_t port = header->d.tcp_listen.dest_port;
        if (port == 0) {
            // allocate an ephemeral port
            uint16_t first = st->tcp.next_ephemeral_port;
            port = first;
            while (true) {
                uint16_t next_port;
                if (port == EPHEMERAL_PORT_END) {
                    next_port = EPHEMERAL_PORT_START;
                } else {
                    next_port = port + 1;
                }
                if (collections_hash_find(st->tcp.listen_table, port) == NULL) {
                    st->tcp.next_ephemeral_port = next_port;
                    break;
                }
                port = next_port;
                if (port == first) {
                    port = 0;
                    break;
                }
            }
        }
        if (port == 0) {
            err = LIB_ERR_NET_ALLOC_PORT;
        } else if (collections_hash_find(st->tcp.listen_table, port) != NULL) {
            err = LIB_ERR_NET_PORT_IN_USE;
        } else {
            ENET_DEBUG("Listening on TCP port %d\n", port);
            reply->ump_header.ret.d.tcp_listen.dest_port = port;
            collections_hash_insert(st->tcp.listen_table, port, client);
            err = SYS_ERR_OK;
        }
    } else if (header->op == UMP_NET_OP_TCP_LISTEN_STOP) {
        uint16_t port = header->d.tcp_listen_stop.dest_port;
        if (collections_hash_find(st->tcp.listen_table, port) != client) {
            err = LIB_ERR_NET_NOT_LISTENING;
        } else {
            ENET_DEBUG("Stopped listening on TCP port %d\n", port);
            collections_hash_delete(st->tcp.listen_table, port);
            err = SYS_ERR_OK;
        }
    } else {
        debug_printf("Error: unknown UMP op\n");
        simpleslab_free(&client->ump_reply_slab, reply);
        ump_recv_payload(client->recv_chan, NULL, netump_payload_ignore, client);
        return;
    }

    netump_send_reply(reply, err);
    ump_recv_payload(client->recv_chan, NULL, netump_payload_ignore, client);
}

static errval_t netump_connect (void *arg, struct capref cap) {
    errval_t err;

    ENET_DEBUG("Incoming UMP connection\n");

    struct ump_client *client = malloc(sizeof(struct ump_client));
    if (client == NULL) return LIB_ERR_MALLOC_FAIL;

    err = ump_chan_init(UMP_ROLE_SERVER, &client->send_chan, &client->recv_chan,
        sizeof(struct ump_net_out_header), cap, get_default_waitset());
    if (err_is_fail(err)) return err;

    simpleslab_init(&client->ump_reply_slab, sizeof(struct ump_reply_buf),
        &client->ump_reply_slab_buf, sizeof(client->ump_reply_slab_buf));
    client->blocked_on_reply_slab = false;

    netump_next(client);

    return SYS_ERR_OK;
}

void enet_loop (void) {
    errval_t err;

    st->arp.tx_queue_len = 0;
    st->arp.tx_queue_tail = 0;
    collections_hash_create(&st->arp.table, NULL);

    st->ip.next_id = 0;

    collections_hash_create_with_buckets(&st->udp.listen_table, 100, NULL);
    st->udp.next_ephemeral_port = EPHEMERAL_PORT_START;

    collections_hash_create_with_buckets(&st->tcp.listen_table, 100, NULL);
    st->tcp.next_ephemeral_port = EPHEMERAL_PORT_START;

    simpleslab_init(&st->tcp.tcp_reply_slab, sizeof(struct tcp_reply_buf),
        &st->tcp.tcp_reply_slab_buf, sizeof(st->tcp.tcp_reply_slab_buf));

    simpleslab_init(&st->rx_meta_slab, sizeof(union rx_meta_slab_data),
        &st->rx_meta_slab_buf, sizeof(st->rx_meta_slab_buf));

    struct ump_binding_server server;
    err = ump_binding_register(&server, UMP_SERVER_NET, netump_connect, NULL);
    if (err_is_fail(err)) USER_PANIC_ERR(err, "Failed to register UMP server");

    struct devq_buf buf;
    struct waitset *default_ws = get_default_waitset();
    while (true) {
        bool made_progress = false;

        err = event_dispatch_non_block(default_ws);
        if (err_is_fail(err) && err != LIB_ERR_NO_EVENT) {
            DEBUG_ERR(err, "in event_dispatch");
            abort();
        }
        // TODO: set made_progress if we handled an event. Currently,
        // UMP always registers a callback, so we would always see an event.
        // See TODO in poll_channels_disabled.

        err = devq_dequeue((struct devq*) st->rxq, &buf.rid, &buf.offset,
                           &buf.length, &buf.valid_data, &buf.valid_length,
                           &buf.flags);
        if (err_is_ok(err)) {
            made_progress = true;
            rx_handle(&buf);
        }

        if (st->tx_alloc_queue_head != NULL) {
            err = devq_dequeue((struct devq*) st->txq, &buf.rid, &buf.offset,
                               &buf.length, &buf.valid_data, &buf.valid_length,
                               &buf.flags);
            if (err_is_ok(err)) {
                made_progress = true;
                struct tx_alloc_queue_entry *entry = st->tx_alloc_queue_head;
                st->tx_alloc_queue_head = entry->next;
                if (entry->next == NULL) {
                    st->tx_alloc_queue_tail = NULL;
                }
                struct region_entry *region_entry = enet_get_region(st->txq, buf.rid);
                assert(region_entry != NULL);
                void *vaddr = (void*)region_entry->mem.vbase + buf.offset;
                entry->cb(entry->cb_arg, &buf, vaddr);
            }
        }

        if (!made_progress) {
            thread_yield();
        }
    }
}