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

#define MINIZ_NO_STDIO
#define MINIZ_NO_TIME
#define MINIZ_NO_ARCHIVE_APIS
#define MINIZ_NO_ARCHIVE_WRITING_APIS
#define MINIZ_NO_DEFLATE_APIS

#include <miniz.h>
#include <SDL3/SDL.h>

#include "pxl8_ase.h"
#include "pxl8_io.h"
#include "pxl8_macros.h"

static u16 read_u16_le(const u8* data) {
    return data[0] | (data[1] << 8);
}

static u32 read_u32_le(const u8* data) {
    return data[0] | (data[1] << 8) | (data[2] << 16) | (data[3] << 24);
}

static i16 read_i16_le(const u8* data) {
    return (i16)read_u16_le(data);
}

static pxl8_result parse_ase_header(const u8* data, pxl8_ase_header* header) {
    header->file_size = read_u32_le(data);
    header->magic = read_u16_le(data + 4);
    header->frames = read_u16_le(data + 6);
    header->width = read_u16_le(data + 8);
    header->height = read_u16_le(data + 10);
    header->color_depth = read_u16_le(data + 12);
    header->flags = read_u32_le(data + 14);
    header->speed = read_u16_le(data + 18);
    header->transparent_index = read_u32_le(data + 24);
    header->n_colors = data[28];
    header->pixel_width = data[29];
    header->pixel_height = data[30];
    header->grid_x = read_i16_le(data + 31);
    header->grid_y = read_i16_le(data + 33);
    header->grid_width = read_u16_le(data + 35);
    header->grid_height = read_u16_le(data + 37);

    if (header->magic != PXL8_ASE_MAGIC) {
        pxl8_error("Invalid ASE file magic: 0x%04X", header->magic);
        return PXL8_ERROR_ASE_INVALID_MAGIC;
    }

    return PXL8_OK;
}

static pxl8_result parse_old_palette_chunk(const u8* data, pxl8_ase_palette* palette) {
    u16 packet_count = read_u16_le(data);
    const u8* packet_data = data + 2;
    
    u32 total_colors = 0;
    const u8* temp_data = packet_data;
    for (u16 packet = 0; packet < packet_count; packet++) {
        u8 skip_colors = temp_data[0];
        u8 colors_in_packet = temp_data[1];
        u32 actual_colors = (colors_in_packet == 0) ? 256 : colors_in_packet;
        total_colors = skip_colors + actual_colors;
        temp_data += 2 + (actual_colors * 3);
    }
    
    palette->entry_count = total_colors;
    palette->first_color = 0;
    palette->last_color = total_colors - 1;
    palette->colors = (u32*)SDL_malloc(total_colors * sizeof(u32));
    if (!palette->colors) {
        return PXL8_ERROR_OUT_OF_MEMORY;
    }
    
    u32 color_index = 0;
    for (u16 packet = 0; packet < packet_count; packet++) {
        u8 skip_colors = packet_data[0];
        u8 colors_in_packet = packet_data[1];
        u32 actual_colors = (colors_in_packet == 0) ? 256 : colors_in_packet;
        
        for (u32 i = 0; i < skip_colors && color_index < total_colors; i++, color_index++) {
            palette->colors[color_index] = 0xFF000000;
        }
        
        packet_data += 2;
        
        for (u32 i = 0; i < actual_colors && color_index < total_colors; i++, color_index++) {
            u8 r = packet_data[0];
            u8 g = packet_data[1];
            u8 b = packet_data[2];

            palette->colors[color_index] = 0xFF000000 | (b << 16) | (g << 8) | r;
            packet_data += 3;
        }
    }
    
    return PXL8_OK;
}

static pxl8_result parse_layer_chunk(const u8* data, pxl8_ase_layer* layer) {
    layer->flags = read_u16_le(data);
    layer->layer_type = read_u16_le(data + 2);
    layer->child_level = read_u16_le(data + 4);
    layer->blend_mode = read_u16_le(data + 10);
    layer->opacity = data[12];

    u16 name_len = read_u16_le(data + 16);
    if (name_len > 0) {
        layer->name = (char*)SDL_malloc(name_len + 1);
        if (!layer->name) return PXL8_ERROR_OUT_OF_MEMORY;
        memcpy(layer->name, data + 18, name_len);
        layer->name[name_len] = '\0';
    } else {
        layer->name = NULL;
    }
    
    return PXL8_OK;
}

static pxl8_result parse_palette_chunk(const u8* data, pxl8_ase_palette* palette) {
    palette->entry_count = read_u32_le(data);
    palette->first_color = read_u32_le(data + 4);
    palette->last_color = read_u32_le(data + 8);
    
    u32 color_count = palette->entry_count;
    palette->colors = (u32*)SDL_malloc(color_count * sizeof(u32));
    if (!palette->colors) {
        return PXL8_ERROR_OUT_OF_MEMORY;
    }

    const u8* color_data = data + 20;
    for (u32 i = 0; i < color_count; i++) {
        u16 flags = read_u16_le(color_data);
        u8 r = color_data[2];
        u8 g = color_data[3];
        u8 b = color_data[4];
        u8 a = color_data[5];

        palette->colors[i] = (a << 24) | (b << 16) | (g << 8) | r;
        color_data += 6;
        
        if (flags & 1) {
            u16 name_len = read_u16_le(color_data);
            color_data += 2 + name_len;
        }
    }

    return PXL8_OK;
}

static pxl8_result parse_cel_chunk(const u8* data, u32 chunk_size, pxl8_ase_cel* cel) {
    if (chunk_size < 9) {
        return PXL8_ERROR_ASE_MALFORMED_CHUNK;
    }

    cel->layer_index = read_u16_le(data);
    cel->x = read_i16_le(data + 2);
    cel->y = read_i16_le(data + 4);
    cel->opacity = data[6];
    cel->cel_type = read_u16_le(data + 7);

    if (cel->cel_type == 2) {
        if (chunk_size < 20) {
            return PXL8_ERROR_ASE_MALFORMED_CHUNK;
        }

        cel->width = read_u16_le(data + 16);
        cel->height = read_u16_le(data + 18);
        
        u32 pixel_data_size = cel->width * cel->height;
        u32 compressed_data_size = chunk_size - 20;
        
        cel->pixel_data = (u8*)SDL_malloc(pixel_data_size);
        if (!cel->pixel_data) {
            return PXL8_ERROR_OUT_OF_MEMORY;
        }

        mz_ulong dest_len = pixel_data_size;
        int result = mz_uncompress(cel->pixel_data, &dest_len, data + 20, compressed_data_size);
        if (result != MZ_OK) {
            pxl8_error("Failed to decompress cel data: miniz error %d", result);
            SDL_free(cel->pixel_data);
            cel->pixel_data = NULL;
            return PXL8_ERROR_ASE_MALFORMED_CHUNK;
        }
        
        if (dest_len != pixel_data_size) {
            pxl8_warn("Decompressed size mismatch: expected %u, got %lu", pixel_data_size, dest_len);
        }
    }

    return PXL8_OK;
}

pxl8_result pxl8_ase_load(const char* filepath, pxl8_ase_file* ase_file) {
    if (!filepath || !ase_file) {
        return PXL8_ERROR_NULL_POINTER;
    }

    memset(ase_file, 0, sizeof(pxl8_ase_file));

    u8* file_data;
    size_t file_size;
    pxl8_result result = pxl8_io_read_binary_file(filepath, &file_data, &file_size);
    if (result != PXL8_OK) {
        return result;
    }

    if (file_size < 128) {
        pxl8_io_free_binary_data(file_data);
        return PXL8_ERROR_ASE_TRUNCATED_FILE;
    }

    result = parse_ase_header(file_data, &ase_file->header);
    if (result != PXL8_OK) {
        pxl8_io_free_binary_data(file_data);
        return result;
    }

    ase_file->frame_count = ase_file->header.frames;
    ase_file->frames = (pxl8_ase_frame*)SDL_calloc(ase_file->frame_count, sizeof(pxl8_ase_frame));
    if (!ase_file->frames) {
        pxl8_io_free_binary_data(file_data);
        return PXL8_ERROR_OUT_OF_MEMORY;
    }

    const u8* frame_data = file_data + 128;
    for (u16 frame_idx = 0; frame_idx < ase_file->header.frames; frame_idx++) {
        pxl8_ase_frame_header frame_header;
        frame_header.frame_bytes = read_u32_le(frame_data);
        frame_header.magic = read_u16_le(frame_data + 4);
        frame_header.chunks = read_u16_le(frame_data + 6);
        frame_header.duration = read_u16_le(frame_data + 8);

        if (frame_header.magic != PXL8_ASE_FRAME_MAGIC) {
            pxl8_error("Invalid frame magic: 0x%04X", frame_header.magic);
            result = PXL8_ERROR_ASE_INVALID_FRAME_MAGIC;
            break;
        }

        pxl8_ase_frame* frame = &ase_file->frames[frame_idx];
        frame->frame_id = frame_idx;
        frame->width = ase_file->header.width;
        frame->height = ase_file->header.height;
        frame->duration = frame_header.duration;
        
        u32 pixel_count = frame->width * frame->height;
        frame->pixels = (u8*)SDL_calloc(pixel_count, sizeof(u8));
        if (!frame->pixels) {
            result = PXL8_ERROR_OUT_OF_MEMORY;
            break;
        }

        const u8* chunk_data = frame_data + 16;
        for (u16 chunk_idx = 0; chunk_idx < frame_header.chunks; chunk_idx++) {
            pxl8_ase_chunk_header chunk_header;
            chunk_header.chunk_size = read_u32_le(chunk_data);
            chunk_header.chunk_type = read_u16_le(chunk_data + 4);

            const u8* chunk_payload = chunk_data + 6;
            
            pxl8_debug("Found chunk: type=0x%04X, size=%d", chunk_header.chunk_type, chunk_header.chunk_size);
            switch (chunk_header.chunk_type) {
                case PXL8_ASE_CHUNK_OLD_PALETTE:  // 0x0004
                    if (!ase_file->palette.colors) {
                        result = parse_old_palette_chunk(chunk_payload, &ase_file->palette);
                        pxl8_debug("Parsed old palette: %d colors, indices %d-%d",
                                    ase_file->palette.entry_count, ase_file->palette.first_color, ase_file->palette.last_color);
                    } else {
                        pxl8_debug("Ignoring old palette (0x0004) - new palette (0x2019) already loaded");
                    }
                    break;
                    
                case PXL8_ASE_CHUNK_LAYER: {  // 0x2004
                    ase_file->layers =
                        (pxl8_ase_layer*)SDL_realloc(ase_file->layers,
                            (ase_file->layer_count + 1) * sizeof(pxl8_ase_layer));
                    if (!ase_file->layers) {
                        result = PXL8_ERROR_OUT_OF_MEMORY;
                        break;
                    }
                    
                    result = parse_layer_chunk(chunk_payload, &ase_file->layers[ase_file->layer_count]);
                    if (result == PXL8_OK) {
                        pxl8_debug("Parsed layer %d: '%s', blend_mode=%d, opacity=%d", 
                                   ase_file->layer_count,
                                   ase_file->layers[ase_file->layer_count].name ? ase_file->layers[ase_file->layer_count].name : "(unnamed)",
                                   ase_file->layers[ase_file->layer_count].blend_mode,
                                   ase_file->layers[ase_file->layer_count].opacity);
                        ase_file->layer_count++;
                    }
                    break;
                }
                    
                case PXL8_ASE_CHUNK_CEL: {  // 0x2005
                    pxl8_ase_cel cel = {0};
                    pxl8_debug("Found CEL chunk: size=%d", chunk_header.chunk_size - 6);
                    result = parse_cel_chunk(chunk_payload, chunk_header.chunk_size - 6, &cel);
                    if (result == PXL8_OK && cel.pixel_data) {
                        pxl8_debug("CEL data loaded: %dx%d, type=%d, first pixel = %d", cel.width, cel.height, cel.cel_type, cel.pixel_data[0]);
                        u32 copy_width = (cel.width < frame->width) ? cel.width : frame->width;
                        u32 copy_height = (cel.height < frame->height) ? cel.height : frame->height;
                        
                        pxl8_debug("Copying cel to frame: cel_pos=(%d,%d), copy_size=%dx%d", cel.x, cel.y, copy_width, copy_height);
                        for (u32 y = 0; y < copy_height; y++) {
                            u32 src_offset = y * cel.width;
                            u32 dst_offset = (y + cel.y) * frame->width + cel.x;
                            if (dst_offset + copy_width <= pixel_count) {
                                for (u32 x = 0; x < copy_width; x++) {
                                    u8 src_pixel = cel.pixel_data[src_offset + x];
                                    bool is_transparent = false;
                                    
                                    if (src_pixel < ase_file->palette.entry_count && ase_file->palette.colors) {
                                        u32 color = ase_file->palette.colors[src_pixel];
                                        is_transparent = ((color >> 24) & 0xFF) == 0;
                                    }
                                    
                                    if (!is_transparent) {
                                        frame->pixels[dst_offset + x] = src_pixel;
                                    }
                                }
                                if (y < 3) {
                                    pxl8_debug("Row %d: cel[%d]=%d, frame[%d]=%d",
                                                y, src_offset, cel.pixel_data[src_offset], dst_offset, frame->pixels[dst_offset]);
                                }
                            }
                        }
                        SDL_free(cel.pixel_data);
                    }
                    break;
                }
                    
                case PXL8_ASE_CHUNK_PALETTE:  // 0x2019
                    if (ase_file->palette.colors) {
                        SDL_free(ase_file->palette.colors);
                    }
                    result = parse_palette_chunk(chunk_payload, &ase_file->palette);
                    pxl8_debug("Parsed new palette: %d colors, indices %d-%d",
                                ase_file->palette.entry_count, ase_file->palette.first_color, ase_file->palette.last_color);
                    break;
                    
                default:
                    break;
            }
            
            if (result != PXL8_OK) break;

            chunk_data += chunk_header.chunk_size;
        }
        
        if (result != PXL8_OK) break;

        frame_data += frame_header.frame_bytes;
    }

    pxl8_io_free_binary_data(file_data);

    if (result != PXL8_OK) {
        pxl8_ase_free(ase_file);
    }

    return result;
}

void pxl8_ase_free(pxl8_ase_file* ase_file) {
    if (!ase_file) return;

    if (ase_file->frames) {
        for (u32 i = 0; i < ase_file->frame_count; i++) {
            if (ase_file->frames[i].pixels) {
                SDL_free(ase_file->frames[i].pixels);
            }
        }
        SDL_free(ase_file->frames);
    }

    if (ase_file->palette.colors) {
        SDL_free(ase_file->palette.colors);
    }

    if (ase_file->layers) {
        for (u32 i = 0; i < ase_file->layer_count; i++) {
            if (ase_file->layers[i].name) {
                SDL_free(ase_file->layers[i].name);
            }
        }
        SDL_free(ase_file->layers);
    }

    memset(ase_file, 0, sizeof(pxl8_ase_file));
}