#include "twn_textures_c.h"
#include "twn_util.h"
#include "twn_util_c.h"
#include "twn_engine_context_c.h"

#include <SDL2/SDL.h>
#include <physfs.h>
#include <physfsrwops.h>
#include <stb_ds.h>
#include <stb_rect_pack.h>
#include <stb_image.h>

#include <stdbool.h>

typedef struct {
    SDL_RWops *rwops;
    size_t size;
    size_t position;
} TextureLoadingContext;


static int load_read_callback(void *user, char *data, int size) {
    TextureLoadingContext *context = user;
    int read = (int)SDL_RWread(context->rwops, data, 1, size);
    context->position += read;
    if (read == 0)
        CRY_SDL( "Error in streamed texture load.");
    return read;
}

static void load_skip_callback(void *user, int n) {
    TextureLoadingContext *context = user;
    context->position += n;
    Sint64 result = SDL_RWseek(context->rwops, n, RW_SEEK_CUR);
    SDL_assert_always(result != -1);
}

static int load_eof_callback(void *user) {
    TextureLoadingContext *context = user;
    return context->position == context->size;
}

SDL_Surface *textures_load_surface(const char *path) {
    SDL_RWops *handle = PHYSFSRWOPS_openRead(path);
    if (handle == NULL)
        goto ERR_CANNOT_OPEN_FILE;

    TextureLoadingContext context = {
        .rwops = handle,
        .size = SDL_RWsize(handle),
    };

    stbi_io_callbacks callbacks = {
        .read = load_read_callback,
        .skip = load_skip_callback,
        .eof = load_eof_callback,
    };

    int width, height, channels;
    void *image_mem = stbi_load_from_callbacks(&callbacks, &context, &width, &height, &channels, 0);
    SDL_FreeRW(handle);
    if (!image_mem)
        goto ERR_CANNOT_READ_IMAGE;

    Uint32 rmask, gmask, bmask, amask;

    #if SDL_BYTEORDER == SDL_BIG_ENDIAN
        rmask = 0xff000000;
        gmask = 0x00ff0000;
        bmask = 0x0000ff00;
        amask = 0x000000ff;
    #else
        rmask = 0x000000ff;
        gmask = 0x0000ff00;
        bmask = 0x00ff0000;
        amask = 0xff000000;
    #endif

    SDL_Surface* surface = SDL_CreateRGBSurfaceFrom(image_mem, width, height,
                                                    channels * 8,
                                                    width * channels,
                                                    rmask, gmask, bmask, amask);
    if (surface == NULL)
        goto ERR_CANNOT_CREATE_SURFACE;

    SDL_SetSurfaceBlendMode(surface, SDL_BLENDMODE_NONE);
    SDL_SetSurfaceRLE(surface, true);

    return surface;

ERR_CANNOT_CREATE_SURFACE:
    stbi_image_free(image_mem);

ERR_CANNOT_READ_IMAGE:
ERR_CANNOT_OPEN_FILE:
    CRY(path, "Failed to load image. Aborting...");
    die_abruptly();
}


static SDL_Surface *create_surface(int width, int height) {
    Uint32 rmask, gmask, bmask, amask;

    #if SDL_BYTEORDER == SDL_BIG_ENDIAN
        rmask = 0xff000000;
        gmask = 0x00ff0000;
        bmask = 0x0000ff00;
        amask = 0x000000ff;
    #else
        rmask = 0x000000ff;
        gmask = 0x0000ff00;
        bmask = 0x00ff0000;
        amask = 0xff000000;
    #endif

    SDL_Surface *surface = SDL_CreateRGBSurface(0,
                                                width,
                                                height,
                                                TEXTURE_ATLAS_BIT_DEPTH,
                                                rmask,
                                                gmask,
                                                bmask,
                                                amask);

    SDL_SetSurfaceBlendMode(surface, SDL_BLENDMODE_NONE);
    SDL_SetSurfaceRLE(surface, true);

    return surface;
}


/* adds a new, blank atlas surface to the cache */
static void add_new_atlas(TextureCache *cache) {
    SDL_Surface *new_atlas = create_surface((int)ctx.texture_atlas_size, (int)ctx.texture_atlas_size);
    arrput(cache->atlas_surfaces, new_atlas);
    arrput(cache->atlas_textures, create_gpu_texture(TEXTURE_FILTER_NEAREAST, true));
}


static void upload_texture_from_surface(GPUTexture texture, SDL_Surface *surface) {
    SDL_LockSurface(surface);

    upload_gpu_texture(texture, surface->pixels, surface->format->BytesPerPixel, surface->w, surface->h);

    SDL_UnlockSurface(surface);
}


static void recreate_current_atlas_texture(TextureCache *cache) {
    /* TODO: should surfaces be freed after they cannot be referenced in atlas builing? */
    /*       for example, if full page of 64x64 tiles was already filled, there's no real reason to process them further */
    SDL_Surface *atlas_surface = cache->atlas_surfaces[cache->atlas_index];

    /* clear */
    SDL_FillRect(atlas_surface, NULL, 0);

    /* blit the texture surfaces onto the atlas */
    for (size_t i = 0; i < shlenu(cache->hash); ++i) {
        /* skip all that aren't part of currently built one */
        if (cache->hash[i].value.atlas_index != cache->atlas_index)
            continue;

        /* skip loners */
        if (cache->hash[i].value.loner_texture != 0)
            continue;

        SDL_BlitSurface(cache->hash[i].value.data,
                        NULL,
                        atlas_surface,
                        &(SDL_Rect){
                            .x = (int)cache->hash[i].value.srcrect.x,
                            .y = (int)cache->hash[i].value.srcrect.y,
                            .w = (int)cache->hash[i].value.srcrect.w,
                            .h = (int)cache->hash[i].value.srcrect.h,
                        });
    }

    /* texturize it! */
    upload_texture_from_surface(cache->atlas_textures[cache->atlas_index], atlas_surface);
}


/* uses the textures currently in the cache to create an array of stbrp_rects */
static stbrp_rect *create_rects_from_cache(TextureCache *cache) {
    stbrp_rect *rects = NULL;
    for (size_t i = 0; i < shlenu(cache->hash); ++i) {
        if (cache->hash[i].value.loner_texture != 0)
            continue;

        const SDL_Surface *surface_data = cache->hash[i].value.data;
        stbrp_rect new_rect = {
            .w = surface_data->w,
            .h = surface_data->h,
        };

        arrput(rects, new_rect);
    }
    return rects;
}


/* returns an array which contains a _copy_ of every unpacked rect in rects. */
/* each of these copies will have their original index in rects saved in */
/* their `id` field, which is an int. */
static stbrp_rect *filter_unpacked_rects(stbrp_rect *rects) {
    stbrp_rect *unpacked_rects = NULL;
    for (size_t i = 0; i < arrlenu(rects); ++i) {
        /* already packed */
        if (rects[i].was_packed)
            continue;

        arrput(unpacked_rects, rects[i]);
        /* stb_rect_pack mercifully gives you a free userdata int */
        /* the index is saved there so the original array can be updated later */
        unpacked_rects[arrlenu(unpacked_rects)-1].id = (int)i;
    }
    return unpacked_rects;
}


/* updates the original rects array with the data from packed_rects */
/* returns true if all rects were packed successfully */
static bool update_rects(TextureCache *cache, stbrp_rect *rects, stbrp_rect *packed_rects) {
    /* !!! do not grow either of the arrays !!! */
    /* the reallocation will try to reassign the array pointer, to no effect. */
    /* see stb_ds.h */
    bool packed_all = true;
    
    for (size_t i = 0; i < arrlenu(packed_rects); ++i) {
        /* we can check if any rects failed to be packed right here */
        /* it's not ideal, but it avoids another iteration */
        if (!packed_rects[i].was_packed) {
            packed_all = false;
            continue;
        }

        rects[packed_rects[i].id] = packed_rects[i];
        /* while the order of the elements in the hash map is unknown to us, */
        /* their equivalents in `rects` are in that same (unknown) order, which means */
        /* we can use the index we had saved to find the original texture struct */
        cache->hash[packed_rects[i].id].value.atlas_index = cache->atlas_index;
    }

    return packed_all;
}


/* updates the atlas location of every rect in the cache */
static void update_texture_rects_in_atlas(TextureCache *cache, stbrp_rect *rects) {
    int r = 0;
    for (size_t i = 0; i < shlenu(cache->hash); ++i) {
        if (cache->hash[i].value.loner_texture != 0)
            continue;

        cache->hash[i].value.srcrect = (Rect) {
            .x = (float)rects[r].x,
            .y = (float)rects[r].y,
            .w = (float)rects[r].w,
            .h = (float)rects[r].h,
        };

        r++;
    }
}


void textures_cache_init(TextureCache *cache, SDL_Window *window) {
    cache->window = window;
    sh_new_arena(cache->hash);

    cache->node_buffer = SDL_calloc(ctx.texture_atlas_size, sizeof *cache->node_buffer);

    add_new_atlas(cache);
    recreate_current_atlas_texture(cache);
}


void textures_cache_deinit(TextureCache *cache) {
    /* free atlas textures */
    for (size_t i = 0; i < arrlenu(cache->atlas_textures); ++i) {
        delete_gpu_texture(cache->atlas_textures[i]);
    }
    arrfree(cache->atlas_textures);

    /* free atlas surfaces */
    for (size_t i = 0; i < arrlenu(cache->atlas_surfaces); ++i) {
        SDL_FreeSurface(cache->atlas_surfaces[i]);
    }
    arrfree(cache->atlas_surfaces);

    /* free cache hashes */
    for (size_t i = 0; i < shlenu(cache->hash); ++i) {
        stbi_image_free(cache->hash[i].value.data->pixels);
        SDL_FreeSurface(cache->hash[i].value.data);
    }
    shfree(cache->hash);

    SDL_free(cache->node_buffer);
}


static enum TextureMode infer_texture_mode(SDL_Surface *surface) {
    const uint32_t amask = surface->format->Amask;
    if (amask == 0)
        return TEXTURE_MODE_OPAQUE;

    enum TextureMode result = TEXTURE_MODE_OPAQUE;

    SDL_LockSurface(surface);

    for (int i = 0; i < surface->w * surface->h; ++i) {
        /* TODO: don't assume 32 bit depth ? */
        Color color;
        SDL_GetRGBA(((uint32_t *)surface->pixels)[i], surface->format, &color.r, &color.g, &color.b, &color.a);

        if (color.a == 0)
            result = TEXTURE_MODE_SEETHROUGH;
        else if (color.a != 255) {
            result = TEXTURE_MODE_GHOSTLY;
            break;
        }
    }

    SDL_UnlockSurface(surface);

    return result;
}


static TextureKey textures_load(TextureCache *cache, const char *path) {
    /* no need to do anything if it was loaded already */
    const ptrdiff_t i = shgeti(cache->hash, path);
    if (i >= 0)
        return (TextureKey){ (uint16_t)i };

    SDL_Surface *surface = textures_load_surface(path);
    Texture new_texture = {
        .data = surface,
        .mode = infer_texture_mode(surface),
    };

    /* it's a "loner texture," it doesn't fit in an atlas so it's not in one */
    if (surface->w >= (int)ctx.texture_atlas_size || surface->h >= (int)ctx.texture_atlas_size) {
        if (ctx.game.debug) {
            if (surface->w > 2048 || surface->h > 2048)
                log_warn("Unportable texture dimensions for %s, use 2048x2048 at max", path);
            if (!is_power_of_two(surface->w) || !is_power_of_two(surface->h))
                log_warn("Unportable texture dimensions for %s, should be powers of 2", path);
        }
        new_texture.loner_texture = create_gpu_texture(TEXTURE_FILTER_NEAREAST, true);
        upload_texture_from_surface(new_texture.loner_texture, surface); 
        new_texture.srcrect = (Rect) { .w = (float)surface->w, .h = (float)surface->h };
    } else {
        /* will be fully populated as the atlas updates */
        new_texture.atlas_index = cache->atlas_index;
        cache->is_dirty = true;
    }
    shput(cache->hash, path, new_texture);
    return (TextureKey){ (uint16_t)shgeti(cache->hash, path) };
}


void textures_update_atlas(TextureCache *cache) {
    if (!cache->is_dirty)
        return;

    /* this function makes a lot more sense if you read stb_rect_pack.h */
    stbrp_context pack_ctx; /* target info */
    stbrp_init_target(&pack_ctx,
                      (int)ctx.texture_atlas_size,
                      (int)ctx.texture_atlas_size,
                      cache->node_buffer,
                      (int)ctx.texture_atlas_size);

    stbrp_rect *rects = create_rects_from_cache(cache);

    /* we have to keep packing, and creating atlases if necessary, */
    /* until all rects have been packed. */
    /* ideally, this will not iterate more than once. */
    bool textures_remaining = true;
    while (textures_remaining) {
        stbrp_rect *rects_to_pack = filter_unpacked_rects(rects);
        stbrp_pack_rects(&pack_ctx, rects_to_pack, (int)arrlen(rects_to_pack));

        textures_remaining = !update_rects(cache, rects, rects_to_pack);
        arrfree(rects_to_pack); /* got what we needed */

        /* some textures couldn't be packed */
        if (textures_remaining) {
            update_texture_rects_in_atlas(cache, rects);
            recreate_current_atlas_texture(cache);

            /* need a new atlas for next time */
            add_new_atlas(cache);
            ++cache->atlas_index;
        }
    }

    update_texture_rects_in_atlas(cache, rects);
    recreate_current_atlas_texture(cache);

    cache->is_dirty = false;

    arrfree(rects);
}

/* EXPERIMANTAL: LIKELY TO BE REMOVED! */
#if defined(__linux__) /* use rodata elf section for fast lookups of repeating textures */

#include "system/linux/twn_elf.h"

static const char *rodata_start;
static const char *rodata_stop;

static const char *last_path = NULL;
static TextureKey last_texture;
static struct PtrToTexture {
    const void *key;
    TextureKey value;
} *ptr_to_texture;

/* TODO: separate and reuse */
TextureKey textures_get_key(TextureCache *cache, const char *path) {
    if (rodata_stop == NULL)
        if (!infer_elf_section_bounds(".rodata", &rodata_start, &rodata_stop))
            CRY("Section inference", ".rodata section lookup failed");

    /* the fastest path */
    if (path == last_path)
        return last_texture;
    else {
        /* moderately fast path, by pointer hashing */
        const ptrdiff_t texture = hmgeti(ptr_to_texture, path);
        if (texture != -1) {
            if (path >= rodata_start && path < rodata_stop)
                last_path = path;
            last_texture = ptr_to_texture[texture].value;
            return last_texture;
        }
    }

    /* try loading */
    last_texture = textures_load(cache, path);
    hmput(ptr_to_texture, path, last_texture);

    if (path >= rodata_start && path < rodata_stop)
        last_path = path;

    return last_texture;
}

#else
TextureKey textures_get_key(TextureCache *cache, const char *path) {
    /* hash tables are assumed to be stable, so we just return indices */
    const ptrdiff_t texture = shgeti(cache->hash, path);

    /* load it if it isn't */
    if (texture == -1) {
        return textures_load(cache, path);
    } else
        return (TextureKey){ (uint16_t)texture };
}

#endif /* generic implementation of textures_get_key() */

int32_t textures_get_atlas_id(const TextureCache *cache, TextureKey key) {
    if (m_texture_key_is_valid(key)) {
        if (cache->hash[key.id].value.loner_texture != 0)
            return -cache->hash[key.id].value.loner_texture;
        else
            return cache->hash[key.id].value.atlas_index;
    } else {
        CRY("Texture lookup failed.",
            "Tried to get atlas id that isn't loaded.");
        return 0;
    }
}

Rect textures_get_srcrect(const TextureCache *cache, TextureKey key) {
    if (m_texture_key_is_valid(key)) {
        return cache->hash[key.id].value.srcrect;
    } else {
        CRY("Texture lookup failed.",
            "Tried to get texture that isn't loaded.");
        return (Rect){ 0, 0, 0, 0 };
    }
}


Rect textures_get_dims(const TextureCache *cache, TextureKey key) {
    if (m_texture_key_is_valid(key)) {
        if (cache->hash[key.id].value.loner_texture != 0)
            return cache->hash[key.id].value.srcrect;
        else
            return (Rect){ .w = (float)ctx.texture_atlas_size, .h = (float)ctx.texture_atlas_size };
    } else {
        CRY("Texture lookup failed.",
            "Tried to get texture that isn't loaded.");
        return (Rect){ 0, 0, 0, 0 };
    }
}


void textures_bind(const TextureCache *cache, TextureKey key) {
    if (m_texture_key_is_valid(key)) {
        if (cache->hash[key.id].value.loner_texture == 0)
            bind_gpu_texture(cache->atlas_textures[cache->hash[key.id].value.atlas_index]);
        else
            bind_gpu_texture(cache->hash[key.id].value.loner_texture);
    } else if (key.id == 0) {
        CRY("Texture binding failed.",
            "Tried to get texture that isn't loaded.");
    }
}


/* TODO: alternative schemes, such as: array texture, fragment shader and geometry division */
void textures_bind_repeating(const TextureCache *cache, TextureKey key) {
    if (m_texture_key_is_valid(key)) {
        if (cache->hash[key.id].value.loner_texture == 0) {

            /* already allocated */
            if (cache->hash[key.id].value.repeating_texture != 0) {
                bind_gpu_texture(cache->hash[key.id].value.repeating_texture);
                return;
            }

            const Texture texture = cache->hash[key.id].value;

            const GPUTexture repeating_texture = create_gpu_texture(TEXTURE_FILTER_NEAREAST, false);

            SDL_LockSurface(texture.data);

            upload_gpu_texture(repeating_texture,
                               texture.data->pixels,
                               4,
                               texture.data->w,
                               texture.data->h);

            SDL_UnlockSurface(texture.data);

            cache->hash[key.id].value.repeating_texture = repeating_texture;

            bind_gpu_texture(repeating_texture);

        } else
            bind_gpu_texture(cache->hash[key.id].value.loner_texture);

    } else if (key.id == 0) {
        CRY("Texture binding failed.",
            "Tried to get texture that isn't loaded.");
    }
}


TextureMode textures_get_mode(const TextureCache *cache, TextureKey key) {
    if (m_texture_key_is_valid(key)) {
        return cache->hash[key.id].value.mode;
    } else {
        CRY("Texture binding failed.",
            "Tried to get texture that isn't loaded.");
        return TEXTURE_MODE_GHOSTLY;
    }
}


size_t textures_get_num_atlases(const TextureCache *cache) {
    return cache->atlas_index + 1;    
}

void textures_reset_state(void) {
#if defined(__linux__) && !defined(HOT_RELOAD_SUPPORT)
    last_path = NULL;
    last_texture = (TextureKey){0};
    shfree(ptr_to_texture);

#endif
}