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

#include "nvs_flash.h"
#include "nvs.h"

#include "lib/video/vram.h"
#include "boot/answer_code.h"
#include "boot/entry.h"

vram_manager_t *vram_sys = NULL;

int init_vram(apios_main_table_t *apios_table)
{
    (void)apios_table;

    if (vram_sys != NULL) return OS_WRN_NOT_NEED_INIT;

    int vram_kb[3] = {128, 512, 1024};
    int8_t val     = 1;

    nvs_handle_t h;
    if (nvs_open("config", NVS_READONLY, &h) == ESP_OK) {
        if (nvs_get_i8(h, "vram", &val) != ESP_OK || val < 0 || val > 2) {
            val = 1;
        }
        nvs_close(h);
    } else {
        if (load_log) printf("VRAM: No saved config, using default 512kb\n");
        val = 1;
    }

    vram_sys = (vram_manager_t*)malloc(sizeof(vram_manager_t));
    if (!vram_sys) return OS_ERR_NO_MEMORY;

    memset(vram_sys, 0, sizeof(vram_manager_t));

    vram_sys->capacity    = (uint32_t)(vram_kb[val] * 1024);
    vram_sys->next_id     = 1;
    vram_sys->max_blocks  = 256;
    vram_sys->blocks_count = 1;

    vram_sys->base_ptr = (uint8_t*)malloc(vram_sys->capacity);
    vram_sys->blocks   = (vram_block_t*)malloc(sizeof(vram_block_t) * vram_sys->max_blocks);

    if (!vram_sys->base_ptr || !vram_sys->blocks) {
        if (vram_sys->base_ptr) free(vram_sys->base_ptr);
        if (vram_sys->blocks)   free(vram_sys->blocks);
        free(vram_sys);
        vram_sys = NULL;
        return OS_ERR_NO_MEMORY;
    }

    vram_sys->blocks[0].offset  = 0;
    vram_sys->blocks[0].size    = vram_sys->capacity;
    vram_sys->blocks[0].is_free = true;
    vram_sys->blocks[0].id      = -1;

    if (load_log) {
        printf("VRAM: Initialized %d kb at %p\n",
               vram_kb[val],
               (void*)vram_sys->base_ptr);
    }

    return OS_OK;
}

vram_handle_t vram_alloc(uint32_t size)
{
    if (!vram_sys) return -1;

    size = (size + 3) & ~3;

    for (int i = 0; i < vram_sys->blocks_count; i++) {
        if (vram_sys->blocks[i].is_free && vram_sys->blocks[i].size >= size) {
            uint32_t remaining = vram_sys->blocks[i].size - size;

            vram_sys->blocks[i].size    = size;
            vram_sys->blocks[i].is_free = false;
            vram_sys->blocks[i].id      = vram_sys->next_id++;

            if (remaining > 0 && vram_sys->blocks_count < vram_sys->max_blocks) {
                for (int j = vram_sys->blocks_count; j > i + 1; j--) {
                    vram_sys->blocks[j] = vram_sys->blocks[j - 1];
                }
                vram_sys->blocks[i + 1].offset  = vram_sys->blocks[i].offset + size;
                vram_sys->blocks[i + 1].size    = remaining;
                vram_sys->blocks[i + 1].is_free = true;
                vram_sys->blocks[i + 1].id      = -1;
                vram_sys->blocks_count++;
            }

            return vram_sys->blocks[i].id;
        }
    }

    if (load_log) printf("VRAM: No space for %lu bytes! Need defrag.\n", (unsigned long)size);
    return VRAM_WRN_DEFRAG_REQ;
}

void* vram_get_ptr(vram_handle_t handle)
{
    if (!vram_sys || handle <= 0) return NULL;

    for (int i = 0; i < vram_sys->blocks_count; i++) {
        if (vram_sys->blocks[i].id == handle) {
            return (void*)(vram_sys->base_ptr + vram_sys->blocks[i].offset);
        }
    }
    return NULL;
}

void vram_free(vram_handle_t handle)
{
    if (!vram_sys) return;

    for (int i = 0; i < vram_sys->blocks_count; i++) {
        if (vram_sys->blocks[i].id == handle) {
            vram_sys->blocks[i].is_free = true;
            vram_sys->blocks[i].id      = -1;

            if (i + 1 < vram_sys->blocks_count && vram_sys->blocks[i + 1].is_free) {
                vram_sys->blocks[i].size += vram_sys->blocks[i + 1].size;
                for (int j = i + 1; j < vram_sys->blocks_count - 1; j++) {
                    vram_sys->blocks[j] = vram_sys->blocks[j + 1];
                }
                vram_sys->blocks_count--;
            }

            if (i > 0 && vram_sys->blocks[i - 1].is_free) {
                vram_sys->blocks[i - 1].size += vram_sys->blocks[i].size;
                for (int j = i; j < vram_sys->blocks_count - 1; j++) {
                    vram_sys->blocks[j] = vram_sys->blocks[j + 1];
                }
                vram_sys->blocks_count--;
            }

            return;
        }
    }
}

void vram_defrag(void)
{
    if (!vram_sys) return;

    uint32_t current_offset = 0;
    int      write_idx      = 0;

    for (int i = 0; i < vram_sys->blocks_count; i++) {
        if (!vram_sys->blocks[i].is_free) {
            if (vram_sys->blocks[i].offset != current_offset) {
                memmove(vram_sys->base_ptr + current_offset,
                        vram_sys->base_ptr + vram_sys->blocks[i].offset,
                        vram_sys->blocks[i].size);
                vram_sys->blocks[i].offset = current_offset;
            }

            vram_sys->blocks[write_idx] = vram_sys->blocks[i];
            current_offset += vram_sys->blocks[write_idx].size;
            write_idx++;
        }
    }

    vram_sys->blocks_count = write_idx;

    if (current_offset < vram_sys->capacity) {
        vram_sys->blocks[write_idx].offset  = current_offset;
        vram_sys->blocks[write_idx].size    = vram_sys->capacity - current_offset;
        vram_sys->blocks[write_idx].is_free = true;
        vram_sys->blocks[write_idx].id      = -1;
        vram_sys->blocks_count++;
    }

    if (load_log) {
        printf("VRAM: Defrag complete. Free: %lu bytes\n",
               (unsigned long)(vram_sys->capacity - current_offset));
    }
}

int vram_write(vram_handle_t handle, const void *src, uint32_t size)
{
    if (!vram_sys || handle <= 0 || !src) return -1;

    for (int i = 0; i < vram_sys->blocks_count; i++) {
        if (vram_sys->blocks[i].id == handle) {
            uint32_t write_size = (size < vram_sys->blocks[i].size)
                                  ? size
                                  : vram_sys->blocks[i].size;
            memcpy(vram_sys->base_ptr + vram_sys->blocks[i].offset, src, write_size);
            return OS_OK;
        }
    }

    return VRAM_ERR_WR_or_RD;
}

int vram_read(vram_handle_t handle, void *dst, uint32_t size)
{
    if (!vram_sys || handle <= 0 || !dst) return -1;

    for (int i = 0; i < vram_sys->blocks_count; i++) {
        if (vram_sys->blocks[i].id == handle) {
            uint32_t read_size = (size < vram_sys->blocks[i].size)
                                 ? size
                                 : vram_sys->blocks[i].size;
            memcpy(dst, vram_sys->base_ptr + vram_sys->blocks[i].offset, read_size);
            return OS_OK;
        }
    }

    return VRAM_ERR_WR_or_RD;
}