/*
 * wip, a small TCP/IP stack.
 * Copyright (C) 2025  Xavier Del Campo Romero
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#include <wip/wip.h>
#include <wip/prv/io.h>
#include <wip/prv/log.h>
#include <wip/prv/icmp.h>
#include <wip/prv/tcp.h>
#include <wip/prv/udp.h>
#include <wip/prv/routines.h>
#include <stddef.h>

#include <time.h>

unsigned valid_packets;

enum {MIN_HEADER_SZ = 20};

static const struct wip_prot_ops protocols[] =
{
    {
        WIP_P_ICMP,
        wip_icmp_alloc,
        wip_icmp_free,
        wip_icmp_rx,
        wip_icmp_tx
    },

    {
        WIP_P_TCP,
        wip_tcp_alloc,
        wip_tcp_free,
        wip_tcp_rx,
        wip_tcp_tx
    },

    {
        WIP_P_UDP,
        wip_udp_alloc,
        wip_udp_free,
        wip_udp_rx,
        wip_udp_tx
    }
};

static enum wip_state get_data(struct wip *const w)
{
    struct wip_rx *const rx = &w->rx;
    struct wip_rx_sm_h *const h = &rx->sm.h;
    struct wip_sm_io *const io = &h->io;
    const unsigned short dlen = wip_be16(&h->len) - h->header_sz;
    const struct wip_prot_ops *const p = h->prot;
    enum wip_state n;

    if ((n = wip_io_read(w, io)))
        return n;
    else if (p && (n = p->rx(w, rx->buf, io->n, h->args)))
        return n;

#if 0
    wip_log("got new data, h->i=%lu, io->n=%lu, dlen=%u\n",
        (unsigned long)h->i, (unsigned long)io->n, dlen);
#endif

    if ((h->i += io->n) >= dlen)
    {
#if 1
        wip_log("finished getting data (%u bytes)\n", (unsigned)dlen);
#endif
        if (h->prot)
        {
            h->prot->free(w, h->args);
            h->prot = h->args = NULL;
        }

        wip_rx(w);
        valid_packets++;
    }
    else
    {
        struct wip_sm_io nio = {0};
        const size_t n = dlen - h->i;

#if 0
        wip_log("still need more data! (%u bytes)\n", n);
#endif

        nio.buf = rx->buf;
        nio.n = n > sizeof rx->buf ? sizeof rx->buf : n;
        *io = nio;
    }

    return WIP_OK;
}

static enum wip_state get_options(struct wip *const w)
{
    unsigned char b;
    struct wip_rx *const rx = &w->rx;
    struct wip_rx_sm_h *const h = &rx->sm.h;
    struct wip_sm_io io = {0};
    enum wip_state n;

    io.buf = &b;
    io.n = sizeof b;

    if ((n = wip_io_read(w, &io)))
        return n;
    else if (++h->i >= h->header_sz - MIN_HEADER_SZ)
    {
        struct wip_sm_io nio = {0};
        const unsigned short dlen = wip_be16(&h->len);
        const size_t n = dlen - h->i;

        nio.buf = rx->buf;
        nio.n = n > sizeof rx->buf ? sizeof rx->buf : n;
        h->io = nio;
        h->i = 0;
#if 0
        wip_log("%s: will read %lu bytes of data\n", __func__, (unsigned long)h->io.n);
#endif
        rx->next = get_data;
    }

    h->chksum += b;
    return WIP_OK;
}

static enum wip_state get_dst(struct wip *const w)
{
    struct wip_rx *const rx = &w->rx;
    struct wip_rx_sm_h *const h = &rx->sm.h;
    const enum wip_state n = wip_io_read(w, &h->io);

    if (n)
        return n;

#if 0
    wip_log("got dst address=%hhu.%hhu.%hhu.%hhu\n",
        h->src.v[0], h->src.v[1], h->src.v[2], h->src.v[3]);
#endif
    h->chksum += wip_be32(&h->dst);

    if (h->header_sz > MIN_HEADER_SZ)
    {
#if 0
        wip_log("will parse options\n");
#endif
        rx->next = get_options;
    }
    else
    {
        struct wip_sm_io io = {0};
        const size_t n = wip_be16(&h->len) - h->header_sz;

        io.buf = rx->buf;
        io.n = n >= sizeof rx->buf ? sizeof rx->buf : n;
#if 0
        wip_log("%s: expecting %lu bytes of data\n", __func__, (unsigned long)io.n);
#endif
        h->i = 0;
        h->io = io;
        rx->next = get_data;
    }

    return WIP_OK;
}

static enum wip_state get_src(struct wip *const w)
{
    struct wip_rx *const rx = &w->rx;
    struct wip_rx_sm_h *const h = &rx->sm.h;
    const enum wip_state n = wip_io_read(w, &h->io);

    if (n)
        return n;
    else
    {
        struct wip_sm_io io = {0};

        io.buf = &h->dst;
        io.n = sizeof h->dst;
        h->io = io;
    }

#if 0
    wip_log("got source address=%hhu.%hhu.%hhu.%hhu\n",
        h->src.v[0], h->src.v[1], h->src.v[2], h->src.v[3]);
#endif
    h->chksum += wip_be32(&h->src);
    rx->next = get_dst;
    return WIP_OK;
}

static enum wip_state get_chksum(struct wip *const w)
{
    struct wip_rx *const rx = &w->rx;
    struct wip_rx_sm_h *const h = &rx->sm.h;
    const enum wip_state n = wip_io_read(w, &h->io);

    if (n)
        return n;
    else
    {
        struct wip_sm_io io = {0};

        io.buf = &h->src;
        io.n = sizeof h->src;
        h->io = io;
    }

#if 0
    wip_log("got expected checksum=%hu\n", wip_be16(&h->rchksum));
#endif
    rx->next = get_src;
    return WIP_OK;
}

static const struct wip_prot_ops *get_ops(const unsigned char protocol)
{
    size_t i;

    for (i = 0; i < sizeof protocols / sizeof *protocols; i++)
    {
        const struct wip_prot_ops *const p = &protocols[i];

        if (protocol == p->id)
            return p;
    }

    return NULL;
}

static enum wip_state get_protocol(struct wip *const w)
{
    struct wip_rx *const rx = &w->rx;
    struct wip_rx_sm_h *const h = &rx->sm.h;
    struct wip_sm_io io = {0};
    enum wip_state n;

    io.buf = &h->protocol;
    io.n = sizeof h->protocol;

    if ((n = wip_io_read(w, &io)))
        return n;
    else
    {
        struct wip_sm_io io = {0};

        io.buf = &h->rchksum;
        io.n = sizeof h->rchksum;
        h->io = io;
    }

#if 0
    wip_log("got protocol: %#hhx\n", h->protocol);
#endif
    if ((h->prot = get_ops(h->protocol))
        && !(h->args = h->prot->alloc(w)))
        return WIP_FATAL;

    h->chksum += h->protocol;
    rx->next = get_chksum;
    return WIP_OK;
}

static enum wip_state get_ttl(struct wip *const w)
{
    struct wip_rx *const rx = &w->rx;
    struct wip_rx_sm_h *const h = &rx->sm.h;
    struct wip_sm_io io = {0};
    enum wip_state n;

    io.buf = &h->ttl;
    io.n = sizeof h->ttl;

    if ((n = wip_io_read(w, &io)))
        return n;

#if 0
    wip_log("got TTL\n");
#endif
    h->chksum += h->ttl;
    rx->next = get_protocol;
    return WIP_OK;
}

static enum wip_state get_flags_offset(struct wip *const w)
{
    unsigned char flags;
    unsigned short fl_off;
    struct wip_rx *const rx = &w->rx;
    struct wip_rx_sm_h *const h = &rx->sm.h;
    const enum wip_state n = wip_io_read(w, &h->io);

    if (n)
        return n;

    fl_off = wip_be16(&h->fl_off);

    if ((flags = fl_off >> 13) & 1)
        return WIP_INVALID;

#if 0
    wip_log("got flags offset: %hu\n", fl_off);
#endif
    h->chksum += fl_off;
    rx->next = get_ttl;
    return WIP_OK;
}

static enum wip_state get_id(struct wip *const w)
{
    struct wip_rx *const rx = &w->rx;
    struct wip_rx_sm_h *const h = &rx->sm.h;
    const enum wip_state n = wip_io_read(w, &h->io);

    /* TODO: use ID? */
    if (n)
        return n;
    else
    {
        struct wip_sm_io io = {0};

        io.buf = &h->fl_off;
        io.n = sizeof h->fl_off;
        h->io = io;
    }

#if 0
    wip_log("got id: %hu\n", wip_be16(&h->id));
#endif
    h->chksum += wip_be16(&h->id);
    rx->next = get_flags_offset;
    return WIP_OK;
}

static enum wip_state get_len(struct wip *const w)
{
    unsigned short len;
    struct wip_rx *const rx = &w->rx;
    struct wip_rx_sm_h *const h = &rx->sm.h;
    const enum wip_state n = wip_io_read(w, &h->io);

    if (n)
        return n;
    else if ((len = wip_be16(&h->len)) < MIN_HEADER_SZ)
        return WIP_INVALID;
    else if (len > sizeof rx->buf)
    {
#if 0
        wip_log("data too big: %hu (max %lu)\n", len,
            (unsigned long)sizeof rx->buf);
#endif
        return WIP_INVALID;
    }
    else
    {
        struct wip_sm_io io = {0};

        io.buf = &h->id;
        io.n = sizeof h->id;
        h->io = io;
    }

#if 0
    wip_log("got length: %hu\n", len);
#endif
    h->chksum = len;
    rx->next = get_id;
    return WIP_OK;
}

static enum wip_state get_dscp_ecn(struct wip *const w)
{
    unsigned char b;
    struct wip_rx *const rx = &w->rx;
    struct wip_rx_sm_h *const h = &rx->sm.h;
    struct wip_sm_io io = {0};
    enum wip_state n;

    io.buf = &b;
    io.n = sizeof b;

    /* TODO: support DSCP and ECN? */
    if ((n = wip_io_read(w, &io)))
        return n;
    else
    {
        struct wip_rx_sm_h *const h = &rx->sm.h;
        struct wip_sm_io io = {0};

        io.buf = &h->len;
        io.n = sizeof h->len;
        h->io = io;
    }

#if 0
    wip_log("got DSCP and ECN\n");
#endif
    h->chksum += b;
    rx->next = get_len;
    return WIP_OK;
}

static enum wip_state get_version_ihl(struct wip *const w)
{
    enum {VERSION = 4};
    unsigned char b, version, ihl;
    struct wip_rx *const rx = &w->rx;
    struct wip_rx_sm_h *const h = &rx->sm.h;
    struct wip_sm_io io = {0};
    enum wip_state n;

    io.buf = &b;
    io.n = sizeof b;

    if ((n = wip_io_read(w, &io)))
        return n;
    else if ((version = b >> 4) != 4
        || (ihl = b & 0xf) < MIN_HEADER_SZ / 4)
        return WIP_INVALID;

#if 0
    wip_log("got IPv4 header, length=%u\n", ihl * 4);
#endif
    h->chksum += b;
    h->header_sz = ihl * 4;
    rx->next = get_dscp_ecn;
    return WIP_OK;
}

void wip_rx(struct wip *const w)
{
    const struct wip_rx_sm_h h = {0};
    struct wip_rx *const rx = &w->rx;
    struct wip_rx_sm_h *const ph = &rx->sm.h;

    *ph = h;
    rx->next = get_version_ihl;
}