#include "pxl8_tilemap.h"
#include "pxl8_tilesheet.h"
#include "pxl8_macros.h"
#include <stdlib.h>
#include <string.h>

static inline u32 pxl8_chunk_index(u32 x, u32 y, u32 chunks_wide) {
    return (y >> 4) * chunks_wide + (x >> 4);
}

static pxl8_tile_chunk* pxl8_get_or_create_chunk(pxl8_tilemap_layer* layer, u32 x, u32 y) {
    u32 chunk_x = x >> 4;
    u32 chunk_y = y >> 4;
    u32 idx = chunk_y * layer->chunks_wide + chunk_x;

    if (idx >= layer->chunks_wide * layer->chunks_high) return NULL;

    if (!layer->chunks[idx]) {
        layer->chunks[idx] = calloc(1, sizeof(pxl8_tile_chunk));
        if (!layer->chunks[idx]) return NULL;

        layer->chunks[idx]->chunk_x = chunk_x;
        layer->chunks[idx]->chunk_y = chunk_y;
        layer->chunks[idx]->empty = true;
        layer->allocated_chunks++;
    }

    return layer->chunks[idx];
}

pxl8_result pxl8_tilemap_init(pxl8_tilemap* tilemap, u32 width, u32 height, u32 tile_size) {
    if (!tilemap) return PXL8_ERROR_NULL_POINTER;
    if (width > PXL8_MAX_TILEMAP_WIDTH || height > PXL8_MAX_TILEMAP_HEIGHT) {
        return PXL8_ERROR_INVALID_SIZE;
    }

    memset(tilemap, 0, sizeof(pxl8_tilemap));
    tilemap->width = width;
    tilemap->height = height;
    tilemap->tile_size = tile_size ? tile_size : PXL8_TILE_SIZE;
    tilemap->active_layers = 1;

    u32 chunks_wide = (width + PXL8_CHUNK_SIZE - 1) / PXL8_CHUNK_SIZE;
    u32 chunks_high = (height + PXL8_CHUNK_SIZE - 1) / PXL8_CHUNK_SIZE;

    for (u32 i = 0; i < PXL8_MAX_TILE_LAYERS; i++) {
        pxl8_tilemap_layer* layer = &tilemap->layers[i];
        layer->chunks_wide = chunks_wide;
        layer->chunks_high = chunks_high;
        layer->chunk_count = chunks_wide * chunks_high;
        layer->visible = (i == 0);
        layer->opacity = 255;

        layer->chunks = calloc(layer->chunk_count, sizeof(pxl8_tile_chunk*));
        if (!layer->chunks) {
            for (u32 j = 0; j < i; j++) {
                free(tilemap->layers[j].chunks);
            }
            return PXL8_ERROR_OUT_OF_MEMORY;
        }
    }

    return PXL8_OK;
}

void pxl8_tilemap_free(pxl8_tilemap* tilemap) {
    if (!tilemap) return;

    for (u32 i = 0; i < PXL8_MAX_TILE_LAYERS; i++) {
        pxl8_tilemap_layer* layer = &tilemap->layers[i];
        if (layer->chunks) {
            for (u32 j = 0; j < layer->chunk_count; j++) {
                if (layer->chunks[j]) {
                    free(layer->chunks[j]);
                }
            }
            free(layer->chunks);
            layer->chunks = NULL;
        }
    }

    if (tilemap->tilesheet) {
        pxl8_tilesheet_unref(tilemap->tilesheet);
        tilemap->tilesheet = NULL;
    }
}

pxl8_result pxl8_tilemap_set_tilesheet(pxl8_tilemap* tilemap, pxl8_tilesheet* tilesheet) {
    if (!tilemap || !tilesheet) return PXL8_ERROR_NULL_POINTER;

    if (tilemap->tilesheet) {
        pxl8_tilesheet_unref(tilemap->tilesheet);
    }

    tilemap->tilesheet = tilesheet;
    pxl8_tilesheet_ref(tilesheet);

    if (tilesheet->tile_size != tilemap->tile_size) {
        pxl8_warn("Tilesheet tile size (%d) differs from tilemap tile size (%d)",
                  tilesheet->tile_size, tilemap->tile_size);
        tilemap->tile_size = tilesheet->tile_size;
    }

    return PXL8_OK;
}

pxl8_result pxl8_tilemap_set_tile(pxl8_tilemap* tilemap, u32 layer, u32 x, u32 y, u16 tile_id, u8 flags) {
    if (!tilemap) return PXL8_ERROR_NULL_POINTER;
    if (layer >= PXL8_MAX_TILE_LAYERS) return PXL8_ERROR_INVALID_ARGUMENT;
    if (x >= tilemap->width || y >= tilemap->height) return PXL8_ERROR_INVALID_COORDINATE;

    pxl8_tilemap_layer* l = &tilemap->layers[layer];
    pxl8_tile_chunk* chunk = pxl8_get_or_create_chunk(l, x, y);
    if (!chunk) return PXL8_ERROR_OUT_OF_MEMORY;

    u32 local_x = x & PXL8_CHUNK_MASK;
    u32 local_y = y & PXL8_CHUNK_MASK;
    u32 idx = local_y * PXL8_CHUNK_SIZE + local_x;

    chunk->tiles[idx] = pxl8_tile_pack(tile_id, flags, 0);

    if (tile_id != 0) {
        chunk->empty = false;
    }

    if (layer >= tilemap->active_layers) {
        tilemap->active_layers = layer + 1;
        l->visible = true;
    }

    return PXL8_OK;
}

pxl8_tile pxl8_tilemap_get_tile(const pxl8_tilemap* tilemap, u32 layer, u32 x, u32 y) {
    if (!tilemap || layer >= PXL8_MAX_TILE_LAYERS) return 0;
    if (x >= tilemap->width || y >= tilemap->height) return 0;

    const pxl8_tilemap_layer* l = &tilemap->layers[layer];
    u32 chunk_idx = pxl8_chunk_index(x, y, l->chunks_wide);

    if (chunk_idx >= l->chunk_count || !l->chunks[chunk_idx]) return 0;

    const pxl8_tile_chunk* chunk = l->chunks[chunk_idx];
    u32 local_x = x & PXL8_CHUNK_MASK;
    u32 local_y = y & PXL8_CHUNK_MASK;
    u32 idx = local_y * PXL8_CHUNK_SIZE + local_x;

    return chunk->tiles[idx];
}

void pxl8_tilemap_set_camera(pxl8_tilemap* tilemap, i32 x, i32 y) {
    if (!tilemap) return;
    tilemap->camera_x = x;
    tilemap->camera_y = y;
}

void pxl8_tilemap_get_view(const pxl8_tilemap* tilemap, const pxl8_gfx_ctx* gfx, pxl8_tilemap_view* view) {
    if (!tilemap || !gfx || !view) return;

    view->x = -tilemap->camera_x;
    view->y = -tilemap->camera_y;
    view->width = gfx->framebuffer_width;
    view->height = gfx->framebuffer_height;

    view->tile_start_x = pxl8_max(0, tilemap->camera_x / (i32)tilemap->tile_size);
    view->tile_start_y = pxl8_max(0, tilemap->camera_y / (i32)tilemap->tile_size);
    view->tile_end_x = pxl8_min((i32)tilemap->width,
                                 (tilemap->camera_x + view->width) / (i32)tilemap->tile_size + 1);
    view->tile_end_y = pxl8_min((i32)tilemap->height,
                                 (tilemap->camera_y + view->height) / (i32)tilemap->tile_size + 1);
}

void pxl8_tilemap_render_tile(const pxl8_tilemap* tilemap, pxl8_gfx_ctx* gfx, u16 tile_id, i32 x, i32 y, u8 flags) {
    if (!tilemap || !gfx || !tilemap->tilesheet) return;
    pxl8_tilesheet_render_tile(tilemap->tilesheet, gfx, tile_id, x, y, flags);
}

void pxl8_tilemap_render_layer(const pxl8_tilemap* tilemap, pxl8_gfx_ctx* gfx, u32 layer) {
    if (!tilemap || !gfx || layer >= tilemap->active_layers) return;
    if (!tilemap->tilesheet) {
        pxl8_warn("No tilesheet set for tilemap");
        return;
    }

    const pxl8_tilemap_layer* l = &tilemap->layers[layer];
    if (!l->visible) return;

    i32 view_left = tilemap->camera_x / tilemap->tile_size;
    i32 view_top = tilemap->camera_y / tilemap->tile_size;
    i32 view_right = (tilemap->camera_x + gfx->framebuffer_width) / tilemap->tile_size + 1;
    i32 view_bottom = (tilemap->camera_y + gfx->framebuffer_height) / tilemap->tile_size + 1;

    u32 chunk_left = pxl8_max(0, view_left >> 4);
    u32 chunk_top = pxl8_max(0, view_top >> 4);
    u32 chunk_right = pxl8_min((tilemap->width + 15) >> 4, (view_right >> 4) + 1);
    u32 chunk_bottom = pxl8_min((tilemap->height + 15) >> 4, (view_bottom >> 4) + 1);

    for (u32 cy = chunk_top; cy < chunk_bottom; cy++) {
        for (u32 cx = chunk_left; cx < chunk_right; cx++) {
            u32 chunk_idx = cy * l->chunks_wide + cx;
            if (chunk_idx >= l->chunk_count || !l->chunks[chunk_idx]) continue;

            const pxl8_tile_chunk* chunk = l->chunks[chunk_idx];
            if (chunk->empty) continue;

            u32 tile_start_x = pxl8_max(0, view_left - (i32)(cx << 4));
            u32 tile_end_x = pxl8_min(PXL8_CHUNK_SIZE, view_right - (i32)(cx << 4));
            u32 tile_start_y = pxl8_max(0, view_top - (i32)(cy << 4));
            u32 tile_end_y = pxl8_min(PXL8_CHUNK_SIZE, view_bottom - (i32)(cy << 4));

            for (u32 ty = tile_start_y; ty < tile_end_y; ty++) {
                for (u32 tx = tile_start_x; tx < tile_end_x; tx++) {
                    u32 idx = ty * PXL8_CHUNK_SIZE + tx;
                    pxl8_tile tile = chunk->tiles[idx];
                    u16 tile_id = pxl8_tile_get_id(tile);
                    if (tile_id == 0) continue;

                    i32 world_x = (cx << 4) + tx;
                    i32 world_y = (cy << 4) + ty;
                    i32 screen_x = world_x * tilemap->tile_size - tilemap->camera_x;
                    i32 screen_y = world_y * tilemap->tile_size - tilemap->camera_y;

                    u8 flags = pxl8_tile_get_flags(tile);
                    if (flags & PXL8_TILE_ANIMATED && tilemap->tilesheet) {
                        tile_id = pxl8_tilesheet_get_animated_frame(tilemap->tilesheet, tile_id);
                    }

                    pxl8_tilesheet_render_tile(tilemap->tilesheet, gfx, tile_id, screen_x, screen_y, flags);
                }
            }
        }
    }
}

void pxl8_tilemap_render(const pxl8_tilemap* tilemap, pxl8_gfx_ctx* gfx) {
    if (!tilemap || !gfx) return;

    for (u32 layer = 0; layer < tilemap->active_layers; layer++) {
        pxl8_tilemap_render_layer(tilemap, gfx, layer);
    }
}

bool pxl8_tilemap_is_solid(const pxl8_tilemap* tilemap, u32 x, u32 y) {
    if (!tilemap) return true;

    u32 tile_x = x / tilemap->tile_size;
    u32 tile_y = y / tilemap->tile_size;

    for (u32 layer = 0; layer < tilemap->active_layers; layer++) {
        pxl8_tile tile = pxl8_tilemap_get_tile(tilemap, layer, tile_x, tile_y);
        if (pxl8_tile_get_flags(tile) & PXL8_TILE_SOLID) return true;
    }

    return false;
}

bool pxl8_tilemap_check_collision(const pxl8_tilemap* tilemap, i32 x, i32 y, i32 w, i32 h) {
    if (!tilemap) return true;

    i32 left = x / tilemap->tile_size;
    i32 top = y / tilemap->tile_size;
    i32 right = (x + w - 1) / tilemap->tile_size;
    i32 bottom = (y + h - 1) / tilemap->tile_size;

    for (i32 ty = top; ty <= bottom; ty++) {
        for (i32 tx = left; tx <= right; tx++) {
            if (tx < 0 || tx >= (i32)tilemap->width ||
                ty < 0 || ty >= (i32)tilemap->height) return true;

            for (u32 layer = 0; layer < tilemap->active_layers; layer++) {
                pxl8_tile tile = pxl8_tilemap_get_tile(tilemap, layer, tx, ty);
                if (pxl8_tile_get_flags(tile) & PXL8_TILE_SOLID) return true;
            }
        }
    }

    return false;
}

pxl8_tilemap* pxl8_tilemap_new(u32 width, u32 height, u32 tile_size) {
    pxl8_tilemap* tilemap = calloc(1, sizeof(pxl8_tilemap));
    if (!tilemap) {
        pxl8_error("Failed to allocate tilemap");
        return NULL;
    }

    pxl8_result result = pxl8_tilemap_init(tilemap, width, height, tile_size);
    if (result != PXL8_OK) {
        pxl8_error("Failed to initialize tilemap");
        free(tilemap);
        return NULL;
    }

    return tilemap;
}

void pxl8_tilemap_destroy(pxl8_tilemap* tilemap) {
    if (!tilemap) return;
    pxl8_tilemap_free(tilemap);
    free(tilemap);
}

u16 pxl8_tilemap_get_tile_id(const pxl8_tilemap* tilemap, u32 layer, u32 x, u32 y) {
    if (!tilemap) return 0;
    pxl8_tile tile = pxl8_tilemap_get_tile(tilemap, layer, x, y);
    return pxl8_tile_get_id(tile);
}

void pxl8_tilemap_update(pxl8_tilemap* tilemap, f32 delta_time) {
    if (!tilemap || !tilemap->tilesheet) return;
    pxl8_tilesheet_update_animations(tilemap->tilesheet, delta_time);
}

static u8 pxl8_tilemap_get_neighbors(const pxl8_tilemap* tilemap, u32 layer, u32 x, u32 y, u16 match_id) {
    u8 neighbors = 0;

    if (y > 0 && pxl8_tilemap_get_tile_id(tilemap, layer, x, y - 1) == match_id)
        neighbors |= 1 << 0;
    if (x < tilemap->width - 1 && pxl8_tilemap_get_tile_id(tilemap, layer, x + 1, y) == match_id)
        neighbors |= 1 << 1;
    if (y < tilemap->height - 1 && pxl8_tilemap_get_tile_id(tilemap, layer, x, y + 1) == match_id)
        neighbors |= 1 << 2;
    if (x > 0 && pxl8_tilemap_get_tile_id(tilemap, layer, x - 1, y) == match_id)
        neighbors |= 1 << 3;

    if (y > 0 && x < tilemap->width - 1 && pxl8_tilemap_get_tile_id(tilemap, layer, x + 1, y - 1) == match_id)
        neighbors |= 1 << 4;
    if (y < tilemap->height - 1 && x < tilemap->width - 1 && pxl8_tilemap_get_tile_id(tilemap, layer, x + 1, y + 1) == match_id)
        neighbors |= 1 << 5;
    if (y < tilemap->height - 1 && x > 0 && pxl8_tilemap_get_tile_id(tilemap, layer, x - 1, y + 1) == match_id)
        neighbors |= 1 << 6;
    if (y > 0 && x > 0 && pxl8_tilemap_get_tile_id(tilemap, layer, x - 1, y - 1) == match_id)
        neighbors |= 1 << 7;

    return neighbors;
}

pxl8_result pxl8_tilemap_set_tile_auto(pxl8_tilemap* tilemap, u32 layer, u32 x, u32 y,
                                       u16 base_tile_id, u8 flags) {
    if (!tilemap || !tilemap->tilesheet) return PXL8_ERROR_NULL_POINTER;

    pxl8_result result = pxl8_tilemap_set_tile(tilemap, layer, x, y, base_tile_id, flags | PXL8_TILE_AUTOTILE);
    if (result != PXL8_OK) return result;

    pxl8_tilemap_update_autotiles(tilemap, layer,
                                  x > 0 ? x - 1 : x, y > 0 ? y - 1 : y,
                                  x < tilemap->width - 1 ? 3 : 2,
                                  y < tilemap->height - 1 ? 3 : 2);

    return PXL8_OK;
}

void pxl8_tilemap_update_autotiles(pxl8_tilemap* tilemap, u32 layer, u32 x, u32 y, u32 w, u32 h) {
    if (!tilemap || !tilemap->tilesheet) return;

    for (u32 ty = y; ty < y + h && ty < tilemap->height; ty++) {
        for (u32 tx = x; tx < x + w && tx < tilemap->width; tx++) {
            pxl8_tile tile = pxl8_tilemap_get_tile(tilemap, layer, tx, ty);
            u8 flags = pxl8_tile_get_flags(tile);

            if (flags & PXL8_TILE_AUTOTILE) {
                u16 base_id = pxl8_tile_get_id(tile);
                u8 neighbors = pxl8_tilemap_get_neighbors(tilemap, layer, tx, ty, base_id);
                u16 new_id = pxl8_tilesheet_apply_autotile(tilemap->tilesheet, base_id, neighbors);

                if (new_id != base_id) {
                    pxl8_tilemap_layer* l = &tilemap->layers[layer];
                    pxl8_tile_chunk* chunk = pxl8_get_or_create_chunk(l, tx, ty);
                    if (chunk) {
                        u32 local_x = tx & PXL8_CHUNK_MASK;
                        u32 local_y = ty & PXL8_CHUNK_MASK;
                        u32 idx = local_y * PXL8_CHUNK_SIZE + local_x;
                        chunk->tiles[idx] = pxl8_tile_pack(new_id, flags, pxl8_tile_get_palette(tile));
                    }
                }
            }
        }
    }
}

u32 pxl8_tilemap_get_memory_usage(const pxl8_tilemap* tilemap) {
    if (!tilemap) return 0;

    u32 total = sizeof(pxl8_tilemap);

    for (u32 i = 0; i < PXL8_MAX_TILE_LAYERS; i++) {
        const pxl8_tilemap_layer* layer = &tilemap->layers[i];
        total += layer->chunk_count * sizeof(pxl8_tile_chunk*);
        total += layer->allocated_chunks * sizeof(pxl8_tile_chunk);
    }

    return total;
}

void pxl8_tilemap_compress(pxl8_tilemap* tilemap) {
    if (!tilemap) return;

    for (u32 i = 0; i < PXL8_MAX_TILE_LAYERS; i++) {
        pxl8_tilemap_layer* layer = &tilemap->layers[i];

        for (u32 j = 0; j < layer->chunk_count; j++) {
            pxl8_tile_chunk* chunk = layer->chunks[j];
            if (!chunk) continue;

            bool has_tiles = false;
            for (u32 k = 0; k < PXL8_CHUNK_SIZE * PXL8_CHUNK_SIZE; k++) {
                if (pxl8_tile_get_id(chunk->tiles[k]) != 0) {
                    has_tiles = true;
                    break;
                }
            }

            if (!has_tiles) {
                free(chunk);
                layer->chunks[j] = NULL;
                layer->allocated_chunks--;
            } else {
                chunk->empty = false;
            }
        }
    }
}