#include <stdlib.h>
#include <stdio.h>
#import "memory.h"

// Finds the real beginning of the memory block (for dealloc/realloc)
#define REAL_BLOCK(mem) (MEM_METADATA(mem)->real_start)

// Retrieves a pointer to the metadata field for a block of memory.
#define MEM_METADATA(mem) ((memory_metadata *)mem - 1)

// Gets the retain count for a block of memory.
#define RCOUNT(mem) (MEM_METADATA(mem)->retain_count)

// If the memory has been allocated with a custom struct, this will return a pointer to it.
#define CUSTOM_STRUCT(mem, type) ((type *)MEM_METADATA(mem) - 1)

typedef struct {
    unsigned int retain_count;
    dealloc_function dealloc;
    void *real_start;
} memory_metadata;

void *alloc_with_custom_struct_and_dealloc(size_t s, size_t struct_size, dealloc_function d)
{
    // Allocate enough space for everything.
    void *real_block = malloc(s + sizeof(memory_metadata) + struct_size);
    memory_metadata *mem = (memory_metadata *)((char *)real_block + struct_size);
    
    mem->retain_count = 1;
    mem->dealloc = d;
    mem->real_start = real_block;
    
    return (void *)(mem + 1);
}

void *alloc_with_custom_struct(size_t s, size_t struct_size)
{
    return alloc_with_custom_struct_and_dealloc(s, struct_size, NULL);
}

void *alloc_with_dealloc(size_t s, dealloc_function d)
{
    return alloc_with_custom_struct_and_dealloc(s, 0, NULL);
}

void *alloc(size_t s)
{
    return alloc_with_dealloc(s, NULL);
}

void dealloc(void *mem)
{
    // Check if we need to call a dealloc function.
    dealloc_function d = MEM_METADATA(mem)->dealloc;
    if (d != NULL) {
        d(mem);
    }
    free(REAL_BLOCK(mem));
}

typedef struct {
    // The number of elements in the array.
    unsigned int count;
    
    // The amount of memory (in number of elements) that has been allocated for the array.
    unsigned int size;
    
    // The amount of memory to be allocated per element.
    size_t cell_size;
} array_metadata;

#define ARRAY_STRUCT(array) CUSTOM_STRUCT(array, array_metadata)

void *alloc_array(size_t cell_size, unsigned int initial_count)
{
    void *array = alloc_with_custom_struct(cell_size * initial_count, sizeof(array_metadata));
    ARRAY_STRUCT(array)->count = initial_count;
    ARRAY_STRUCT(array)->size = initial_count;
    ARRAY_STRUCT(array)->cell_size = cell_size;
    return array;
}

unsigned int array_count(void *array)
{
    return ARRAY_STRUCT(array)->count;
}

void *resize_array(void *array, unsigned int new_count)
{
    if (new_count > ARRAY_STRUCT(array)->size) {
        unsigned int new_size = ARRAY_STRUCT(array)->size;
        while (new_size < new_count) {
            if (new_size == 0) {
                new_size = 1;
            } else {
                new_size <<= 1;
            }
        }
        //array = realloc(REAL_BLOCK(array), new_size * ARRAY_STRUCT(array)->cell_size);
        void * new_array = alloc_array(ARRAY_STRUCT(array)->cell_size, new_size);
        int i;
        for (i=0; i<ARRAY_STRUCT(array)->size*ARRAY_STRUCT(array)->cell_size; i++) {
            ((char *)new_array)[i] = ((char *)array)[i];
        }
        RCOUNT(new_array) = RCOUNT(array);
        MEM_METADATA(new_array)->dealloc = MEM_METADATA(array)->dealloc;
        free(REAL_BLOCK(array));
        array = new_array;
    }
    ARRAY_STRUCT(array)->count = new_count;
    return array;
}

unsigned int retain_count(void *mem)
{
    return RCOUNT(mem);
}

void *retain(void *mem)
{
    RCOUNT(mem)++;
    return mem;
}

void release(void *mem)
{
    if (RCOUNT(mem)-- <= 0) {
        dealloc(mem);
    }
}

typedef struct __a_node {
    void *mem;
    struct __a_node *next;
} autorelease_node;

typedef struct __a_stack {
    autorelease_node *node;
    struct __a_stack *up;
} autorelease_stack;

autorelease_stack *auto_stack = NULL;

void push_autorelease()
{
    autorelease_stack *new_stack = (autorelease_stack *)malloc(sizeof(auto_stack));
    new_stack->up = auto_stack;
    new_stack->node = NULL;
    auto_stack = new_stack;
}

void pop_autorelease()
{
    if (auto_stack == NULL) {
        // Error: stack underflow.
    } else {
        // Iterate through the nodes, releasing as we go (until we get to a NULL node which means we're done).
        while (auto_stack->node != NULL) {
            autorelease_node *dead = auto_stack->node;
            auto_stack->node = auto_stack->node->next;
            release(dead->mem); 
            free(dead);
        }
        // Deallocate the current autorelease pool and replace it with the one "under" it on the stack.
        autorelease_stack *parent = auto_stack->up;
        free(auto_stack);
        auto_stack = parent;
    }
}

void *autorelease(void *mem)
{
    // Allocate and add a node to the end of the current autorelease pool.
    autorelease_node *new_node = malloc(sizeof(autorelease_node));
    new_node->mem = mem;
    new_node->next = NULL;
    if (auto_stack->node != NULL) {
        new_node->next = auto_stack->node->next;
    }
    auto_stack->node = new_node;
    return mem;
}