#include "network/chainbuffer.h"

#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/uio.h>

#define CHAINBUFFER_DEFAULT_CHUNK 4096
#define CHAINBUFFER_MAX_IOV 16

struct chain_buffer_node {
    struct chain_buffer_node *next;
    size_t start;
    size_t end;
    size_t cap;
    uint8_t data[];
};

static chain_buffer_node_t *alloc_node(size_t cap) {
    chain_buffer_node_t *node = (chain_buffer_node_t *)malloc(sizeof(*node) + cap);
    if (!node) {
        return NULL;
    }
    node->next = NULL;
    node->start = 0;
    node->end = 0;
    node->cap = cap;
    return node;
}

void chain_buffer_init(chain_buffer_t *buf, size_t chunk_size) {
    if (!buf) {
        return;
    }
    memset(buf, 0, sizeof(*buf));
    buf->chunk_size = chunk_size ? chunk_size : CHAINBUFFER_DEFAULT_CHUNK;
}

void chain_buffer_reset(chain_buffer_t *buf) {
    if (!buf) {
        return;
    }

    chain_buffer_node_t *node = buf->head;
    while (node) {
        chain_buffer_node_t *next = node->next;
        free(node);
        node = next;
    }

    free(buf->linear_cache);
    buf->linear_cache = NULL;
    buf->linear_cap = 0;
    buf->head = NULL;
    buf->tail = NULL;
    buf->total_len = 0;
}

size_t chain_buffer_len(const chain_buffer_t *buf) {
    return buf ? buf->total_len : 0;
}

int chain_buffer_append(chain_buffer_t *buf, const void *data, size_t len) {
    const uint8_t *src = (const uint8_t *)data;
    if (!buf || (!src && len > 0)) {
        errno = EINVAL;
        return -1;
    }
    if (len == 0) {
        return 0;
    }
    if (buf->total_len > (size_t)-1 - len) {
        errno = EOVERFLOW;
        return -1;
    }

    size_t remain = len;
    while (remain > 0) {
        chain_buffer_node_t *tail = buf->tail;
        size_t writable = 0;

        if (tail && tail->end < tail->cap) {
            writable = tail->cap - tail->end;
        }

        if (writable == 0) {
            size_t cap = remain > buf->chunk_size ? remain : buf->chunk_size;
            chain_buffer_node_t *node = alloc_node(cap);
            if (!node) {
                errno = ENOMEM;
                return -1;
            }

            if (buf->tail) {
                buf->tail->next = node;
                buf->tail = node;
            } else {
                buf->head = node;
                buf->tail = node;
            }
            tail = node;
            writable = tail->cap;
        }

        size_t n = remain < writable ? remain : writable;
        memcpy(tail->data + tail->end, src, n);
        tail->end += n;
        src += n;
        remain -= n;
        buf->total_len += n;
    }

    return 0;
}

size_t chain_buffer_drain(chain_buffer_t *buf, size_t len) {
    if (!buf || len == 0 || buf->total_len == 0) {
        return 0;
    }

    size_t remain = len;
    size_t drained = 0;

    while (remain > 0 && buf->head) {
        chain_buffer_node_t *node = buf->head;
        size_t avail = node->end - node->start;

        if (remain < avail) {
            node->start += remain;
            buf->total_len -= remain;
            drained += remain;
            break;
        }

        remain -= avail;
        drained += avail;
        buf->total_len -= avail;
        buf->head = node->next;
        if (!buf->head) {
            buf->tail = NULL;
        }
        free(node);
    }

    return drained;
}

const uint8_t *chain_buffer_linearize(chain_buffer_t *buf, size_t *out_len) {
    if (!buf) {
        return NULL;
    }

    if (out_len) {
        *out_len = buf->total_len;
    }

    if (buf->total_len == 0) {
        return NULL;
    }

    if (buf->head == buf->tail && buf->head) {
        return buf->head->data + buf->head->start;
    }

    if (buf->linear_cap < buf->total_len) {
        uint8_t *new_cache = (uint8_t *)realloc(buf->linear_cache, buf->total_len);
        if (!new_cache) {
            return NULL;
        }
        buf->linear_cache = new_cache;
        buf->linear_cap = buf->total_len;
    }

    size_t offset = 0;
    for (chain_buffer_node_t *node = buf->head; node; node = node->next) {
        size_t avail = node->end - node->start;
        if (avail == 0) {
            continue;
        }
        memcpy(buf->linear_cache + offset, node->data + node->start, avail);
        offset += avail;
    }

    return buf->linear_cache;
}

ssize_t chain_buffer_send_fd(chain_buffer_t *buf, int fd, int flags) {
    if (!buf) {
        errno = EINVAL;
        return -1;
    }
    if (buf->total_len == 0 || !buf->head) {
        return 0;
    }

    struct iovec iov[CHAINBUFFER_MAX_IOV];
    size_t iovcnt = 0;

    for (chain_buffer_node_t *node = buf->head;
         node && iovcnt < CHAINBUFFER_MAX_IOV;
         node = node->next) {
        size_t avail = node->end - node->start;
        if (avail == 0) {
            continue;
        }
        iov[iovcnt].iov_base = (void *)(node->data + node->start);
        iov[iovcnt].iov_len = avail;
        iovcnt++;
    }

    if (iovcnt == 0) {
        return 0;
    }

    struct msghdr msg;
    memset(&msg, 0, sizeof(msg));
    msg.msg_iov = iov;
    msg.msg_iovlen = iovcnt;

    ssize_t n = sendmsg(fd, &msg, flags);
    if (n > 0) {
        chain_buffer_drain(buf, (size_t)n);
    }

    return n;
}