#include "pxl8_colormap.h"
#include <string.h>

static u8 find_closest_color(const u32* palette, u8 target_r, u8 target_g, u8 target_b) {
    u8 best_idx = 1;
    u32 best_dist = 0xFFFFFFFF;

    u8 dynamic_end = PXL8_DYNAMIC_RANGE_START + PXL8_DYNAMIC_RANGE_COUNT;

    for (u32 i = 1; i < PXL8_FULLBRIGHT_START; i++) {
        if (i >= PXL8_DYNAMIC_RANGE_START && i < dynamic_end) {
            continue;
        }

        u32 c = palette[i];
        u8 pr = c & 0xFF;
        u8 pg = (c >> 8) & 0xFF;
        u8 pb = (c >> 16) & 0xFF;

        i32 dr = (i32)target_r - (i32)pr;
        i32 dg = (i32)target_g - (i32)pg;
        i32 db = (i32)target_b - (i32)pb;
        u32 dist = (u32)(dr * dr + dg * dg + db * db);

        if (dist < best_dist) {
            best_dist = dist;
            best_idx = (u8)i;
            if (dist == 0) break;
        }
    }

    return best_idx;
}

void pxl8_set_colormap(pxl8_colormap* cm, const u8* data, u32 size) {
    if (!cm || !data || size == 0) return;
    u32 copy_size = size > PXL8_COLORMAP_SIZE ? PXL8_COLORMAP_SIZE : size;
    memcpy(cm->table, data, copy_size);
}

static void generate_light_table(pxl8_colormap* cm, const u32* palette, pxl8_light_color light_color) {
    pxl8_rgb light = pxl8_light_colors[light_color];
    u32 base_row = (u32)light_color * PXL8_LIGHT_LEVELS;

    for (u32 level = 0; level < PXL8_LIGHT_LEVELS; level++) {
        f32 brightness = (f32)level / (f32)(PXL8_LIGHT_LEVELS - 1);
        u32 row = base_row + level;

        for (u32 pal_idx = 0; pal_idx < 256; pal_idx++) {
            u8 result_idx;

            if (pal_idx == PXL8_TRANSPARENT) {
                result_idx = PXL8_TRANSPARENT;
            } else if (pal_idx >= PXL8_FULLBRIGHT_START) {
                result_idx = (u8)pal_idx;
            } else {
                u32 c = palette[pal_idx];
                u8 r = c & 0xFF;
                u8 g = (c >> 8) & 0xFF;
                u8 b = (c >> 16) & 0xFF;

                f32 lr = (f32)light.r / 255.0f;
                f32 lg = (f32)light.g / 255.0f;
                f32 lb = (f32)light.b / 255.0f;

                u8 target_r = (u8)(r * brightness * lr);
                u8 target_g = (u8)(g * brightness * lg);
                u8 target_b = (u8)(b * brightness * lb);

                result_idx = find_closest_color(palette, target_r, target_g, target_b);
            }

            cm->table[row * 256 + pal_idx] = result_idx;
        }
    }
}

static void generate_blend_table(pxl8_colormap* cm, const u32* palette) {
    for (u32 src = 0; src < 256; src++) {
        u32 row = PXL8_LIGHT_ROWS + src;

        u8 sr, sg, sb;
        if (src == PXL8_TRANSPARENT) {
            sr = sg = sb = 0;
        } else {
            u32 sc = palette[src];
            sr = sc & 0xFF;
            sg = (sc >> 8) & 0xFF;
            sb = (sc >> 16) & 0xFF;
        }

        for (u32 dst = 0; dst < 256; dst++) {
            u8 result_idx;

            if (src == PXL8_TRANSPARENT) {
                result_idx = (u8)dst;
            } else {
                u32 dc = palette[dst];
                u8 dr = dc & 0xFF;
                u8 dg = (dc >> 8) & 0xFF;
                u8 db = (dc >> 16) & 0xFF;

                u8 blend_r = (u8)((sr + dr) / 2);
                u8 blend_g = (u8)((sg + dg) / 2);
                u8 blend_b = (u8)((sb + db) / 2);

                result_idx = find_closest_color(palette, blend_r, blend_g, blend_b);
            }

            cm->table[row * 256 + dst] = result_idx;
        }
    }
}

void pxl8_colormap_generate(pxl8_colormap* cm, const u32* palette) {
    if (!cm || !palette) return;

    for (u32 light = 0; light < PXL8_LIGHT_COLORS; light++) {
        generate_light_table(cm, palette, (pxl8_light_color)light);
    }

    generate_blend_table(cm, palette);
}