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

#if defined(__APPLE__) || defined(__WIN32__)
/* Linux + Kos define this, OSX does not, so just use malloc there */
static inline void* memalign(size_t alignment, size_t size) {
    return malloc(size);
}
#else
    #include <malloc.h>
#endif

#include "aligned_vector.h"

void aligned_vector_init(AlignedVector* vector, unsigned int element_size) {
    vector->size = vector->capacity = 0;
    vector->element_size = element_size;
    vector->data = NULL;

    /* Reserve some initial capacity */
    aligned_vector_reserve(vector, ALIGNED_VECTOR_INITIAL_CAPACITY);
}

#define VECTOR_ALIGN_COUNT 256
#define _VECTOR_ALIGN_COUNT (VECTOR_ALIGN_COUNT - 1)

void aligned_vector_reserve(AlignedVector* vector, unsigned int element_count) {
    if(element_count <= vector->capacity) {
        return;
    }
    element_count = ((element_count+_VECTOR_ALIGN_COUNT) & ~_VECTOR_ALIGN_COUNT);
   
    unsigned int original_byte_size = vector->size * vector->element_size;
    unsigned int new_byte_size = element_count * vector->element_size;
    unsigned char* original_data = vector->data;
    vector->data = (unsigned char*) memalign(0x20, new_byte_size);

    if(original_data) {
        if( !(*vector->data % 32) && !(*original_data % 32)){
            if (original_byte_size % 4)
                original_byte_size = (original_byte_size & 0xfffffffc) + 4;
            sq_cpy(vector->data, original_data, original_byte_size);
        } else {
            memcpy(vector->data, original_data, original_byte_size);
        }
	    
        free(original_data);
    }

    vector->capacity = element_count;
}

void* aligned_vector_push_back(AlignedVector* vector, const void* objs, unsigned int count) {
    /* Resize enough room */

    unsigned int initial_size = vector->size;
    aligned_vector_resize(vector, vector->size + count);

    unsigned char* dest = vector->data + (vector->element_size * initial_size);

    /* Copy the objects in */
    memcpy(dest, objs, vector->element_size * count);
    //sq_cpy(dest, objs, ((vector->element_size * count) & 0xfffffffc) + 4);

    return dest;
}

void* aligned_vector_resize(AlignedVector* vector, const unsigned int element_count) {
    unsigned int previousCount = vector->size;

    /* Don't change memory when resizing downwards, just change the size */
    if(element_count <= vector->size) {
        vector->size = element_count;
        return NULL;
    }

    if(vector->capacity < element_count) {
        aligned_vector_reserve(vector, element_count);
    }

    vector->size = element_count;

    if(previousCount < vector->size) {
        return aligned_vector_at(vector, previousCount);
    } else {
        return NULL;
    }
}

void* aligned_vector_at(const AlignedVector* vector, const unsigned int index) {
    assert(index < vector->size);
    return &vector->data[index * vector->element_size];
}

void* aligned_vector_back(AlignedVector* vector) {
    return aligned_vector_at(vector, vector->size - 1);
}

void* aligned_vector_extend(AlignedVector* vector, const unsigned int additional_count) {
    const unsigned int current = vector->size;
    aligned_vector_resize(vector, vector->size + additional_count);
    return aligned_vector_at(vector, current);
}

void aligned_vector_clear(AlignedVector* vector) {
    vector->size = 0;
}

void aligned_vector_shrink_to_fit(AlignedVector* vector) {
    if(vector->size == 0) {
        free(vector->data);
        vector->data = NULL;
        vector->capacity = 0;
    } else {
        unsigned int new_byte_size = vector->size * vector->element_size;
        unsigned char* original_data = vector->data;
        vector->data = (unsigned char*) memalign(0x20, new_byte_size);

        if(original_data) {
            memcpy(vector->data, original_data, new_byte_size);
            //sq_cpy(vector->data, original_data, ((new_byte_size) & 0xfffffffc) + 4);
            
            free(original_data);
        }

        vector->capacity = vector->size;
    }
}

void aligned_vector_cleanup(AlignedVector* vector) {
    aligned_vector_clear(vector);
    aligned_vector_shrink_to_fit(vector);
}