#include "pxl8_cart.h"

#include <dirent.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

#include "pxl8_log.h"

#define PXL8_CART_MAGIC 0x43585850
#define PXL8_CART_VERSION 1
#define PXL8_CART_TRAILER_MAGIC 0x544E4250

typedef struct {
    u32 magic;
    u16 version;
    u16 flags;
    u32 file_count;
    u32 toc_size;
} pxl8_cart_header;

typedef struct {
    u32 offset;
    u32 size;
    u16 path_len;
} pxl8_cart_entry;

typedef struct {
    u32 magic;
    u32 cart_offset;
    u32 cart_size;
} pxl8_cart_trailer;

typedef struct {
    char* path;
    u32 offset;
    u32 size;
} pxl8_cart_file;

struct pxl8_cart {
    u8* data;
    u32 data_size;
    pxl8_cart_file* files;
    u32 file_count;
    char* base_path;
    char* title;
    pxl8_resolution resolution;
    pxl8_size window_size;
    pxl8_pixel_mode pixel_mode;
    bool is_folder;
    bool is_mounted;
};

static pxl8_cart* pxl8_current_cart = NULL;
static char* pxl8_original_cwd = NULL;

static bool is_directory(const char* path) {
    struct stat st;
    return stat(path, &st) == 0 && S_ISDIR(st.st_mode);
}

static bool is_pxc_file(const char* path) {
    size_t len = strlen(path);
    return len > 4 && strcmp(path + len - 4, ".pxc") == 0;
}

static bool has_main_script(const char* base_path) {
    char path[512];
    snprintf(path, sizeof(path), "%s/main.fnl", base_path);
    if (access(path, F_OK) == 0) return true;
    snprintf(path, sizeof(path), "%s/main.lua", base_path);
    return access(path, F_OK) == 0;
}

static void collect_files_recursive(const char* dir_path, const char* prefix,
                                    char*** paths, u32* count, u32* capacity) {
    DIR* dir = opendir(dir_path);
    if (!dir) return;

    struct dirent* entry;
    while ((entry = readdir(dir)) != NULL) {
        if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0) continue;

        char full_path[1024];
        char rel_path[1024];
        snprintf(full_path, sizeof(full_path), "%s/%s", dir_path, entry->d_name);
        snprintf(rel_path, sizeof(rel_path), "%s%s", prefix, entry->d_name);

        struct stat st;
        if (stat(full_path, &st) == 0) {
            if (S_ISDIR(st.st_mode)) {
                char new_prefix[1025];
                snprintf(new_prefix, sizeof(new_prefix), "%s/", rel_path);
                collect_files_recursive(full_path, new_prefix, paths, count, capacity);
            } else {
                if (*count >= *capacity) {
                    *capacity = (*capacity == 0) ? 64 : (*capacity * 2);
                    *paths = realloc(*paths, *capacity * sizeof(char*));
                }
                (*paths)[(*count)++] = strdup(rel_path);
            }
        }
    }
    closedir(dir);
}

static pxl8_cart_file* find_file(const pxl8_cart* cart, const char* path) {
    if (!cart || !cart->files) return NULL;
    for (u32 i = 0; i < cart->file_count; i++) {
        if (strcmp(cart->files[i].path, path) == 0) {
            return &cart->files[i];
        }
    }
    return NULL;
}

static pxl8_result load_packed_cart(pxl8_cart* cart, const u8* data, u32 size) {
    if (size < sizeof(pxl8_cart_header)) return PXL8_ERROR_INVALID_FORMAT;

    const pxl8_cart_header* header = (const pxl8_cart_header*)data;
    if (header->magic != PXL8_CART_MAGIC) return PXL8_ERROR_INVALID_FORMAT;
    if (header->version > PXL8_CART_VERSION) return PXL8_ERROR_INVALID_FORMAT;

    cart->data = malloc(size);
    if (!cart->data) return PXL8_ERROR_OUT_OF_MEMORY;
    memcpy(cart->data, data, size);
    cart->data_size = size;

    cart->file_count = header->file_count;
    cart->files = calloc(cart->file_count, sizeof(pxl8_cart_file));
    if (!cart->files) return PXL8_ERROR_OUT_OF_MEMORY;

    const u8* toc = cart->data + sizeof(pxl8_cart_header);
    for (u32 i = 0; i < cart->file_count; i++) {
        const pxl8_cart_entry* entry = (const pxl8_cart_entry*)toc;
        toc += sizeof(pxl8_cart_entry);

        cart->files[i].path = malloc(entry->path_len + 1);
        memcpy(cart->files[i].path, toc, entry->path_len);
        cart->files[i].path[entry->path_len] = '\0';
        toc += entry->path_len;

        cart->files[i].offset = entry->offset;
        cart->files[i].size = entry->size;
    }

    cart->is_folder = false;
    return PXL8_OK;
}

pxl8_cart* pxl8_cart_create(void) {
    pxl8_cart* cart = calloc(1, sizeof(pxl8_cart));
    if (cart) {
        cart->resolution = PXL8_RESOLUTION_640x360;
        cart->window_size = (pxl8_size){1280, 720};
        cart->pixel_mode = PXL8_PIXEL_INDEXED;
    }
    return cart;
}

pxl8_cart* pxl8_get_cart(void) {
    return pxl8_current_cart;
}

void pxl8_cart_destroy(pxl8_cart* cart) {
    if (!cart) return;
    pxl8_cart_unload(cart);
    free(cart);
}

pxl8_result pxl8_cart_load(pxl8_cart* cart, const char* path) {
    if (!cart || !path) return PXL8_ERROR_NULL_POINTER;
    pxl8_cart_unload(cart);

    if (is_directory(path)) {
        cart->base_path = realpath(path, NULL);
        if (!cart->base_path) {
            pxl8_error("Failed to resolve cart path: %s", path);
            return PXL8_ERROR_FILE_NOT_FOUND;
        }
        cart->is_folder = true;

        if (!has_main_script(cart->base_path)) {
            pxl8_error("No main.fnl or main.lua found in cart: %s", path);
            return PXL8_ERROR_FILE_NOT_FOUND;
        }

        pxl8_info("Loaded cart");
        return PXL8_OK;
    }

    if (is_pxc_file(path)) {
        FILE* file = fopen(path, "rb");
        if (!file) {
            pxl8_error("Failed to open cart: %s", path);
            return PXL8_ERROR_FILE_NOT_FOUND;
        }

        fseek(file, 0, SEEK_END);
        u32 size = (u32)ftell(file);
        fseek(file, 0, SEEK_SET);

        u8* data = malloc(size);
        if (!data || fread(data, 1, size, file) != size) {
            free(data);
            fclose(file);
            return PXL8_ERROR_SYSTEM_FAILURE;
        }
        fclose(file);

        pxl8_result result = load_packed_cart(cart, data, size);
        free(data);

        if (result == PXL8_OK) {
            pxl8_info("Loaded cart");
        }
        return result;
    }

    pxl8_error("Unknown cart format: %s", path);
    return PXL8_ERROR_INVALID_FORMAT;
}

pxl8_result pxl8_cart_load_embedded(pxl8_cart* cart, const char* exe_path) {
    if (!cart || !exe_path) return PXL8_ERROR_NULL_POINTER;

    FILE* file = fopen(exe_path, "rb");
    if (!file) return PXL8_ERROR_FILE_NOT_FOUND;

    fseek(file, -(i32)sizeof(pxl8_cart_trailer), SEEK_END);
    pxl8_cart_trailer trailer;
    if (fread(&trailer, sizeof(trailer), 1, file) != 1) {
        fclose(file);
        return PXL8_ERROR_INVALID_FORMAT;
    }

    if (trailer.magic != PXL8_CART_TRAILER_MAGIC) {
        fclose(file);
        return PXL8_ERROR_INVALID_FORMAT;
    }

    fseek(file, trailer.cart_offset, SEEK_SET);
    u8* data = malloc(trailer.cart_size);
    if (!data || fread(data, 1, trailer.cart_size, file) != trailer.cart_size) {
        free(data);
        fclose(file);
        return PXL8_ERROR_SYSTEM_FAILURE;
    }
    fclose(file);

    pxl8_result result = load_packed_cart(cart, data, trailer.cart_size);
    free(data);

    if (result == PXL8_OK) {
        pxl8_info("Loaded embedded cart");
    }
    return result;
}

void pxl8_cart_unload(pxl8_cart* cart) {
    if (!cart) return;

    if (cart->title) {
        pxl8_info("Unloaded cart: %s", cart->title);
        pxl8_cart_unmount(cart);
        free(cart->title);
        cart->title = NULL;
    } else if (cart->base_path || cart->data) {
        pxl8_info("Unloaded cart");
        pxl8_cart_unmount(cart);
    }

    if (cart->files) {
        for (u32 i = 0; i < cart->file_count; i++) {
            free(cart->files[i].path);
        }
        free(cart->files);
        cart->files = NULL;
    }
    cart->file_count = 0;

    free(cart->data);
    cart->data = NULL;
    cart->data_size = 0;

    free(cart->base_path);
    cart->base_path = NULL;
    cart->is_folder = false;
}

pxl8_result pxl8_cart_mount(pxl8_cart* cart) {
    if (!cart) return PXL8_ERROR_NULL_POINTER;
    if (cart->is_mounted) return PXL8_OK;

    if (pxl8_current_cart) {
        pxl8_cart_unmount(pxl8_current_cart);
    }

    if (cart->is_folder) {
        pxl8_original_cwd = getcwd(NULL, 0);
        if (chdir(cart->base_path) != 0) {
            pxl8_error("Failed to change to cart directory: %s", cart->base_path);
            free(pxl8_original_cwd);
            pxl8_original_cwd = NULL;
            return PXL8_ERROR_FILE_NOT_FOUND;
        }
    }

    cart->is_mounted = true;
    pxl8_current_cart = cart;
    if (cart->title) {
        pxl8_info("Mounted cart: %s", cart->title);
    } else {
        pxl8_info("Mounted cart");
    }
    return PXL8_OK;
}

void pxl8_cart_unmount(pxl8_cart* cart) {
    if (!cart || !cart->is_mounted) return;

    if (pxl8_original_cwd) {
        chdir(pxl8_original_cwd);
        free(pxl8_original_cwd);
        pxl8_original_cwd = NULL;
    }

    cart->is_mounted = false;
    if (pxl8_current_cart == cart) {
        pxl8_current_cart = NULL;
    }
}

const char* pxl8_cart_get_base_path(const pxl8_cart* cart) {
    return cart ? cart->base_path : NULL;
}

const char* pxl8_cart_get_title(const pxl8_cart* cart) {
    return cart ? cart->title : NULL;
}

void pxl8_cart_set_title(pxl8_cart* cart, const char* title) {
    if (!cart || !title) return;
    free(cart->title);
    cart->title = strdup(title);
}

pxl8_resolution pxl8_cart_get_resolution(const pxl8_cart* cart) {
    return cart ? cart->resolution : PXL8_RESOLUTION_640x360;
}

void pxl8_cart_set_resolution(pxl8_cart* cart, pxl8_resolution resolution) {
    if (cart) cart->resolution = resolution;
}

pxl8_size pxl8_cart_get_window_size(const pxl8_cart* cart) {
    return cart ? cart->window_size : (pxl8_size){1280, 720};
}

void pxl8_cart_set_window_size(pxl8_cart* cart, pxl8_size size) {
    if (cart) cart->window_size = size;
}

pxl8_pixel_mode pxl8_cart_get_pixel_mode(const pxl8_cart* cart) {
    return cart ? cart->pixel_mode : PXL8_PIXEL_INDEXED;
}

void pxl8_cart_set_pixel_mode(pxl8_cart* cart, pxl8_pixel_mode mode) {
    if (cart) cart->pixel_mode = mode;
}

bool pxl8_cart_is_packed(const pxl8_cart* cart) {
    return cart && !cart->is_folder;
}

bool pxl8_cart_has_embedded(const char* exe_path) {
    FILE* file = fopen(exe_path, "rb");
    if (!file) return false;

    fseek(file, -(i32)sizeof(pxl8_cart_trailer), SEEK_END);
    pxl8_cart_trailer trailer;
    bool has = (fread(&trailer, sizeof(trailer), 1, file) == 1 &&
                trailer.magic == PXL8_CART_TRAILER_MAGIC);
    fclose(file);
    return has;
}

bool pxl8_cart_file_exists(const pxl8_cart* cart, const char* path) {
    if (!cart || !path) return false;

    if (cart->is_folder) {
        char full_path[512];
        if (!pxl8_cart_resolve_path(cart, path, full_path, sizeof(full_path))) return false;
        return access(full_path, F_OK) == 0;
    }

    return find_file(cart, path) != NULL;
}

bool pxl8_cart_resolve_path(const pxl8_cart* cart, const char* relative_path, char* out_path, size_t out_size) {
    if (!cart || !relative_path || !out_path || out_size == 0) return false;

    if (cart->is_folder && cart->base_path) {
        i32 written = snprintf(out_path, out_size, "%s/%s", cart->base_path, relative_path);
        return written >= 0 && (size_t)written < out_size;
    }

    i32 written = snprintf(out_path, out_size, "%s", relative_path);
    return written >= 0 && (size_t)written < out_size;
}

pxl8_result pxl8_cart_read_file(const pxl8_cart* cart, const char* path, u8** data_out, u32* size_out) {
    if (!cart || !path || !data_out || !size_out) return PXL8_ERROR_NULL_POINTER;

    if (cart->is_folder) {
        char full_path[512];
        if (!pxl8_cart_resolve_path(cart, path, full_path, sizeof(full_path))) {
            return PXL8_ERROR_FILE_NOT_FOUND;
        }

        FILE* file = fopen(full_path, "rb");
        if (!file) return PXL8_ERROR_FILE_NOT_FOUND;

        fseek(file, 0, SEEK_END);
        *size_out = (u32)ftell(file);
        fseek(file, 0, SEEK_SET);

        *data_out = malloc(*size_out);
        if (!*data_out || fread(*data_out, 1, *size_out, file) != *size_out) {
            free(*data_out);
            *data_out = NULL;
            fclose(file);
            return PXL8_ERROR_SYSTEM_FAILURE;
        }
        fclose(file);
        return PXL8_OK;
    }

    pxl8_cart_file* cf = find_file(cart, path);
    if (!cf) return PXL8_ERROR_FILE_NOT_FOUND;

    *size_out = cf->size;
    *data_out = malloc(cf->size);
    if (!*data_out) return PXL8_ERROR_OUT_OF_MEMORY;

    memcpy(*data_out, cart->data + cf->offset, cf->size);
    return PXL8_OK;
}

void pxl8_cart_free_file(u8* data) {
    free(data);
}

pxl8_result pxl8_cart_pack(const char* folder_path, const char* output_path) {
    if (!folder_path || !output_path) return PXL8_ERROR_NULL_POINTER;

    if (!is_directory(folder_path)) {
        pxl8_error("Cart folder not found: %s", folder_path);
        return PXL8_ERROR_FILE_NOT_FOUND;
    }

    if (!has_main_script(folder_path)) {
        pxl8_error("No main.fnl or main.lua found in cart: %s", folder_path);
        return PXL8_ERROR_FILE_NOT_FOUND;
    }

    char** paths = NULL;
    u32 count = 0, capacity = 0;
    collect_files_recursive(folder_path, "", &paths, &count, &capacity);

    if (count == 0) {
        pxl8_error("No files found in cart folder");
        return PXL8_ERROR_FILE_NOT_FOUND;
    }

    u32 toc_size = 0;
    for (u32 i = 0; i < count; i++) {
        toc_size += sizeof(pxl8_cart_entry) + strlen(paths[i]);
    }

    u32 data_offset = sizeof(pxl8_cart_header) + toc_size;
    u32 total_size = data_offset;

    u32* file_sizes = malloc(count * sizeof(u32));
    for (u32 i = 0; i < count; i++) {
        char full_path[1024];
        snprintf(full_path, sizeof(full_path), "%s/%s", folder_path, paths[i]);
        struct stat st;
        if (stat(full_path, &st) == 0) {
            file_sizes[i] = (u32)st.st_size;
            total_size += file_sizes[i];
        } else {
            file_sizes[i] = 0;
        }
    }

    u8* buffer = calloc(1, total_size);
    if (!buffer) {
        free(file_sizes);
        for (u32 i = 0; i < count; i++) free(paths[i]);
        free(paths);
        return PXL8_ERROR_OUT_OF_MEMORY;
    }

    pxl8_cart_header* header = (pxl8_cart_header*)buffer;
    header->magic = PXL8_CART_MAGIC;
    header->version = PXL8_CART_VERSION;
    header->flags = 0;
    header->file_count = count;
    header->toc_size = toc_size;

    u8* toc = buffer + sizeof(pxl8_cart_header);
    u32 file_offset = data_offset;

    pxl8_info("Packing cart: %s -> %s", folder_path, output_path);

    for (u32 i = 0; i < count; i++) {
        pxl8_cart_entry* entry = (pxl8_cart_entry*)toc;
        entry->offset = file_offset;
        entry->size = file_sizes[i];
        entry->path_len = (u16)strlen(paths[i]);
        toc += sizeof(pxl8_cart_entry);

        memcpy(toc, paths[i], entry->path_len);
        toc += entry->path_len;

        char full_path[1024];
        snprintf(full_path, sizeof(full_path), "%s/%s", folder_path, paths[i]);
        FILE* file = fopen(full_path, "rb");
        if (file) {
            fread(buffer + file_offset, 1, file_sizes[i], file);
            fclose(file);
            pxl8_info("  %s (%u bytes)", paths[i], file_sizes[i]);
        }

        file_offset += file_sizes[i];
        free(paths[i]);
    }
    free(paths);
    free(file_sizes);

    FILE* out = fopen(output_path, "wb");
    if (!out) {
        free(buffer);
        return PXL8_ERROR_SYSTEM_FAILURE;
    }

    fwrite(buffer, 1, total_size, out);
    fclose(out);
    free(buffer);

    pxl8_info("Cart packed: %u files, %u bytes", count, total_size);
    return PXL8_OK;
}

pxl8_result pxl8_cart_bundle(const char* input_path, const char* output_path, const char* exe_path) {
    if (!input_path || !output_path || !exe_path) return PXL8_ERROR_NULL_POINTER;

    u8* cart_data = NULL;
    u32 cart_size = 0;
    bool free_cart = false;

    if (is_directory(input_path)) {
        char temp_pxc[256];
        snprintf(temp_pxc, sizeof(temp_pxc), "/tmp/pxl8_bundle_%d.pxc", getpid());
        pxl8_result result = pxl8_cart_pack(input_path, temp_pxc);
        if (result != PXL8_OK) return result;

        FILE* f = fopen(temp_pxc, "rb");
        if (!f) return PXL8_ERROR_SYSTEM_FAILURE;
        fseek(f, 0, SEEK_END);
        cart_size = (u32)ftell(f);
        fseek(f, 0, SEEK_SET);
        cart_data = malloc(cart_size);
        fread(cart_data, 1, cart_size, f);
        fclose(f);
        unlink(temp_pxc);
        free_cart = true;
    } else if (is_pxc_file(input_path)) {
        FILE* f = fopen(input_path, "rb");
        if (!f) return PXL8_ERROR_FILE_NOT_FOUND;
        fseek(f, 0, SEEK_END);
        cart_size = (u32)ftell(f);
        fseek(f, 0, SEEK_SET);
        cart_data = malloc(cart_size);
        fread(cart_data, 1, cart_size, f);
        fclose(f);
        free_cart = true;
    } else {
        return PXL8_ERROR_INVALID_FORMAT;
    }

    FILE* exe = fopen(exe_path, "rb");
    if (!exe) {
        if (free_cart) free(cart_data);
        return PXL8_ERROR_FILE_NOT_FOUND;
    }
    fseek(exe, 0, SEEK_END);
    u32 exe_size = (u32)ftell(exe);
    fseek(exe, 0, SEEK_SET);

    u8* exe_data = malloc(exe_size);
    fread(exe_data, 1, exe_size, exe);
    fclose(exe);

    FILE* out = fopen(output_path, "wb");
    if (!out) {
        free(exe_data);
        if (free_cart) free(cart_data);
        return PXL8_ERROR_SYSTEM_FAILURE;
    }

    fwrite(exe_data, 1, exe_size, out);
    free(exe_data);

    u32 cart_offset = exe_size;
    fwrite(cart_data, 1, cart_size, out);
    if (free_cart) free(cart_data);

    pxl8_cart_trailer trailer = {
        .magic = PXL8_CART_TRAILER_MAGIC,
        .cart_offset = cart_offset,
        .cart_size = cart_size
    };
    fwrite(&trailer, sizeof(trailer), 1, out);
    fclose(out);

    chmod(output_path, 0755);

    pxl8_info("Bundle created: %s", output_path);
    return PXL8_OK;
}