#include "byteio.h"
#include "memory.h"
#include "arena.h"

/* ========================== *
 * Writer
 * ========================== */

INTERNAL b32 write_check_overflow(struct byte_writer *bw, i64 amount)
{
    b32 overflowed = bw->overflowed;
    if (overflowed || amount < 0) {
        overflowed = true;
        bw->overflowed = true;
    } else {
        u8 *end = bw->buff.text + bw->buff.len;
        u8 *new_pos = bw->at + amount;
        i64 new_space_left = end - new_pos;
        if (new_space_left < 0) {
            struct arena *arena = bw->arena;
            if (arena) {
                if ((arena->base + arena->pos) == end) {
                    arena_push_array(arena, u8, -new_space_left);
                    bw->buff.len += -new_space_left;
                } else {
                    /* Writer memory must be contiguous in arena */
                    overflowed = true;
                    bw->overflowed = true;
                }
            } else {
                overflowed = true;
                bw->overflowed = true;
            }
        }
    }
    ASSERT(!overflowed);
    return overflowed;
}

INTERNAL void write(struct byte_writer *bw, void *v, u64 size)
{
    if (write_check_overflow(bw, size)) {
        return;
    }
    MEMCPY(bw->at, v, size);
    bw->at += size;
}

struct byte_writer bw_from_buffer(struct string buff)
{
    struct byte_writer bw = ZI;
    bw.buff = buff;
    bw.at = buff.text;
    return bw;
}

/* Returns a writer that will allocate to arena instead of overflowing (writes must stay contiguous in arena) */
struct byte_writer bw_from_arena(struct arena *arena)
{
    struct byte_writer bw = ZI;
    bw.arena = arena;
    bw.buff.text = arena->base + arena->pos;
    bw.buff.len = 0;
    bw.at = bw.buff.text;
    return bw;
}

/* Seeks forward and returns a new writer pointing to the skipped bytes */
struct byte_writer bw_branch(struct byte_writer *bw, u64 size)
{
    struct string buff = STRING(size, bw->at);
    struct byte_writer branch = bw_from_buffer(buff);
    bw_seek(bw, size);
    branch.overflowed = bw->overflowed;
    return branch;
}

void bw_seek(struct byte_writer *bw, u64 amount)
{
    if (write_check_overflow(bw, amount)) {
        return;
    }
    bw->at += amount;
}

void bw_seek_to(struct byte_writer *bw, u64 pos)
{
    if (write_check_overflow(bw, (i64)pos - (i64)(bw->at - bw->buff.text))) {
        return;
    }
    bw->at = bw->buff.text + pos;
}

void bw_write_buffer(struct byte_writer *bw, struct string buff)
{
    write(bw, buff.text, buff.len);
}

void bw_write_u8(struct byte_writer *bw, u8 v)
{
    write(bw, &v, sizeof(v));
}

void bw_write_u16(struct byte_writer *bw, u16 v)
{
    write(bw, &v, sizeof(v));
}

void bw_write_u32(struct byte_writer *bw, u32 v)
{
    write(bw, &v, sizeof(v));
}

void bw_write_u64(struct byte_writer *bw, u64 v)
{
    write(bw, &v, sizeof(v));
}

void bw_write_i8(struct byte_writer *bw, i8 v)
{
    write(bw, &v, sizeof(v));
}

void bw_write_i16(struct byte_writer *bw, i16 v)
{
    write(bw, &v, sizeof(v));
}

void bw_write_i32(struct byte_writer *bw, i32 v)
{
    write(bw, &v, sizeof(v));
}

void bw_write_i64(struct byte_writer *bw, i64 v)
{
    write(bw, &v, sizeof(v));
}

void bw_write_var_uint(struct byte_writer *bw, u64 v)
{
    /* TODO: real varint write */
    bw_write_u64(bw, v);
}

void bw_write_var_sint(struct byte_writer *bw, i64 v)
{
    /* TODO: real varint write */
    bw_write_i64(bw, v);
}

void bw_write_f32(struct byte_writer *bw, f32 v)
{
    write(bw, &v, sizeof(v));
}

void bw_write_f64(struct byte_writer *bw, f64 v)
{
    write(bw, &v, sizeof(v));
}

void bw_write_v2(struct byte_writer *bw, struct v2 v)
{
    bw_write_f32(bw, v.x);
    bw_write_f32(bw, v.y);
}

void bw_write_string(struct byte_writer *bw, struct string str)
{
    bw_write_var_uint(bw, str.len);
    bw_write_buffer(bw, str);
}

/* ========================== *
 * Reader
 * ========================== */

INTERNAL b32 read_check_overflow(struct byte_reader *br, i64 amount)
{
    b32 overflowed = br->overflowed;
    if (overflowed || amount < 0 || (br->at + amount) > (br->buff.text + br->buff.len) ){
        ASSERT(false);
        br->overflowed = true;
        overflowed = true;
    }
    return overflowed;
}

INTERNAL void read(struct byte_reader *br, void *dst, u64 size)
{
    if (read_check_overflow(br, size)) {
        MEMZERO(dst, size);
    }
    MEMCPY(dst, br->at, size);
    br->at += size;
}

struct byte_reader br_from_buffer(struct string buff)
{
    struct byte_reader br = ZI;
    br.buff = buff;
    br.at = buff.text;
    return br;
}

/* Returns pointer to old position, or NULL on overflow */
void *br_seek(struct byte_reader *br, u64 amount)
{
    if (read_check_overflow(br, amount)) {
        return NULL;
    }
    void *ptr = br->at;
    br->at += amount;
    return ptr;
}

/* Returns pointer to old position, or NULL on overflow */
void *br_seek_to(struct byte_reader *br, u64 pos)
{
    if (read_check_overflow(br, (i64)pos - (i64)(br->at - br->buff.text))) {
        return NULL;
    }
    void *ptr = br->at;
    br->at = br->buff.text + pos;
    return ptr;
}

/* Will fill buff with zeroes on overflow */
void br_read_to_buffer(struct byte_reader *br, struct string buff)
{
    read(br, buff.text, buff.len);
}

u8 br_read_u8(struct byte_reader *br)
{
    u8 res;
    read(br, &res, sizeof(res));
    return res;
}

u16 br_read_u16(struct byte_reader *br)
{
    u16 res;
    read(br, &res, sizeof(res));
    return res;
}

u32 br_read_u32(struct byte_reader *br)
{
    u32 res;
    read(br, &res, sizeof(res));
    return res;
}

u64 br_read_u64(struct byte_reader *br)
{
    u64 res;
    read(br, &res, sizeof(res));
    return res;
}

i8 br_read_i8(struct byte_reader *br)
{
    i8 res;
    read(br, &res, sizeof(res));
    return res;
}

i16 br_read_i16(struct byte_reader *br)
{
    i16 res = 0;
    read(br, &res, sizeof(res));
    return res;
}

i32 br_read_i32(struct byte_reader *br)
{
    i32 res = 0;
    read(br, &res, sizeof(res));
    return res;
}

i64 br_read_i64(struct byte_reader *br)
{
    i64 res;
    read(br, &res, sizeof(res));
    return res;
}

u64 br_read_var_uint(struct byte_reader *br)
{
    /* TODO: real variable length read */
    return br_read_u64(br);
}

i64 br_read_var_sint(struct byte_reader *br)
{
    /* TODO: real variable length read */
    return br_read_i64(br);
}

f32 br_read_f32(struct byte_reader *br)
{
    f32 res;
    read(br, &res, sizeof(res));
    return res;
}

f64 br_read_f64(struct byte_reader *br)
{
    f64 res;
    read(br, &res, sizeof(res));
    return res;
}

struct v2 br_read_v2(struct byte_reader *br)
{
    struct v2 res;
    res.x = br_read_f32(br);
    res.y = br_read_f32(br);
    return res;
}

struct string br_read_string(struct arena *arena, struct byte_reader *br)
{
    struct string res = ZI;
    u64 len = br_read_var_uint(br);
    u8 *text = br_seek(br, len);
    if (text != NULL) {
        res.len = len;
        res.text = arena_push_array(arena, u8, len);
        MEMCPY(res.text, text, len);
    }
    return res;
}