#include "pxl8_atlas.h"

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

#include "pxl8_color.h"
#include "pxl8_macros.h"

typedef struct pxl8_skyline_fit {
    bool found;
    u32 node_idx;
    pxl8_point pos;
} pxl8_skyline_fit;

typedef struct pxl8_skyline_node {
    i32 x, y, width;
} pxl8_skyline_node;

typedef struct pxl8_skyline {
    pxl8_skyline_node* nodes;
    u32 count;
    u32 capacity;
} pxl8_skyline;

struct pxl8_atlas {
    u32 height, width;
    u8* pixels;

    bool dirty;

    u32 entry_capacity, entry_count;
    pxl8_atlas_entry* entries;

    u32 free_capacity, free_count;
    u32* free_list;

    pxl8_skyline skyline;
};

static pxl8_skyline_fit pxl8_skyline_find_position(
    const pxl8_skyline* skyline,
    u32 atlas_w,
    u32 atlas_h,
    u32 rect_w,
    u32 rect_h
) {
    pxl8_skyline_fit result = {.found = false};
    i32 best_y = INT32_MAX;
    i32 best_x = 0;
    u32 best_idx = 0;

    for (u32 i = 0; i < skyline->count; i++) {
        i32 x = skyline->nodes[i].x;
        i32 y = skyline->nodes[i].y;

        if (x + (i32)rect_w > (i32)atlas_w) continue;

        i32 max_y = y;
        for (u32 j = i; j < skyline->count && skyline->nodes[j].x < x + (i32)rect_w; j++) {
            if (skyline->nodes[j].y > max_y) {
                max_y = skyline->nodes[j].y;
            }
        }

        if (max_y + (i32)rect_h > (i32)atlas_h) continue;

        if (max_y < best_y || (max_y == best_y && x < best_x)) {
            best_y = max_y;
            best_x = x;
            best_idx = i;
        }
    }

    if (best_y != INT32_MAX) {
        result.found = true;
        result.pos.x = best_x;
        result.pos.y = best_y;
        result.node_idx = best_idx;
    }

    return result;
}

static bool pxl8_skyline_add_rect(pxl8_skyline* skyline, pxl8_point pos, u32 w, u32 h) {
    u32 node_idx = 0;
    for (u32 i = 0; i < skyline->count; i++) {
        if (skyline->nodes[i].x == pos.x) {
            node_idx = i;
            break;
        }
    }

    u32 nodes_to_remove = 0;
    for (u32 i = node_idx; i < skyline->count; i++) {
        if (skyline->nodes[i].x < pos.x + (i32)w) {
            nodes_to_remove++;
        } else {
            break;
        }
    }

    if (skyline->count - nodes_to_remove + 1 > skyline->capacity) {
        u32 new_capacity = (skyline->count - nodes_to_remove + 1) * 2;
        pxl8_skyline_node* new_nodes = (pxl8_skyline_node*)realloc(
            skyline->nodes,
            new_capacity * sizeof(pxl8_skyline_node)
        );
        if (!new_nodes) return false;
        skyline->nodes = new_nodes;
        skyline->capacity = new_capacity;
    }

    if (nodes_to_remove > 0) {
        memmove(
            &skyline->nodes[node_idx + 1],
            &skyline->nodes[node_idx + nodes_to_remove],
            (skyline->count - node_idx - nodes_to_remove) * sizeof(pxl8_skyline_node)
        );
    }

    skyline->nodes[node_idx] = (pxl8_skyline_node){pos.x, pos.y + (i32)h, (i32)w};
    skyline->count = skyline->count - nodes_to_remove + 1;
    return true;
}

static void pxl8_skyline_compact(pxl8_skyline* skyline) {
    for (u32 i = 0; i < skyline->count - 1; ) {
        if (skyline->nodes[i].y == skyline->nodes[i + 1].y) {
            skyline->nodes[i].width += skyline->nodes[i + 1].width;
            memmove(
                &skyline->nodes[i + 1],
                &skyline->nodes[i + 2],
                (skyline->count - i - 2) * sizeof(pxl8_skyline_node)
            );
            skyline->count--;
        } else {
            i++;
        }
    }
}

pxl8_atlas* pxl8_atlas_create(u32 width, u32 height, pxl8_pixel_mode pixel_mode) {
    pxl8_atlas* atlas = (pxl8_atlas*)calloc(1, sizeof(pxl8_atlas));
    if (!atlas) return NULL;

    atlas->height = height;
    atlas->width = width;

    i32 bytes_per_pixel = pxl8_bytes_per_pixel(pixel_mode);
    atlas->pixels = (u8*)calloc(width * height, bytes_per_pixel);
    if (!atlas->pixels) {
        free(atlas);
        return NULL;
    }

    atlas->entry_capacity = PXL8_DEFAULT_ATLAS_ENTRY_CAPACITY;
    atlas->entries = (pxl8_atlas_entry*)calloc(atlas->entry_capacity, sizeof(pxl8_atlas_entry));
    if (!atlas->entries) {
        free(atlas->pixels);
        free(atlas);
        return NULL;
    }

    atlas->free_capacity = 16;
    atlas->free_list = (u32*)calloc(atlas->free_capacity, sizeof(u32));
    if (!atlas->free_list) {
        free(atlas->entries);
        free(atlas->pixels);
        free(atlas);
        return NULL;
    }

    atlas->skyline.capacity = 16;
    atlas->skyline.nodes =
        (pxl8_skyline_node*)calloc(atlas->skyline.capacity, sizeof(pxl8_skyline_node));
    if (!atlas->skyline.nodes) {
        free(atlas->free_list);
        free(atlas->entries);
        free(atlas->pixels);
        free(atlas);
        return NULL;
    }

    atlas->skyline.nodes[0] = (pxl8_skyline_node){0, 0, (i32)width};
    atlas->skyline.count = 1;

    return atlas;
}

void pxl8_atlas_destroy(pxl8_atlas* atlas) {
    if (!atlas) return;

    free(atlas->entries);
    free(atlas->free_list);
    free(atlas->pixels);
    free(atlas->skyline.nodes);
    free(atlas);
}

bool pxl8_atlas_expand(pxl8_atlas* atlas, pxl8_pixel_mode pixel_mode) {
    if (!atlas || atlas->width >= 4096) return false;

    i32 bytes_per_pixel = pxl8_bytes_per_pixel(pixel_mode);
    u32 new_size = atlas->width * 2;
    u32 old_width = atlas->width;

    u8* new_pixels = (u8*)calloc(new_size * new_size, bytes_per_pixel);
    if (!new_pixels) return false;

    pxl8_skyline new_skyline;
    new_skyline.nodes = (pxl8_skyline_node*)calloc(16, sizeof(pxl8_skyline_node));
    if (!new_skyline.nodes) {
        free(new_pixels);
        return false;
    }

    new_skyline.nodes[0] = (pxl8_skyline_node){0, 0, (i32)new_size};
    new_skyline.count = 1;
    new_skyline.capacity = 16;

    for (u32 i = 0; i < atlas->entry_count; i++) {
        if (!atlas->entries[i].active) continue;

        pxl8_skyline_fit fit = pxl8_skyline_find_position(
            &new_skyline,
            new_size,
            new_size,
            atlas->entries[i].w,
            atlas->entries[i].h
        );

        if (!fit.found) {
            free(new_skyline.nodes);
            free(new_pixels);
            return false;
        }

        for (u32 y = 0; y < (u32)atlas->entries[i].h; y++) {
            for (u32 x = 0; x < (u32)atlas->entries[i].w; x++) {
                u32 src_idx = (atlas->entries[i].y + y) * old_width + (atlas->entries[i].x + x);
                u32 dst_idx = (fit.pos.y + y) * new_size + (fit.pos.x + x);
                if (bytes_per_pixel == 2) {
                    ((u16*)new_pixels)[dst_idx] = ((u16*)atlas->pixels)[src_idx];
                } else {
                    new_pixels[dst_idx] = atlas->pixels[src_idx];
                }
            }
        }

        atlas->entries[i].x = fit.pos.x;
        atlas->entries[i].y = fit.pos.y;

        if (!pxl8_skyline_add_rect(&new_skyline, fit.pos, atlas->entries[i].w, atlas->entries[i].h)) {
            free(new_skyline.nodes);
            free(new_pixels);
            return false;
        }
        pxl8_skyline_compact(&new_skyline);
    }

    free(atlas->pixels);
    free(atlas->skyline.nodes);

    atlas->pixels = new_pixels;
    atlas->skyline = new_skyline;
    atlas->width = new_size;
    atlas->height = new_size;
    atlas->dirty = true;

    pxl8_debug("Atlas expanded to %ux%u", atlas->width, atlas->height);
    return true;
}

u32 pxl8_atlas_add_texture(
    pxl8_atlas* atlas,
    const u8* pixels,
    u32 w,
    u32 h,
    pxl8_pixel_mode pixel_mode
) {
    if (!atlas || !pixels) return UINT32_MAX;

    pxl8_skyline_fit fit =
        pxl8_skyline_find_position(&atlas->skyline, atlas->width, atlas->height, w, h);
    if (!fit.found) {
        if (!pxl8_atlas_expand(atlas, pixel_mode)) {
            return UINT32_MAX;
        }

        fit = pxl8_skyline_find_position(&atlas->skyline, atlas->width, atlas->height, w, h);

        if (!fit.found) return UINT32_MAX;
    }

    u32 texture_id;
    if (atlas->free_count > 0) {
        texture_id = atlas->free_list[--atlas->free_count];
    } else {
        if (atlas->entry_count >= atlas->entry_capacity) {
            u32 new_capacity = atlas->entry_capacity * 2;
            pxl8_atlas_entry* new_entries = (pxl8_atlas_entry*)realloc(
                atlas->entries,
                new_capacity * sizeof(pxl8_atlas_entry)
            );
            if (!new_entries) return UINT32_MAX;
            atlas->entries = new_entries;
            atlas->entry_capacity = new_capacity;
        }
        texture_id = atlas->entry_count++;
    }

    pxl8_atlas_entry* entry = &atlas->entries[texture_id];
    entry->active = true;
    entry->texture_id = texture_id;
    entry->x = fit.pos.x;
    entry->y = fit.pos.y;
    entry->w = w;
    entry->h = h;

    i32 bytes_per_pixel = pxl8_bytes_per_pixel(pixel_mode);
    for (u32 y = 0; y < h; y++) {
        for (u32 x = 0; x < w; x++) {
            u32 src_idx = y * w + x;
            u32 dst_idx = (fit.pos.y + y) * atlas->width + (fit.pos.x + x);

            if (bytes_per_pixel == 2) {
                ((u16*)atlas->pixels)[dst_idx] = ((const u16*)pixels)[src_idx];
            } else {
                atlas->pixels[dst_idx] = pixels[src_idx];
            }
        }
    }

    if (!pxl8_skyline_add_rect(&atlas->skyline, fit.pos, w, h)) {
        entry->active = false;
        return UINT32_MAX;
    }
    pxl8_skyline_compact(&atlas->skyline);

    atlas->dirty = true;

    return texture_id;
}

const pxl8_atlas_entry* pxl8_atlas_get_entry(const pxl8_atlas* atlas, u32 id) {
    if (!atlas || id >= atlas->entry_count) return NULL;
    return &atlas->entries[id];
}

u32 pxl8_atlas_get_entry_count(const pxl8_atlas* atlas) {
    return atlas ? atlas->entry_count : 0;
}

u32 pxl8_atlas_get_height(const pxl8_atlas* atlas) {
    return atlas ? atlas->height : 0;
}

const u8* pxl8_atlas_get_pixels(const pxl8_atlas* atlas) {
    return atlas ? atlas->pixels : NULL;
}

u32 pxl8_atlas_get_width(const pxl8_atlas* atlas) {
    return atlas ? atlas->width : 0;
}

bool pxl8_atlas_is_dirty(const pxl8_atlas* atlas) {
    return atlas ? atlas->dirty : false;
}

void pxl8_atlas_mark_clean(pxl8_atlas* atlas) {
    if (atlas) {
        atlas->dirty = false;
    }
}