#include "pxl8_palette.h"

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

#include "pxl8_ase.h"
#include "pxl8_color.h"
#include "pxl8_log.h"

#define PXL8_PALETTE_HASH_SIZE 512

typedef struct {
    u32 color;
    i16 index;
} pxl8_palette_hash_entry;

struct pxl8_palette {
    u32 base_colors[PXL8_MAX_CYCLES][PXL8_MAX_CYCLE_LEN];
    u32 colors[PXL8_PALETTE_SIZE];
    u8 color_ramp[PXL8_PALETTE_SIZE];
    u16 color_count;
    pxl8_cycle_range cycles[PXL8_MAX_CYCLES];
    i8 directions[PXL8_MAX_CYCLES];
    f32 phases[PXL8_MAX_CYCLES];
    pxl8_palette_hash_entry hash[PXL8_PALETTE_HASH_SIZE];
};

static inline u32 pxl8_palette_hash(u32 color) {
    color ^= color >> 16;
    color *= 0x85ebca6b;
    color ^= color >> 13;
    return color & (PXL8_PALETTE_HASH_SIZE - 1);
}

static void pxl8_palette_rebuild_hash(pxl8_palette* pal) {
    for (u32 i = 0; i < PXL8_PALETTE_HASH_SIZE; i++) {
        pal->hash[i].index = -1;
    }
    for (u32 i = 0; i < PXL8_PALETTE_SIZE; i++) {
        u32 color = pal->colors[i];
        u32 slot = pxl8_palette_hash(color);
        while (pal->hash[slot].index >= 0) {
            slot = (slot + 1) & (PXL8_PALETTE_HASH_SIZE - 1);
        }
        pal->hash[slot].color = color;
        pal->hash[slot].index = (i16)i;
    }
}

static inline u32 pack_rgb(u8 r, u8 g, u8 b) {
    return 0xFF000000 | ((u32)b << 16) | ((u32)g << 8) | r;
}

static inline u32 pack_rgba(u8 r, u8 g, u8 b, u8 a) {
    return ((u32)a << 24) | ((u32)b << 16) | ((u32)g << 8) | r;
}

static inline void unpack_rgba(u32 c, u8* r, u8* g, u8* b, u8* a) {
    *r = c & 0xFF;
    *g = (c >> 8) & 0xFF;
    *b = (c >> 16) & 0xFF;
    *a = (c >> 24) & 0xFF;
}

static f32 apply_easing(pxl8_easing easing, f32 t) {
    switch (easing) {
        case PXL8_EASE_IN:
            return t * t;
        case PXL8_EASE_IN_OUT:
            return t < 0.5f ? 2.0f * t * t : 1.0f - (-2.0f * t + 2.0f) * (-2.0f * t + 2.0f) * 0.5f;
        case PXL8_EASE_LINEAR:
            return t;
        case PXL8_EASE_OUT:
            return 1.0f - (1.0f - t) * (1.0f - t);
        default:
            return t;
    }
}

static void rgb_to_hsl(u8 r, u8 g, u8 b, f32* h, f32* s, f32* l) {
    f32 rf = r / 255.0f;
    f32 gf = g / 255.0f;
    f32 bf = b / 255.0f;

    f32 max = rf > gf ? (rf > bf ? rf : bf) : (gf > bf ? gf : bf);
    f32 min = rf < gf ? (rf < bf ? rf : bf) : (gf < bf ? gf : bf);
    f32 delta = max - min;

    *l = (max + min) / 2.0f;

    if (delta < 0.001f) {
        *h = 0.0f;
        *s = 0.0f;
    } else {
        *s = *l > 0.5f ? delta / (2.0f - max - min) : delta / (max + min);

        if (max == rf) {
            *h = (gf - bf) / delta + (gf < bf ? 6.0f : 0.0f);
        } else if (max == gf) {
            *h = (bf - rf) / delta + 2.0f;
        } else {
            *h = (rf - gf) / delta + 4.0f;
        }
        *h /= 6.0f;
    }
}

typedef struct {
    u8 index;
    f32 hue;
    f32 sat;
    f32 lum;
} palette_sort_entry;

static int palette_sort_cmp(const void* a, const void* b) {
    const palette_sort_entry* entry_a = (const palette_sort_entry*)a;
    const palette_sort_entry* entry_b = (const palette_sort_entry*)b;

    u8 a_gray = entry_a->sat < 0.1f;
    u8 b_gray = entry_b->sat < 0.1f;
    if (a_gray != b_gray) return a_gray - b_gray;

    if (a_gray) {
        if (entry_a->lum < entry_b->lum) return -1;
        if (entry_a->lum > entry_b->lum) return 1;
        return (int)entry_a->index - (int)entry_b->index;
    }

    if (entry_a->hue < entry_b->hue - 0.02f) return -1;
    if (entry_a->hue > entry_b->hue + 0.02f) return 1;

    if (entry_a->lum < entry_b->lum) return -1;
    if (entry_a->lum > entry_b->lum) return 1;

    return (int)entry_a->index - (int)entry_b->index;
}

static void pxl8_palette_sort_colors(pxl8_palette* pal) {
    if (!pal || pal->color_count < 2) return;

    pal->color_ramp[0] = 0;

    u8 count = pal->color_count;
    palette_sort_entry entries[PXL8_PALETTE_SIZE - 1];
    for (u32 i = 1; i < count; i++) {
        u8 r, g, b, a;
        unpack_rgba(pal->colors[i], &r, &g, &b, &a);

        entries[i - 1].index = (u8)i;
        rgb_to_hsl(r, g, b, &entries[i - 1].hue, &entries[i - 1].sat, &entries[i - 1].lum);
    }

    qsort(entries, count - 1, sizeof(palette_sort_entry), palette_sort_cmp);

    for (u32 i = 1; i < count; i++) {
        pal->color_ramp[i] = entries[i - 1].index;
    }
}

static u32 lerp_color(u32 c0, u32 c1, f32 t) {
    u8 r0, g0, b0, a0, r1, g1, b1, a1;
    unpack_rgba(c0, &r0, &g0, &b0, &a0);
    unpack_rgba(c1, &r1, &g1, &b1, &a1);

    u8 r = (u8)(r0 + (r1 - r0) * t);
    u8 g = (u8)(g0 + (g1 - g0) * t);
    u8 b = (u8)(b0 + (b1 - b0) * t);
    u8 a = (u8)(a0 + (a1 - a0) * t);

    return pack_rgba(r, g, b, a);
}

static void update_cycle_colors(pxl8_palette* pal, u8 slot) {
    pxl8_cycle_range* cycle = &pal->cycles[slot];
    f32 phase = pal->phases[slot];
    f32 eased = apply_easing(cycle->easing, phase);
    u8 len = cycle->len;
    u8 start = cycle->start;

    if (cycle->interpolate) {
        f32 pos = eased * (len - 1);
        u8 idx0 = (u8)pos;
        if (idx0 >= len - 1) idx0 = len - 1;
        f32 frac = pos - idx0;

        for (u8 i = 0; i < len; i++) {
            u8 src0 = (i + idx0) % len;
            u8 src1 = (i + idx0 + 1) % len;
            u32 c0 = pal->base_colors[slot][src0];
            u32 c1 = pal->base_colors[slot][src1];
            pal->colors[start + i] = lerp_color(c0, c1, frac);
        }
    } else {
        u8 offset = (u8)(eased * len) % len;
        for (u8 i = 0; i < len; i++) {
            u8 src = (i + offset) % len;
            pal->colors[start + i] = pal->base_colors[slot][src];
        }
    }
}

pxl8_palette* pxl8_palette_create(void) {
    pxl8_palette* pal = calloc(1, sizeof(pxl8_palette));
    if (!pal) return NULL;

    pal->colors[0] = 0x00000000;
    pal->colors[1] = pack_rgb(0x00, 0x00, 0x00);
    pal->colors[2] = pack_rgb(0xFF, 0xFF, 0xFF);
    pal->color_count = 3;
    pal->color_ramp[0] = 0;
    pal->color_ramp[1] = 1;
    pal->color_ramp[2] = 2;

    for (u8 i = 0; i < PXL8_MAX_CYCLES; i++) {
        pal->cycles[i] = pxl8_cycle_range_disabled();
        pal->phases[i] = 0.0f;
        pal->directions[i] = 1;
    }

    return pal;
}

void pxl8_palette_destroy(pxl8_palette* pal) {
    free(pal);
}

pxl8_result pxl8_palette_load_ase(pxl8_palette* pal, const char* path) {
    if (!pal || !path) return PXL8_ERROR_INVALID_ARGUMENT;

    pxl8_ase_file ase_file;
    pxl8_result result = pxl8_ase_load(path, &ase_file);
    if (result != PXL8_OK) {
        pxl8_error("Failed to load ASE file for palette: %s", path);
        return result;
    }

    if (ase_file.palette.entry_count == 0 || !ase_file.palette.colors) {
        pxl8_error("No palette data in ASE file");
        pxl8_ase_destroy(&ase_file);
        return PXL8_ERROR_INVALID_FORMAT;
    }

    u32 copy_count = ase_file.palette.entry_count < PXL8_PALETTE_SIZE
        ? ase_file.palette.entry_count
        : PXL8_PALETTE_SIZE;

    memcpy(pal->colors, ase_file.palette.colors, copy_count * sizeof(u32));
    pal->color_count = (u16)copy_count;

    pxl8_ase_destroy(&ase_file);
    pxl8_palette_sort_colors(pal);
    pxl8_palette_rebuild_hash(pal);
    return PXL8_OK;
}

u32* pxl8_palette_colors(pxl8_palette* pal) {
    return pal ? pal->colors : NULL;
}

u8* pxl8_palette_color_ramp(pxl8_palette* pal) {
    return pal ? pal->color_ramp : NULL;
}

u16 pxl8_palette_color_count(const pxl8_palette* pal) {
    return pal ? pal->color_count : 0;
}

u8 pxl8_palette_ramp_index(const pxl8_palette* pal, u8 position) {
    if (!pal || position >= pal->color_count) return 0;
    return pal->color_ramp[position];
}

u8 pxl8_palette_ramp_position(const pxl8_palette* pal, u8 index) {
    if (!pal) return 0;
    for (u32 i = 0; i < pal->color_count; i++) {
        if (pal->color_ramp[i] == index) return (u8)i;
    }
    return 0;
}

u8 pxl8_palette_find_closest(const pxl8_palette* pal, u8 r, u8 g, u8 b) {
    if (!pal || pal->color_count < 2) return 1;

    u8 best_idx = 1;
    u32 best_dist = 0xFFFFFFFF;

    for (u32 i = 1; i < pal->color_count; i++) {
        u8 pr, pg, pb, pa;
        unpack_rgba(pal->colors[i], &pr, &pg, &pb, &pa);

        i32 dr = (i32)r - (i32)pr;
        i32 dg = (i32)g - (i32)pg;
        i32 db = (i32)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;
}

u32 pxl8_palette_color(const pxl8_palette* pal, u8 idx) {
    if (!pal) return 0;
    return pxl8_color_to_rgba(pal->colors[idx]);
}

i32 pxl8_palette_index(const pxl8_palette* pal, u32 rgba) {
    if (!pal) return -1;
    if (rgba <= 0xFFFFFF) rgba = (rgba << 8) | 0xFF;
    u32 abgr = pxl8_color_from_rgba(rgba);
    u32 slot = pxl8_palette_hash(abgr);
    for (u32 i = 0; i < PXL8_PALETTE_HASH_SIZE; i++) {
        u32 idx = (slot + i) & (PXL8_PALETTE_HASH_SIZE - 1);
        if (pal->hash[idx].index < 0) return -1;
        if (pal->hash[idx].color == abgr) return pal->hash[idx].index;
    }
    return -1;
}

void pxl8_palette_get_rgb(const pxl8_palette* pal, u8 idx, u8* r, u8* g, u8* b) {
    if (!pal || !r || !g || !b) return;
    u8 a;
    unpack_rgba(pal->colors[idx], r, g, b, &a);
}

void pxl8_palette_get_rgba(const pxl8_palette* pal, u8 idx, u8* r, u8* g, u8* b, u8* a) {
    if (!pal || !r || !g || !b || !a) return;
    unpack_rgba(pal->colors[idx], r, g, b, a);
}

void pxl8_palette_set(pxl8_palette* pal, u8 idx, u32 color) {
    if (pal) pal->colors[idx] = color;
}

void pxl8_palette_set_rgb(pxl8_palette* pal, u8 idx, u8 r, u8 g, u8 b) {
    if (pal) pal->colors[idx] = pack_rgb(r, g, b);
}

void pxl8_palette_set_rgba(pxl8_palette* pal, u8 idx, u8 r, u8 g, u8 b, u8 a) {
    if (pal) pal->colors[idx] = pack_rgba(r, g, b, a);
}

void pxl8_palette_fill_gradient(pxl8_palette* pal, u8 start, u8 count, u32 from, u32 to) {
    if (!pal || count == 0) return;

    u8 r0, g0, b0, a0, r1, g1, b1, a1;
    unpack_rgba(from, &r0, &g0, &b0, &a0);
    unpack_rgba(to, &r1, &g1, &b1, &a1);

    for (u8 i = 0; i < count; i++) {
        f32 t = (count > 1) ? (f32)i / (f32)(count - 1) : 0.0f;
        u8 r = (u8)(r0 + (r1 - r0) * t);
        u8 g = (u8)(g0 + (g1 - g0) * t);
        u8 b = (u8)(b0 + (b1 - b0) * t);
        u8 a = (u8)(a0 + (a1 - a0) * t);
        pal->colors[start + i] = pack_rgba(r, g, b, a);
    }
}

void pxl8_palette_fill_gradient_rgb(pxl8_palette* pal, u8 start, u8 count, u8 r0, u8 g0, u8 b0, u8 r1, u8 g1, u8 b1) {
    pxl8_palette_fill_gradient(pal, start, count, pack_rgb(r0, g0, b0), pack_rgb(r1, g1, b1));
}

void pxl8_palette_reset_cycle(pxl8_palette* pal, u8 slot) {
    if (!pal || slot >= PXL8_MAX_CYCLES) return;
    pal->phases[slot] = 0.0f;
    pal->directions[slot] = 1;
}

void pxl8_palette_set_cycle(pxl8_palette* pal, u8 slot, pxl8_cycle_range range) {
    if (!pal || slot >= PXL8_MAX_CYCLES) return;

    pal->cycles[slot] = range;
    pal->phases[slot] = 0.0f;
    pal->directions[slot] = 1;

    u8 start = range.start;
    u8 len = range.len;
    if (len > PXL8_MAX_CYCLE_LEN) len = PXL8_MAX_CYCLE_LEN;

    for (u8 i = 0; i < len; i++) {
        pal->base_colors[slot][i] = pal->colors[start + i];
    }
}

void pxl8_palette_set_cycle_colors(pxl8_palette* pal, u8 slot, const u32* colors, u8 count) {
    if (!pal || !colors || slot >= PXL8_MAX_CYCLES) return;

    pxl8_cycle_range* cycle = &pal->cycles[slot];
    u8 start = cycle->start;
    u8 len = cycle->len;
    if (len > count) len = count;
    if (len > PXL8_MAX_CYCLE_LEN) len = PXL8_MAX_CYCLE_LEN;

    for (u8 i = 0; i < len; i++) {
        pal->base_colors[slot][i] = colors[i];
        pal->colors[start + i] = colors[i];
    }
}

void pxl8_palette_set_cycle_phase(pxl8_palette* pal, u8 slot, f32 phase) {
    if (!pal || slot >= PXL8_MAX_CYCLES) return;

    if (phase < 0.0f) phase = 0.0f;
    if (phase > 1.0f) phase = 1.0f;

    pal->phases[slot] = phase;
    update_cycle_colors(pal, slot);
}

void pxl8_palette_tick(pxl8_palette* pal, u16 delta_ticks) {
    if (!pal) return;

    for (u8 slot = 0; slot < PXL8_MAX_CYCLES; slot++) {
        pxl8_cycle_range* cycle = &pal->cycles[slot];
        if (cycle->period == 0 || cycle->len < 2) continue;

        f32 delta = (f32)delta_ticks / (f32)cycle->period;
        f32 dir = (f32)pal->directions[slot];
        pal->phases[slot] += delta * dir;

        switch (cycle->mode) {
            case PXL8_CYCLE_LOOP:
                while (pal->phases[slot] >= 1.0f) pal->phases[slot] -= 1.0f;
                while (pal->phases[slot] < 0.0f) pal->phases[slot] += 1.0f;
                break;

            case PXL8_CYCLE_PINGPONG:
                if (pal->phases[slot] >= 1.0f) {
                    pal->phases[slot] = 1.0f - (pal->phases[slot] - 1.0f);
                    pal->directions[slot] = -1;
                } else if (pal->phases[slot] <= 0.0f) {
                    pal->phases[slot] = -pal->phases[slot];
                    pal->directions[slot] = 1;
                }
                break;

            case PXL8_CYCLE_ONCE:
                if (pal->phases[slot] < 0.0f) pal->phases[slot] = 0.0f;
                if (pal->phases[slot] > 1.0f) pal->phases[slot] = 1.0f;
                break;
        }

        update_cycle_colors(pal, slot);
    }
}

pxl8_cycle_range pxl8_cycle_range_new(u8 start, u8 len, u16 period) {
    pxl8_cycle_range range = {
        .easing = PXL8_EASE_LINEAR,
        .interpolate = true,
        .len = len,
        .mode = PXL8_CYCLE_LOOP,
        .period = period,
        .start = start,
    };
    return range;
}

pxl8_cycle_range pxl8_cycle_range_disabled(void) {
    pxl8_cycle_range range = {
        .easing = PXL8_EASE_LINEAR,
        .interpolate = false,
        .len = 0,
        .mode = PXL8_CYCLE_LOOP,
        .period = 0,
        .start = 0,
    };
    return range;
}