#include "../include/async_image.h"
#include "../include/alloc.h"
#include "../include/hashmap.h"
#include <pthread.h>
#include <semaphore.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef struct {
    mgui_async_image *async_image;
} async_image_task;

typedef struct {
    async_image_task *tasks;
    size_t capacity;
    size_t start_index;
    size_t num_occupied_slots;
} async_image_task_fifo;

static bool initialized = false;
static pthread_t images_thread;
static mgui_hashmap images;
static pthread_mutex_t mutex;
static sem_t task_sem;
static async_image_task_fifo tasks;
static uint32_t update_counter = 0;

static size_t align_up(size_t value, size_t alignment) {
    size_t over = value % alignment;
    if(over == 0)
        return value;
    return value - over + alignment;
}

static void async_image_task_fifo_init(async_image_task_fifo *self) {
    self->tasks = NULL;
    self->capacity = 0;
    self->start_index = 0;
    self->num_occupied_slots = 0;
}

static void async_image_task_fifo_deinit(async_image_task_fifo *self) {
    mgui_free(self->tasks);
    self->tasks = NULL;
    self->capacity = 0;
    self->start_index = 0;
    self->num_occupied_slots = 0;
}

static void async_image_task_fifo_ensure_capacity(async_image_task_fifo *self, size_t new_capacity) {
    if(self->capacity >= new_capacity)
        return;

    size_t capacity = self->capacity;
    if(capacity == 0)
        capacity = 8;

    while(capacity < new_capacity) {
        capacity = capacity + (capacity >> 1); /* capacity *= 1.5 */
    }

    capacity = align_up(capacity, sizeof(async_image_task));
    self->tasks = mgui_realloc(self->tasks, capacity);

    /* Straighten circular buffer. For example: convert [4][5][][][1][2][3][][] to [][][][][1][2][3][4][5] */
    if(self->capacity > 0) {
        const size_t num_allocated_slots = self->capacity / sizeof(async_image_task);
        const size_t end_index = (self->start_index + self->num_occupied_slots) % num_allocated_slots;
        if(end_index < self->start_index) {
            for(size_t i = 0; i <= end_index; ++i) {
                self->tasks[self->num_occupied_slots + i] = self->tasks[end_index - i];
            }
        }
    }
    self->capacity = capacity;
}

static void async_image_task_fifo_append(async_image_task_fifo *self, async_image_task *task) {
    async_image_task_fifo_ensure_capacity(self, (self->num_occupied_slots + 1) * sizeof(async_image_task));
    const size_t num_allocated_slots = self->capacity / sizeof(async_image_task);
    const size_t insert_index = (self->start_index + self->num_occupied_slots) % num_allocated_slots;
    self->tasks[insert_index] = *task;
    ++self->num_occupied_slots;
}

/* pop first item if not empty. returns false if empty */
static bool async_image_task_fifo_pop(async_image_task_fifo *self, async_image_task *task_out) {
    if(self->num_occupied_slots == 0)
        return false;

    *task_out = self->tasks[self->start_index];
    const size_t num_allocated_slots = self->capacity / sizeof(async_image_task);
    self->start_index = (self->start_index + 1) % num_allocated_slots;
    --self->num_occupied_slots;
    return true;
}

/* Return false from |callback| to stop the iteration */
static void async_image_task_fifo_for_each(async_image_task_fifo *self, bool (*callback)(async_image_task *task, void *userdata), void *userdata) {
    if(self->capacity == 0)
        return;

    const size_t num_allocated_slots = self->capacity / sizeof(async_image_task);
    for(size_t i = 0; i < self->num_occupied_slots; ++i) {
        const size_t index = (self->start_index + i) % num_allocated_slots;
        if(!callback(&self->tasks[index], userdata))
            return;
    }
}

static void* images_thread_callback(void *arg) {
    (void)arg;
    async_image_task task;
    while(initialized) {
        sem_wait(&task_sem);
        if(!initialized)
            break;

        pthread_mutex_lock(&mutex);
        const bool has_task = async_image_task_fifo_pop(&tasks, &task);
        const bool empty_task = !has_task || !task.async_image;
        if(!empty_task)
            ++task.async_image->ref_count;
        pthread_mutex_unlock(&mutex);

        if(empty_task)
            continue;

        /* TODO: Create thumbnail (resize to target size) and load that instead */
        if(mgl_image_load_from_file(&task.async_image->image, task.async_image->filepath) == 0)
            task.async_image->state = MGUI_ASYNC_IMAGE_LOADED;
        else
            task.async_image->state = MGUI_ASYNC_IMAGE_FAILED_TO_LOAD;

        pthread_mutex_lock(&mutex);
        if(task.async_image->ref_count > 0)
            --task.async_image->ref_count;
        pthread_mutex_unlock(&mutex);
    }
    return NULL;
}

void mgui_async_image_init(void) {
    if(!initialized) {
        initialized = true;
        mgui_hashmap_init(&images);
        async_image_task_fifo_init(&tasks);

        if(sem_init(&task_sem, 0, 0) != 0) {
            fprintf(stderr, "mgui error: failed to initialize images semaphore\n");
            abort();
        }

        if(pthread_mutex_init(&mutex, NULL) != 0) {
            fprintf(stderr, "mgui error: failed to initialize images mutex\n");
            abort();
        }

        if(pthread_create(&images_thread, NULL, images_thread_callback, NULL) != 0) {
            fprintf(stderr, "mgui error: failed to initialize images thread\n");
            abort();
        }
    }
}

static bool images_free(void *value, void *userdata) {
    (void)userdata;
    mgui_async_image *async_image = value;
    mgl_image_unload(&async_image->image);
    mgl_texture_unload(&async_image->texture);
    mgui_free(async_image->filepath);
    mgui_free(async_image);
    return true;
}

void mgui_async_image_deinit(void) {
    if(initialized) {
        initialized = false;
        sem_post(&task_sem);
        pthread_join(images_thread, NULL);
        pthread_mutex_destroy(&mutex);
        sem_destroy(&task_sem);

        async_image_task task;
        while(async_image_task_fifo_pop(&tasks, &task)) {}

        async_image_task_fifo_deinit(&tasks);
        mgui_hashmap_for_each(&images, images_free, NULL);
        mgui_hashmap_deinit(&images);
    }
}

static bool remove_unreferenced_tasks(async_image_task *task, void *userdata) {
    (void)userdata;
    if(task->async_image && task->async_image->updated != update_counter)
        task->async_image = NULL;
    return true;
}

static bool images_unload_unreferenced(void *value, void *userdata) {
    (void)userdata;
    mgui_async_image *async_image = value;
    if(async_image->updated != update_counter) {
        bool remove_image = false;
        if(async_image->state != MGUI_ASYNC_IMAGE_LOADING) {
            async_image->state = MGUI_ASYNC_IMAGE_UNLOADED;
            mgl_image_unload(&async_image->image);
            mgl_texture_unload(&async_image->texture);

            pthread_mutex_lock(&mutex);
            /* TODO: Check if this actually happens */
            if(async_image->ref_count == 0) {
                mgui_free(async_image->filepath);
                mgui_free(async_image);
                remove_image = true;
            }
            pthread_mutex_unlock(&mutex);
        }
        return remove_image;
    } else {
        return false;
    }
}

void mgui_async_image_unload_unreferenced(void) {
    pthread_mutex_lock(&mutex);
    async_image_task_fifo_for_each(&tasks, remove_unreferenced_tasks, NULL);
    pthread_mutex_unlock(&mutex);

    mgui_hashmap_for_each_erase(&images, images_unload_unreferenced, NULL);

    ++update_counter;
}

mgui_async_image* mgui_async_image_get_by_path(const char *filepath) {
    const size_t filepath_len = strlen(filepath);

    pthread_mutex_lock(&mutex);
    void *value;
    if(!mgui_hashmap_get(&images, filepath, filepath_len, &value)) {
        mgui_async_image *async_image = mgui_alloc(sizeof(mgui_async_image));
        async_image->image.data = NULL;
        async_image->texture.id = 0;
        async_image->hash = 0;
        async_image->state = MGUI_ASYNC_IMAGE_LOADING;
        async_image->updated = update_counter;
        async_image->ref_count = 0;
        async_image->filepath = mgui_alloc(filepath_len + 1);
        memcpy(async_image->filepath, filepath, filepath_len);
        async_image->filepath[filepath_len] = '\0';
        mgui_hashmap_insert(&images, filepath, filepath_len, async_image, &async_image->hash);
        value = async_image;

        async_image_task task = {
            .async_image = async_image
        };
        async_image_task_fifo_append(&tasks, &task);
        sem_post(&task_sem);
    }

    mgui_async_image *async_image = value;
    ++async_image->ref_count;
    pthread_mutex_unlock(&mutex);

    return async_image;
}

void mgui_async_image_update(mgui_async_image *self) {
    if(self->state == MGUI_ASYNC_IMAGE_LOADED) {
        self->state = MGUI_ASYNC_IMAGE_APPLIED;
        if(mgl_texture_init(&self->texture) == 0)
            mgl_texture_load_from_image(&self->texture, &self->image, NULL);
        mgl_image_unload(&self->image);
    } else if(self->state == MGUI_ASYNC_IMAGE_UNLOADED) {
        self->state = MGUI_ASYNC_IMAGE_LOADING;
        async_image_task task = {
            .async_image = self
        };
        async_image_task_fifo_append(&tasks, &task);
        sem_post(&task_sem);
    }
    self->updated = update_counter;
}

void mgui_async_image_unref(mgui_async_image *self) {
    pthread_mutex_lock(&mutex);
    if(self->ref_count > 0)
        --self->ref_count;
    pthread_mutex_unlock(&mutex);
}