#include <string.h>
#include <stdbool.h>
#include <sys/queue.h>

#include "log_internal.h"
#include "interval_tree.h"
#include "tcp_reassembly.h"
#include "stellar/stellar.h"

#define TCP_REASSEMBLY_LOG_DEBUG(format, ...) STELLAR_LOG_DEBUG(__thread_local_logger, "tcp_reassembly", format, ##__VA_ARGS__)
#define TCP_REASSEMBLY_LOG_ERROR(format, ...) STELLAR_LOG_ERROR(__thread_local_logger, "tcp_reassembly", format, ##__VA_ARGS__)

struct tcp_segment_private
{
    uint64_t ts;
    uint64_t id;
    struct interval_tree_node node;
    TAILQ_ENTRY(tcp_segment_private) lru;
    struct tcp_segment seg;
    void *data; // flexible array member
};

TAILQ_HEAD(tcp_segment_private_list, tcp_segment_private);

struct tcp_reassembly
{
    uint64_t max_timeout;
    uint64_t max_seg_num;
    uint64_t cur_seg_num;
    uint64_t sum_seg_num;

    struct rb_root_cached rb_root;
    struct tcp_segment_private_list lru_list;
    uint32_t recv_next;
};

struct tcp_segment *tcp_segment_new(uint32_t seq, const void *data, uint32_t len)
{
    struct tcp_segment_private *p = (struct tcp_segment_private *)malloc(sizeof(struct tcp_segment_private) + len);
    if (!p)
    {
        TCP_REASSEMBLY_LOG_ERROR("calloc failed");
        return NULL;
    }

    p->ts = 0;
    p->id = 0;

    p->node.start = seq;
    p->node.last = (uint64_t)seq + (uint64_t)len - 1;
    p->data = (char *)p + sizeof(struct tcp_segment_private);
    memcpy(p->data, data, len);

    p->seg.len = len;
    p->seg.data = p->data;

    return &p->seg;
}

void tcp_segment_free(struct tcp_segment *seg)
{
    if (seg)
    {
        struct tcp_segment_private *p = container_of(seg, struct tcp_segment_private, seg);
        free(p);
    }
}

struct tcp_reassembly *tcp_reassembly_new(uint64_t max_timeout, uint64_t max_seg_num)
{
    struct tcp_reassembly *tcp_reass = (struct tcp_reassembly *)malloc(sizeof(struct tcp_reassembly));
    if (!tcp_reass)
    {
        TCP_REASSEMBLY_LOG_ERROR("calloc failed");
        return NULL;
    }

    tcp_reass->max_timeout = max_timeout;
    tcp_reass->max_seg_num = max_seg_num;
    tcp_reass->cur_seg_num = 0;
    tcp_reass->sum_seg_num = 0;
    tcp_reass->rb_root = RB_ROOT_CACHED;
    tcp_reass->recv_next = 0;
    TAILQ_INIT(&tcp_reass->lru_list);

    return tcp_reass;
}

void tcp_reassembly_free(struct tcp_reassembly *tcp_reass)
{
    if (tcp_reass)
    {
        struct tcp_segment_private *p = NULL;
        while ((p = TAILQ_FIRST(&tcp_reass->lru_list)))
        {
            tcp_reass->cur_seg_num--;
            TAILQ_REMOVE(&tcp_reass->lru_list, p, lru);
            interval_tree_remove(&p->node, &tcp_reass->rb_root);
            free(p);
        }

        free(tcp_reass);
    }
}

/*
 * return:  1: push tcp segment success (segment overlap)
 * return:  0: push tcp segment success
 * return: -1: push tcp segment failed (no space)
 * return: -2: push tcp segment failed (segment repeat)
 */
int tcp_reassembly_push(struct tcp_reassembly *tcp_reass, struct tcp_segment *seg, uint64_t now)
{
    if (tcp_reass == NULL)
    {
        return -1;
    }

    if (tcp_reass->cur_seg_num >= tcp_reass->max_seg_num)
    {
        TCP_REASSEMBLY_LOG_ERROR("tcp_reass %p is full", tcp_reass);
        return -1;
    }

    int ret = 0;
    struct tcp_segment_private *p = container_of(seg, struct tcp_segment_private, seg);
    struct interval_tree_node *node = interval_tree_iter_first(&tcp_reass->rb_root, p->node.start, p->node.last);
    if (node)
    {
        do
        {
            struct tcp_segment_private *t = container_of(node, struct tcp_segment_private, node);
            if (t->node.start == p->node.start && t->node.last == p->node.last)
            {
                TCP_REASSEMBLY_LOG_DEBUG("tcp_reass %p push segment %p [%lu, %lu] failed, segment repeat", tcp_reass, seg, p->node.start, p->node.last);
                return -2;
            }
        } while ((node = interval_tree_iter_next(node, p->node.start, p->node.last)));

        TCP_REASSEMBLY_LOG_DEBUG("tcp_reass %p push segment %p [%lu, %lu], but segment overlap", tcp_reass, seg, p->node.start, p->node.last);
        ret = 1;
    }
    else
    {
        TCP_REASSEMBLY_LOG_DEBUG("tcp_reass %p push segment %p [%lu, %lu]", tcp_reass, seg, p->node.start, p->node.last);
    }

    p->ts = now;
    p->id = tcp_reass->sum_seg_num++;
    TAILQ_INSERT_TAIL(&tcp_reass->lru_list, p, lru);
    interval_tree_insert(&p->node, &tcp_reass->rb_root);

    tcp_reass->cur_seg_num++;

    return ret;
}

struct tcp_segment *tcp_reassembly_pop(struct tcp_reassembly *tcp_reass)
{
    if (tcp_reass == NULL)
    {
        return NULL;
    }

    struct interval_tree_node *node = interval_tree_iter_first(&tcp_reass->rb_root, tcp_reass->recv_next, tcp_reass->recv_next);
    if (node == NULL)
    {
        return NULL;
    }

    uint64_t overlap = 0;
    uint64_t min_id = UINT64_MAX;
    struct tcp_segment_private *oldest = NULL;
    while (node)
    {
        struct tcp_segment_private *p = container_of(node, struct tcp_segment_private, node);
        if (p->id < min_id)
        {
            min_id = p->id;
            oldest = p;
        }
        node = interval_tree_iter_next(node, tcp_reass->recv_next, tcp_reass->recv_next);
    }

    TAILQ_REMOVE(&tcp_reass->lru_list, oldest, lru);
    interval_tree_remove(&oldest->node, &tcp_reass->rb_root);

    tcp_reass->cur_seg_num--;

    if (oldest->node.start < tcp_reass->recv_next)
    {
        // trim overlap
        overlap = tcp_reass->recv_next - oldest->node.start;
        oldest->seg.len -= overlap;
        oldest->seg.data = (char *)oldest->data + overlap;
    }

    TCP_REASSEMBLY_LOG_DEBUG("tcp_reass %p pop segment %p [%lu, %lu], trim overlap %lu", tcp_reass, &oldest->seg, oldest->node.start, oldest->node.last, overlap);

    // update recv_next
    tcp_reass->recv_next = uint32_add(tcp_reass->recv_next, oldest->seg.len);

    return &oldest->seg;
}

struct tcp_segment *tcp_reassembly_expire(struct tcp_reassembly *tcp_reass, uint64_t now)
{
    if (tcp_reass == NULL)
    {
        return NULL;
    }

    struct tcp_segment_private *p = TAILQ_FIRST(&tcp_reass->lru_list);
    if (p && now - p->ts >= tcp_reass->max_timeout)
    {
        tcp_reass->cur_seg_num--;
        TAILQ_REMOVE(&tcp_reass->lru_list, p, lru);
        interval_tree_remove(&p->node, &tcp_reass->rb_root);
        TCP_REASSEMBLY_LOG_DEBUG("tcp_reass %p expire segment %p [%lu, %lu]", tcp_reass, &p->seg, p->node.start, p->node.last);
        return &p->seg;
    }
    else
    {
        return NULL;
    }
}

void tcp_reassembly_inc_recv_next(struct tcp_reassembly *tcp_reass, uint32_t offset)
{
    if (tcp_reass == NULL)
    {
        return;
    }

    tcp_reass->recv_next = uint32_add(tcp_reass->recv_next, offset);
    TCP_REASSEMBLY_LOG_DEBUG("tcp_reass %p inc recv_next %u to %u", tcp_reass, offset, tcp_reass->recv_next);
}

void tcp_reassembly_set_recv_next(struct tcp_reassembly *tcp_reass, uint32_t seq)
{
    if (tcp_reass == NULL)
    {
        return;
    }

    tcp_reass->recv_next = seq;
    TCP_REASSEMBLY_LOG_DEBUG("tcp_reass %p set recv_next %u", tcp_reass, seq);
}

uint32_t tcp_reassembly_get_recv_next(struct tcp_reassembly *tcp_reass)
{
    if (tcp_reass == NULL)
    {
        return 0;
    }

    return tcp_reass->recv_next;
}