#include <string.h>
#include <assert.h>
#include <stdlib.h>
#include "ucore/mbuf.h"

uint8_t *uc_mbuf_put(struct MBuf *mbuf, uint32_t size)
{
    uint8_t *data = mbuf->tail;

    if (uc_mbuf_tailroom(mbuf) < size) {
        return NULL;
    }

    mbuf->tail += size;

    return data;
}

uint8_t *uc_mbuf_push(struct MBuf *mbuf, uint32_t size)
{
    if (uc_mbuf_headroom(mbuf) < size) {
        return NULL;
    }

    mbuf->head -= size;

    return mbuf->head;
}

uint8_t *uc_mbuf_pull(struct MBuf *mbuf, uint32_t size)
{
    uint8_t *data = mbuf->head;

    if (uc_mbuf_len(mbuf) < size) {
        return NULL;
    }

    mbuf->head += size;

    return data;
}

void uc_mbuf_init(struct MBuf *mbuf, uint8_t *restrict buf, uint32_t len, uint32_t headroom, uint32_t flags)
{
    mbuf->bufstart  = buf;
    mbuf->p1        = mbuf->p2 = mbuf->p3 = mbuf->p4 = NULL;
    mbuf->cookie    = NULL;
    mbuf->flags     = flags;
    mbuf->allocated = len + headroom;
    mbuf->head      = mbuf->bufstart + headroom;
    mbuf->tail      = mbuf->head;
    mbuf->entry.next  = NULL;
    mbuf->entry.prev  = NULL;
}

struct MBuf *uc_mbuf_new_hr(uint32_t len, uint32_t headroom)
{
    struct MBuf *mbuf;
    uint8_t *buf;
    
    mbuf = malloc(sizeof *mbuf);
    if (mbuf == NULL) 
        return mbuf;

    buf = malloc(len + headroom);
    if (buf == NULL) {
        free(mbuf);
        return NULL;
    }

    uc_mbuf_init(mbuf, buf, len, headroom, UC_MBUF_FL_MALLOC);

    return mbuf;
}

struct MBuf *uc_mbuf_new_inline_hr(uint32_t len, uint32_t headroom)
{
    struct MBuf *mbuf;
    
    mbuf = malloc(sizeof *mbuf + len + headroom);
    if (mbuf == NULL) 
        return mbuf;

    uc_mbuf_init(mbuf, mbuf->inline_data, len, headroom, UC_MBUF_FL_MALLOC_INLINE);

    return mbuf;
}

void uc_mbuf_free(struct MBuf *mbuf)
{
    if (mbuf != NULL) {
        if (mbuf->flags & (UC_MBUF_FL_MALLOC | UC_MBUF_FL_MALLOC_BUF))
            free(mbuf->bufstart);
        
        if (mbuf->flags & (UC_MBUF_FL_MALLOC | UC_MBUF_FL_MALLOC_INLINE))
            free(mbuf);
    }
}

void uc_mbuf_zero(struct MBuf *mbuf)
{
    memset(mbuf->bufstart, 0, mbuf->allocated);
}

struct MBuf *uc_mbuf_copy_data(struct MBuf *mbuf)
{
    uint32_t len;
    struct MBuf *copy;
    uint8_t *data;

    len = uc_mbuf_len(mbuf);

    copy = uc_mbuf_new_hr(len, uc_mbuf_headroom(mbuf));
    if (copy == NULL) { 
        return NULL; 
    }

    data = uc_mbuf_put(copy, len);
    memcpy(data, mbuf->head, len);
    
    return copy;
}

int uc_mbuf_extend_tailroom(struct MBuf *mbuf, uint32_t len)
{
    uint32_t tot_len;
    uint32_t headroom_len;
    uint32_t data_len;
    uint32_t tailroom_len;
    uint8_t *new_data;
    uint32_t p1_off = 0;
    uint32_t p2_off = 0;
    uint32_t p3_off = 0;
    uint32_t p4_off = 0;

    if ((mbuf->flags & (UC_MBUF_FL_MALLOC | UC_MBUF_FL_MALLOC_BUF)) == 0) {
        return -1;
    }
    
    headroom_len = uc_mbuf_headroom(mbuf);
    data_len     = uc_mbuf_len(mbuf);
    tailroom_len = uc_mbuf_tailroom(mbuf);
    tot_len = headroom_len +
                  data_len     +
                  tailroom_len;

    if (tot_len + len < tot_len) {
        //overflowed
        return -2;
    }
    
    tot_len += len;
    
    //need to re-adjust p1-p4 after the realloc
    //This assumes the user controlled p1-p4 really are
    //set to point into the data, but we document they have
    //to be.
    if (mbuf->p1 != NULL) {
        p1_off = (uint32_t)(mbuf->p1 - mbuf->bufstart);
    }
    if (mbuf->p2 != NULL) {
        p2_off = (uint32_t)(mbuf->p2 - mbuf->bufstart);
    }
    if (mbuf->p3 != NULL) {
        p3_off = (uint32_t)(mbuf->p3 - mbuf->bufstart);
    }
    if (mbuf->p4 != NULL) {
        p4_off = (uint32_t)(mbuf->p4 - mbuf->bufstart);
    }

    assert(tot_len >= mbuf->allocated);

    new_data = realloc(mbuf->bufstart, tot_len);
    if (new_data == NULL) {
        return -3;
    }

    mbuf->allocated = tot_len;
    mbuf->bufstart = new_data;
    mbuf->head = mbuf->bufstart + headroom_len;
    mbuf->tail = mbuf->bufstart + headroom_len + data_len;

    assert(uc_mbuf_len(mbuf) == data_len);
    assert(uc_mbuf_tailroom(mbuf) >= tailroom_len);

    if (mbuf->p1 != NULL) {
        mbuf->p1 = mbuf->bufstart + p1_off;
    }
    if (mbuf->p2 != NULL) {
        mbuf->p2 = mbuf->bufstart + p2_off;
    }
    if (mbuf->p3 != NULL) {
        mbuf->p3 = mbuf->bufstart + p3_off;
    }
    if (mbuf->p4 != NULL) {
        mbuf->p4 = mbuf->bufstart + p4_off;
    }

    return 0;

}