#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <assert.h>

#include "tcp_reassembly.h"
#include "itree.h"

struct tcp_reassembly
{
    // config
    bool enable;
    uint32_t max_timeout;
    uint32_t max_packets;
    uint32_t max_bytes;

    // stat
    struct tcp_reassembly_stat stat;

    // runtime
    struct itree *c2s_itree;
    struct itree *s2c_itree;
    uint64_t c2s_exp_seq;
    uint64_t s2c_exp_seq;

    // used for timeout
    struct segment *head; // del segment from head
    struct segment *tail; // add segment to tail
};

struct segment
{
    struct tcp_reassembly *assembler;
    struct itree *itree;
    struct segment *next;
    struct segment *prev;

    uint64_t time;
    uint32_t offset;
    uint32_t len;
    char *payload; // Flexible array member
};

void *segment_dup(void *p)
{
    return p;
}

void segment_rel(void *p)
{
    struct segment *seg = (struct segment *)p;
    if (seg)
    {
        struct tcp_reassembly *assembler = seg->assembler;
        // delete from list
        if (assembler->head == seg)
        {
            assembler->head = seg->next;
        }
        if (assembler->tail == seg)
        {
            assembler->tail = seg->prev;
        }
        if (seg->prev)
        {
            seg->prev->next = seg->next;
        }
        if (seg->next)
        {
            seg->next->prev = seg->prev;
        }

        assembler->stat.bytes -= seg->len;
        assembler->stat.packets--;

        free(seg);
    }
}

struct tcp_reassembly *tcp_reassembly_new(bool enable, uint32_t max_timeout, uint32_t max_packets, uint32_t max_bytes)
{
    struct tcp_reassembly *assembler = (struct tcp_reassembly *)calloc(1, sizeof(struct tcp_reassembly));
    if (assembler == NULL)
    {
        return NULL;
    }

    assembler->enable = enable;
    assembler->max_timeout = max_timeout;
    assembler->max_packets = max_packets;
    assembler->max_bytes = max_bytes;
    if (!assembler->enable)
    {
        return assembler;
    }

    assembler->c2s_itree = itree_new(segment_dup, segment_rel);
    assembler->s2c_itree = itree_new(segment_dup, segment_rel);
    if (assembler->c2s_itree == NULL || assembler->s2c_itree == NULL)
    {
        goto error_out;
    }

    return assembler;

error_out:
    tcp_reassembly_free(assembler);
    return NULL;
}

void tcp_reassembly_init(struct tcp_reassembly *assembler, uint32_t c2s_init_seq, uint32_t s2c_init_seq)
{
    if (!assembler->enable)
    {
        return;
    }

    if (c2s_init_seq)
    {
        assembler->c2s_exp_seq = c2s_init_seq + 1;
    }

    if (s2c_init_seq)
    {
        assembler->s2c_exp_seq = s2c_init_seq + 1;
    }
}

void tcp_reassembly_free(struct tcp_reassembly *assembler)
{
    if (assembler)
    {
        if (assembler->c2s_itree)
        {
            itree_delete(assembler->c2s_itree);
        }
        if (assembler->s2c_itree)
        {
            itree_delete(assembler->s2c_itree);
        }
        free(assembler);
    }
}

void tcp_reassembly_expire(struct tcp_reassembly *assembler, uint64_t now)
{
    if (!assembler->enable)
    {
        return;
    }

    struct tcp_reassembly_stat *stat = &assembler->stat;
    while (assembler->head)
    {
        struct segment *seg = assembler->head;
        if (now - seg->time < assembler->max_timeout)
        {
            break;
        }

        stat->tcp_segement_timout++;

        struct itree *itree = seg->itree;
        interval_t interval = {
            .low = seg->offset,
            .high = seg->offset + seg->len - 1,
        };
        itree_remove(itree, &interval);
    }
}

void tcp_reassembly_update(struct tcp_reassembly *assembler, int direction, uint32_t offset, const char *payload, uint32_t len, uint64_t now)
{
    if (!assembler->enable)
    {
        return;
    }

    struct itree *itree = (direction == 0x01) ? assembler->c2s_itree : assembler->s2c_itree;
    uint64_t exp_seq = (direction == 0x01) ? assembler->c2s_exp_seq : assembler->s2c_exp_seq;

    if (len == 0 || offset + len < exp_seq)
    {
        return;
    }

    if (assembler->max_packets > 0 && assembler->stat.packets >= assembler->max_packets)
    {
        return;
    }

    if (assembler->max_bytes > 0 && assembler->stat.bytes >= assembler->max_bytes)
    {
        return;
    }

    struct segment *seg = (struct segment *)calloc(1, sizeof(struct segment) + len);
    if (seg == NULL)
    {
        return;
    }

    seg->itree = itree;
    seg->assembler = assembler;
    seg->time = now;
    seg->offset = offset;
    seg->len = len;
    seg->payload = (char *)seg + sizeof(struct segment);
    memcpy(seg->payload, payload, len);

    interval_t interval = {
        .low = seg->offset,
        .high = seg->offset + seg->len - 1,
        .data = seg,
    };
    if (itree_insert(itree, &interval) == 0)
    {
        free(seg);
        return;
    }
    TCP_REASSEMBLE_DEBUG("%s insert [%lu, %lu], segment {ptr: %p, offset: %lu, len: %lu}",
                         (direction == 0x01) ? "C2S" : "S2C",
                         seg->offset, seg->offset + seg->len - 1,
                         seg, seg->offset, seg->len);

    if (assembler->head == NULL)
    {
        assembler->head = seg;
    }
    else
    {
        assembler->tail->next = seg;
        seg->prev = assembler->tail;
    }
    assembler->tail = seg;

    assembler->stat.packets++;
    assembler->stat.bytes += len;

    tcp_reassembly_expire(assembler, now);
}

const char *tcp_reassembly_peek(struct tcp_reassembly *assembler, int direction, uint32_t *len)
{
    *len = 0;

    if (!assembler->enable)
    {
        return NULL;
    }

    struct itree *itree = (direction == 0x01) ? assembler->c2s_itree : assembler->s2c_itree;
    uint64_t exp_seq = (direction == 0x01) ? assembler->c2s_exp_seq : assembler->s2c_exp_seq;

    interval_t interval = {
        .low = exp_seq,
        .high = exp_seq,
    };
    interval_t *result = itree_find(itree, &interval);
    if (result == NULL)
    {
        TCP_REASSEMBLE_DEBUG("%s peek [%lu, +∞]: not found", (direction == 0x01) ? "C2S" : "S2C", exp_seq);
        return NULL;
    }

    struct segment *seg = (struct segment *)result->data;
    assert(seg != NULL);
    // check overlap
    if (seg->offset < exp_seq)
    {
        TCP_REASSEMBLE_DEBUG("%s peek [%lu, +∞], found [%lu, %lu], segment {ptr: %p, offset: %lu, len: %lu, left trim: %lu}",
                             (direction == 0x01) ? "C2S" : "S2C", exp_seq,
                             seg->offset, seg->offset + seg->len - 1,
                             seg, seg->offset, seg->len, exp_seq - seg->offset);
        *len = seg->len - (exp_seq - seg->offset);
        return seg->payload + (exp_seq - seg->offset);
    }
    TCP_REASSEMBLE_DEBUG("%s peek [%lu, +∞], found [%lu, %lu], segment {ptr: %p, offset: %lu, len: %lu}",
                         (direction == 0x01) ? "C2S" : "S2C", exp_seq,
                         seg->offset, seg->offset + seg->len - 1,
                         seg, seg->offset, seg->len);

    *len = seg->len;
    return seg->payload;
}

void tcp_reassembly_consume(struct tcp_reassembly *assembler, int direction, uint32_t len)
{
    if (!assembler->enable)
    {
        return;
    }

    if (len == 0)
    {
        return;
    }

    struct itree *itree = (direction == 0x01) ? assembler->c2s_itree : assembler->s2c_itree;
    uint64_t *exp_seq = (direction == 0x01) ? &assembler->c2s_exp_seq : &assembler->s2c_exp_seq;
    uint64_t old_exp_seq = *exp_seq;

    *exp_seq += len;
    uint64_t new_exp_seq = *exp_seq;

    TCP_REASSEMBLE_DEBUG("%s consume [%lu, %lu], update expect seq %lu -> %lu",
                         (direction == 0x01) ? "C2S" : "S2C",
                         old_exp_seq, new_exp_seq - 1, old_exp_seq, new_exp_seq);

    interval_t interval = {
        .low = 0,
        .high = *exp_seq,
    };

    ilist_t *list = itree_findall(itree, &interval);
    if (list == NULL)
    {
        return;
    }

    interval_t *result;
    int count = ilist_size(list);
    ilisttrav_t *trav = ilisttrav_new(list);
    for (int i = 0; i < count; i++)
    {
        if (i == 0)
        {
            result = (interval_t *)ilisttrav_first(trav);
        }
        else
        {
            result = (interval_t *)ilisttrav_next(trav);
        }
        if (result && result->high < *exp_seq)
        {
            struct segment *seg = (struct segment *)result->data;
            TCP_REASSEMBLE_DEBUG("%s consume [%lu, %lu], delete [%lu, %lu], segment {ptr: %p, offset: %lu, len: %lu}",
                                 (direction == 0x01) ? "C2S" : "S2C", old_exp_seq, new_exp_seq - 1,
                                 result->low, result->high, seg, seg->offset, seg->len);
            itree_remove(itree, result);
        }
    }
    ilisttrav_delete(trav);
    ilist_delete(list);
}

struct tcp_reassembly_stat *tcp_reassembly_get_stat(struct tcp_reassembly *assembler)
{
    return NULL;
}