#include "sock_test.h"

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

/**
 * Grow buf->vecs as needed, to ensure that buf->write_vec is valid
 */
static err_t sock_test_grow_buf (struct io_buf *buf)
{
    size_t read_vec_offset = buf->read_vec ? (buf->read_vec - buf->vecs) : 0;
    size_t write_vec_offset = buf->write_vec ? (buf->write_vec - buf->vecs) : 0;
    struct io_vec *v;
    struct io_vec *vecs_tmp = buf->vecs;

    // don't grow if not full
    if (buf->vecs && buf->write_vec < buf->vecs + buf->count)
        return SUCCESS;

    // new size
    buf->count = buf->count * 2 + 1;

    // grow
    if ((buf->vecs = realloc(buf->vecs, buf->count * sizeof(struct io_vec))) == NULL) {
        // restore old value
        buf->vecs = vecs_tmp;

        return ERR_CALLOC;
    }

    // set vec
    buf->write_vec = buf->vecs + write_vec_offset;
    buf->read_vec = buf->vecs + read_vec_offset;

    // zero new vecs
    for (v = buf->write_vec; v < buf->vecs + buf->count; v++) {
        v->buf = NULL;
        v->len = 0;
    }

    // ok
    return SUCCESS;
}

static err_t sock_test_read (struct sock_stream *base_sock, void *buf_ptr, size_t *len)
{
    struct sock_test *sock = SOCK_FROM_BASE(base_sock, struct sock_test);
    struct io_buf *buf = &sock->recv_buf;
    struct io_vec *vec = buf->read_vec;
    
    // EOF/nonblock if we're past the end of the last vector
    if (!vec || vec == buf->vecs + buf->count || buf->off >= vec->len) {
        if (sock->nonblocking && !sock->eof) {
            // wait for more to be fed in
            *len = 0;
            return SUCCESS;

        } else {
            // EOF!
            return SET_ERROR(SOCK_TEST_ERR(sock), ERR_READ_EOF);
        }
    }
    
    // amount of data available in this iovec
    size_t available = vec->len - buf->off;

    // amount to read
    size_t to_read = *len;
    
    // trim down?
    if (to_read > available)
        to_read = available;

    // copy
    memcpy(buf_ptr, vec->buf + buf->off, to_read);
    
    // consumed the whole vec?
    if (to_read < available) {
        // move offset
        buf->off += to_read;

    } else {
        // next vector
        buf->read_vec++;
    }

    // update len
    *len = to_read;

    // ok
    return SUCCESS;
}

static err_t sock_test_write (struct sock_stream *base_sock, const void *buf_ptr, size_t *len)
{
    struct sock_test *sock = SOCK_FROM_BASE(base_sock, struct sock_test);
    struct io_buf *buf = &sock->send_buf;
    
    // ensure there's room
    assert(sock_test_grow_buf(buf) == SUCCESS);
    
    // the next buffer
    struct io_vec *vec = buf->write_vec;

    // store
    vec->len = *len;
    assert((vec->buf = malloc(vec->len)));
    memcpy(vec->buf, buf_ptr, vec->len);

    // move vec onwards
    buf->write_vec++;

    // ok
    return SUCCESS;
}

static err_t sock_test_event_init (struct sock_stream *base_sock)
{
    struct sock_test *sock = SOCK_FROM_BASE(base_sock, struct sock_test);

    // set the nonblocking flag
    sock->nonblocking = true;

    return SUCCESS;
}

static err_t sock_test_event_enable (struct sock_stream *base_sock, short mask)
{
    struct sock_test *sock = SOCK_FROM_BASE(base_sock, struct sock_test);

    return SUCCESS;
}

static void sock_test_release (struct sock_stream *base_sock)
{
    struct sock_test *sock = SOCK_FROM_BASE(base_sock, struct sock_test);

    sock_test_destroy(sock);
}

/*
 * Our sock_stream_type
 */
static struct sock_stream_type sock_test_type = {
    .methods                = {
        .read               = &sock_test_read,
        .write              = &sock_test_write,
        .event_init         = &sock_test_event_init,
        .event_enable       = &sock_test_event_enable,
        .release            = &sock_test_release,
    },
};

struct sock_test* sock_test_create (void)
{
    struct sock_test *sock;

    // allocate
    assert((sock = calloc(1, sizeof(*sock))));
    
    // initialize base with our sock_stream_type
    sock_stream_init(SOCK_TEST_BASE(sock), &sock_test_type);

    // ok
    return sock;
}

void sock_test_destroy (struct sock_test *sock)
{
    size_t i;
    struct io_buf *sbuf = &sock->send_buf, *rbuf = &sock->recv_buf;
    
    // free the send buffer
    for (i = 0; i < sbuf->count; i++) {
        free(sbuf->vecs[i].buf);
    }

    // free the buffer vector lists
    free(sbuf->vecs);
    free(rbuf->vecs);
    
    // free the sock itself
    free(sock);
}

void sock_test_set_recv_buffer (struct sock_test *sock, struct io_vec *vecs, size_t count, bool eof)
{
    struct io_buf *buf = &sock->recv_buf;

    // allocate + copy
    assert((buf->vecs = calloc(count, sizeof(struct io_vec))));
    memcpy(buf->vecs, vecs, count * sizeof(struct io_vec));
    
    // set
    buf->count = count;
    buf->read_vec = buf->vecs;
    buf->write_vec = buf->vecs + count;
    buf->off = 0;
    
    // set EOF flag?
    if (eof)
        sock->eof = true;
}

void sock_test_add_recv_vec (struct sock_test *sock, struct io_vec new_vec)
{
    struct io_buf *buf = &sock->recv_buf;

    // ensure there's room
    assert(sock_test_grow_buf(buf) == SUCCESS);
    
    // copy    
    *(buf->write_vec++) = new_vec;
}

void sock_test_get_send_data (struct sock_test *sock, char **buf_ptr, size_t *len_ptr)
{
    struct io_buf *buf = &sock->send_buf;
    size_t len = 0, i, off = 0;
    char *out;
    
    // calculate total size
    for (i = 0; i < buf->count; i++) {
        len += buf->vecs[i].len;
    }

    // alloc
    assert((out = malloc(len)));

    // copy
    for (i = 0; i < buf->count; i++) {
        memcpy(out + off, buf->vecs[i].buf, buf->vecs[i].len);
        off += buf->vecs[i].len;
    }
    
    // update return
    *buf_ptr = out;
    *len_ptr = len;
}