#include "pxl8_tilesheet.h"

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

#include "pxl8_ase.h"
#include "pxl8_macros.h"
#include "pxl8_tilemap.h"

struct pxl8_tilesheet {
    u8* data;
    bool* tile_valid;

    u32 height;
    u32 tile_size;
    u32 tiles_per_row;
    u32 total_tiles;
    u32 width;

    pxl8_color_mode color_mode;
    u32 ref_count;

    pxl8_tile_animation* animations;
    u32 animation_count;

    pxl8_tile_properties* properties;

    pxl8_autotile_rule** autotile_rules;
    u32* autotile_rule_counts;
};

pxl8_tilesheet* pxl8_tilesheet_create(u32 tile_size) {
    pxl8_tilesheet* tilesheet = calloc(1, sizeof(pxl8_tilesheet));
    if (!tilesheet) return NULL;

    tilesheet->tile_size = tile_size;
    tilesheet->ref_count = 1;

    return tilesheet;
}

void pxl8_tilesheet_destroy(pxl8_tilesheet* tilesheet) {
    if (!tilesheet) return;

    if (tilesheet->data) {
        free(tilesheet->data);
    }

    if (tilesheet->tile_valid) {
        free(tilesheet->tile_valid);
    }

    if (tilesheet->animations) {
        for (u32 i = 0; i < tilesheet->animation_count; i++) {
            if (tilesheet->animations[i].frames) {
                free(tilesheet->animations[i].frames);
            }
        }
        free(tilesheet->animations);
    }

    if (tilesheet->properties) {
        free(tilesheet->properties);
    }

    if (tilesheet->autotile_rules) {
        for (u32 i = 0; i <= tilesheet->total_tiles; i++) {
            if (tilesheet->autotile_rules[i]) {
                free(tilesheet->autotile_rules[i]);
            }
        }
        free(tilesheet->autotile_rules);
        free(tilesheet->autotile_rule_counts);
    }

    free(tilesheet);
}

pxl8_result pxl8_tilesheet_load(pxl8_tilesheet* tilesheet, const char* filepath, pxl8_gfx* gfx) {
    if (!tilesheet || !filepath || !gfx) return PXL8_ERROR_NULL_POINTER;

    pxl8_ase_file ase_file;
    pxl8_result result = pxl8_ase_load(filepath, &ase_file);
    if (result != PXL8_OK) {
        pxl8_error("Failed to load tilesheet: %s", filepath);
        return result;
    }

    if (tilesheet->data) {
        free(tilesheet->data);
    }

    u32 width = ase_file.header.width;
    u32 height = ase_file.header.height;

    if (ase_file.header.grid_width > 0 && ase_file.header.grid_height > 0) {
        tilesheet->tile_size = ase_file.header.grid_width;
        pxl8_info("Using Aseprite grid size: %dx%d",
                  ase_file.header.grid_width, ase_file.header.grid_height);
    }

    tilesheet->width = width;
    tilesheet->height = height;
    tilesheet->tiles_per_row = width / tilesheet->tile_size;
    tilesheet->total_tiles = (width / tilesheet->tile_size) * (height / tilesheet->tile_size);
    tilesheet->color_mode = pxl8_gfx_get_color_mode(gfx);

    size_t data_size = width * height;
    if (pxl8_gfx_get_color_mode(gfx) == PXL8_COLOR_MODE_HICOLOR) {
        data_size *= 4;
    }

    tilesheet->data = malloc(data_size);
    if (!tilesheet->data) {
        pxl8_ase_destroy(&ase_file);
        return PXL8_ERROR_OUT_OF_MEMORY;
    }

    if (ase_file.frame_count > 0 && ase_file.frames[0].pixels) {
        memcpy(tilesheet->data, ase_file.frames[0].pixels, data_size);
    }

    tilesheet->tile_valid = calloc(tilesheet->total_tiles + 1, sizeof(bool));
    if (!tilesheet->tile_valid) {
        free(tilesheet->data);
        tilesheet->data = NULL;
        pxl8_ase_destroy(&ase_file);
        return PXL8_ERROR_OUT_OF_MEMORY;
    }

    u32 valid_tiles = 0;
    bool is_hicolor = (tilesheet->color_mode == PXL8_COLOR_MODE_HICOLOR);

    for (u32 tile_id = 1; tile_id <= tilesheet->total_tiles; tile_id++) {
        u32 tile_x = ((tile_id - 1) % tilesheet->tiles_per_row) * tilesheet->tile_size;
        u32 tile_y = ((tile_id - 1) / tilesheet->tiles_per_row) * tilesheet->tile_size;

        bool has_content = false;
        for (u32 py = 0; py < tilesheet->tile_size && !has_content; py++) {
            for (u32 px = 0; px < tilesheet->tile_size; px++) {
                if (is_hicolor) {
                    u32 idx = ((tile_y + py) * width + (tile_x + px)) * 4;
                    u8 alpha = tilesheet->data[idx + 3];
                    if (alpha > 0) {
                        has_content = true;
                        break;
                    }
                } else {
                    u32 idx = (tile_y + py) * width + (tile_x + px);
                    if (tilesheet->data[idx] != 0) {
                        has_content = true;
                        break;
                    }
                }
            }
        }
        if (has_content) {
            tilesheet->tile_valid[tile_id] = true;
            valid_tiles++;
        }
    }

    pxl8_info("Loaded tilesheet %s: %dx%d, %d valid tiles out of %d slots (%dx%d each)",
              filepath, width, height, valid_tiles, tilesheet->total_tiles,
              tilesheet->tile_size, tilesheet->tile_size);

    pxl8_ase_destroy(&ase_file);
    return PXL8_OK;
}

void pxl8_tilesheet_render_tile(const pxl8_tilesheet* tilesheet, pxl8_gfx* gfx,
                                u16 tile_id, i32 x, i32 y, u8 flags) {
    if (!tilesheet || !gfx || !tilesheet->data) return;
    if (tile_id == 0 || tile_id > tilesheet->total_tiles) return;
    if (tilesheet->tile_valid && !tilesheet->tile_valid[tile_id]) return;

    u32 tile_x = ((tile_id - 1) % tilesheet->tiles_per_row) * tilesheet->tile_size;
    u32 tile_y = ((tile_id - 1) / tilesheet->tiles_per_row) * tilesheet->tile_size;

    for (u32 py = 0; py < tilesheet->tile_size; py++) {
        for (u32 px = 0; px < tilesheet->tile_size; px++) {
            u32 src_x = tile_x + px;
            u32 src_y = tile_y + py;

            if (flags & PXL8_TILE_FLIP_X) src_x = tile_x + (tilesheet->tile_size - 1 - px);
            if (flags & PXL8_TILE_FLIP_Y) src_y = tile_y + (tilesheet->tile_size - 1 - py);

            u32 src_idx = src_y * tilesheet->width + src_x;
            u8 color_idx = tilesheet->data[src_idx];

            if (color_idx != 0) {
                i32 screen_x = x + px;
                i32 screen_y = y + py;

                if (screen_x >= 0 && screen_x < pxl8_gfx_get_width(gfx) &&
                    screen_y >= 0 && screen_y < pxl8_gfx_get_height(gfx)) {
                    pxl8_pixel(gfx, screen_x, screen_y, color_idx);
                }
            }
        }
    }
}

bool pxl8_tilesheet_is_tile_valid(const pxl8_tilesheet* tilesheet, u16 tile_id) {
    if (!tilesheet || tile_id == 0 || tile_id > tilesheet->total_tiles) return false;
    return tilesheet->tile_valid ? tilesheet->tile_valid[tile_id] : true;
}

void pxl8_tilesheet_ref(pxl8_tilesheet* tilesheet) {
    if (!tilesheet) return;
    tilesheet->ref_count++;
}

void pxl8_tilesheet_unref(pxl8_tilesheet* tilesheet) {
    if (!tilesheet) return;

    if (--tilesheet->ref_count == 0) {
        pxl8_tilesheet_destroy(tilesheet);
    }
}

pxl8_result pxl8_tilesheet_add_animation(pxl8_tilesheet* tilesheet, u16 base_tile_id,
                                         const u16* frames, u16 frame_count, f32 frame_duration) {
    if (!tilesheet || !frames || frame_count == 0) return PXL8_ERROR_INVALID_ARGUMENT;
    if (base_tile_id == 0 || base_tile_id > tilesheet->total_tiles) return PXL8_ERROR_INVALID_ARGUMENT;

    if (!tilesheet->animations) {
        tilesheet->animations = calloc(tilesheet->total_tiles + 1, sizeof(pxl8_tile_animation));
        if (!tilesheet->animations) return PXL8_ERROR_OUT_OF_MEMORY;
    }

    pxl8_tile_animation* anim = &tilesheet->animations[base_tile_id];
    if (anim->frames) free(anim->frames);

    anim->frames = malloc(frame_count * sizeof(u16));
    if (!anim->frames) return PXL8_ERROR_OUT_OF_MEMORY;

    memcpy(anim->frames, frames, frame_count * sizeof(u16));
    anim->frame_count = frame_count;
    anim->current_frame = 0;
    anim->frame_duration = frame_duration;
    anim->time_accumulator = 0;

    tilesheet->animation_count++;
    return PXL8_OK;
}

void pxl8_tilesheet_update_animations(pxl8_tilesheet* tilesheet, f32 delta_time) {
    if (!tilesheet || !tilesheet->animations) return;

    for (u32 i = 1; i <= tilesheet->total_tiles; i++) {
        pxl8_tile_animation* anim = &tilesheet->animations[i];
        if (!anim->frames || anim->frame_count == 0) continue;

        anim->time_accumulator += delta_time;
        while (anim->time_accumulator >= anim->frame_duration) {
            anim->time_accumulator -= anim->frame_duration;
            anim->current_frame = (anim->current_frame + 1) % anim->frame_count;
        }
    }
}

u16 pxl8_tilesheet_get_animated_frame(const pxl8_tilesheet* tilesheet, u16 tile_id) {
    if (!tilesheet || !tilesheet->animations || tile_id == 0 || tile_id > tilesheet->total_tiles) {
        return tile_id;
    }

    const pxl8_tile_animation* anim = &tilesheet->animations[tile_id];
    if (!anim->frames || anim->frame_count == 0) return tile_id;

    return anim->frames[anim->current_frame];
}

pxl8_result pxl8_tilesheet_set_tile_pixels(pxl8_tilesheet* tilesheet, u16 tile_id, const u8* pixels) {
    if (!tilesheet || !pixels || tile_id == 0 || tile_id > tilesheet->total_tiles) return PXL8_ERROR_INVALID_ARGUMENT;
    if (!tilesheet->data) return PXL8_ERROR_NULL_POINTER;

    u32 tile_x = (tile_id - 1) % tilesheet->tiles_per_row;
    u32 tile_y = (tile_id - 1) / tilesheet->tiles_per_row;
    u32 bytes_per_pixel = (tilesheet->color_mode == PXL8_COLOR_MODE_HICOLOR) ? 4 : 1;

    for (u32 py = 0; py < tilesheet->tile_size; py++) {
        for (u32 px = 0; px < tilesheet->tile_size; px++) {
            u32 src_idx = py * tilesheet->tile_size + px;
            u32 dst_x = tile_x * tilesheet->tile_size + px;
            u32 dst_y = tile_y * tilesheet->tile_size + py;
            u32 dst_idx = (dst_y * tilesheet->width + dst_x) * bytes_per_pixel;

            if (bytes_per_pixel == 4) {
                ((u32*)tilesheet->data)[dst_idx / 4] = pixels[src_idx];
            } else {
                tilesheet->data[dst_idx] = pixels[src_idx];
            }
        }
    }

    if (tilesheet->tile_valid) {
        tilesheet->tile_valid[tile_id] = true;
    }

    return PXL8_OK;
}

void pxl8_tilesheet_set_tile_property(pxl8_tilesheet* tilesheet, u16 tile_id,
                                      const pxl8_tile_properties* props) {
    if (!tilesheet || !props || tile_id == 0 || tile_id > tilesheet->total_tiles) return;

    if (!tilesheet->properties) {
        tilesheet->properties = calloc(tilesheet->total_tiles + 1, sizeof(pxl8_tile_properties));
        if (!tilesheet->properties) return;
    }

    tilesheet->properties[tile_id] = *props;
}

const pxl8_tile_properties* pxl8_tilesheet_get_tile_property(const pxl8_tilesheet* tilesheet, u16 tile_id) {
    if (!tilesheet || !tilesheet->properties || tile_id == 0 || tile_id > tilesheet->total_tiles) {
        return NULL;
    }
    return &tilesheet->properties[tile_id];
}

u32 pxl8_tilesheet_get_tile_size(const pxl8_tilesheet* tilesheet) {
    return tilesheet ? tilesheet->tile_size : 0;
}

pxl8_result pxl8_tilesheet_add_autotile_rule(pxl8_tilesheet* tilesheet, u16 base_tile_id,
                                             u8 neighbor_mask, u16 result_tile_id) {
    if (!tilesheet || base_tile_id == 0 || base_tile_id > tilesheet->total_tiles) {
        return PXL8_ERROR_INVALID_ARGUMENT;
    }

    if (!tilesheet->autotile_rules) {
        tilesheet->autotile_rules = calloc(tilesheet->total_tiles + 1, sizeof(pxl8_autotile_rule*));
        tilesheet->autotile_rule_counts = calloc(tilesheet->total_tiles + 1, sizeof(u32));
        if (!tilesheet->autotile_rules || !tilesheet->autotile_rule_counts) {
            return PXL8_ERROR_OUT_OF_MEMORY;
        }
    }

    u32 count = tilesheet->autotile_rule_counts[base_tile_id];
    pxl8_autotile_rule* new_rules = realloc(tilesheet->autotile_rules[base_tile_id],
                                            (count + 1) * sizeof(pxl8_autotile_rule));
    if (!new_rules) return PXL8_ERROR_OUT_OF_MEMORY;

    new_rules[count].neighbor_mask = neighbor_mask;
    new_rules[count].tile_id = result_tile_id;

    tilesheet->autotile_rules[base_tile_id] = new_rules;
    tilesheet->autotile_rule_counts[base_tile_id]++;

    return PXL8_OK;
}

u16 pxl8_tilesheet_apply_autotile(const pxl8_tilesheet* tilesheet, u16 base_tile_id, u8 neighbors) {
    if (!tilesheet || !tilesheet->autotile_rules || base_tile_id == 0 ||
        base_tile_id > tilesheet->total_tiles) {
        return base_tile_id;
    }

    pxl8_autotile_rule* rules = tilesheet->autotile_rules[base_tile_id];
    u32 rule_count = tilesheet->autotile_rule_counts[base_tile_id];

    for (u32 i = 0; i < rule_count; i++) {
        if (rules[i].neighbor_mask == neighbors) {
            return rules[i].tile_id;
        }
    }

    return base_tile_id;
}