#include "private/textures.h"
#include "config.h"
#include "util.h"
#include "textures.h"

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

#include <stdbool.h>
#include <stdlib.h>
#include <stdio.h>


static SDL_Surface *image_to_surface(const char *path) {
    SDL_RWops *handle = PHYSFSRWOPS_openRead(path);
    if (handle == NULL)
        goto fail;
    
    SDL_Surface *result = IMG_Load_RW(handle, true);
    if (result == NULL)
        goto fail;

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

    return result;

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


static GLuint new_gl_texture(void) {
    GLuint texture;
    glGenTextures(1, &texture);
    glBindTexture(GL_TEXTURE_2D, texture);

    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S,     GL_REPEAT);
    glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T,     GL_REPEAT);
    glTexParameteri(GL_TEXTURE_2D, GL_GENERATE_MIPMAP,    GL_TRUE);

    glTexEnvi(GL_TEXTURE_ENV, GL_TEXTURE_ENV_MODE, GL_MODULATE);

    glBindTexture(GL_TEXTURE_2D, 0);
    return texture;
}


/* adds a new, blank atlas surface to the cache */
static void add_new_atlas(struct texture_cache *cache) {
    SDL_PixelFormat *native_format =
        SDL_AllocFormat(SDL_GetWindowPixelFormat(cache->window));

    /* the window format won't have an alpha channel, so we figure this out */
#if SDL_BYTEORDER == SDL_BIG_ENDIAN
    uint32_t a_mask = 0x000000FF;
#else 
    uint32_t a_mask = 0xFF000000;
#endif

    SDL_Surface *new_atlas = SDL_CreateRGBSurface(0,
                                                  TEXTURE_ATLAS_SIZE,
                                                  TEXTURE_ATLAS_SIZE,
                                                  TEXTURE_ATLAS_BIT_DEPTH,
                                                  native_format->Rmask,
                                                  native_format->Gmask,
                                                  native_format->Bmask,
                                                  a_mask);
    SDL_FreeFormat(native_format);

    SDL_SetSurfaceRLE(new_atlas,  true);
    arrput(cache->atlas_surfaces, new_atlas);
    arrput(cache->atlas_textures, new_gl_texture());
}


static void upload_texture_from_surface(GLuint texture, SDL_Surface *surface) {
    Uint32 rmask, gmask, bmask, amask;

    glBindTexture(GL_TEXTURE_2D, texture);

    // glPixelStorei(GL_PACK_ALIGNMENT, 1);

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

    /* TODO: don't do it if format is compatible */
    SDL_Surface* intermediate = SDL_CreateRGBSurface(0,
        surface->w, surface->h, 32, rmask, gmask, bmask, amask);
    SDL_BlitSurface(surface, NULL, intermediate, NULL);
    SDL_LockSurface(intermediate);

    glTexImage2D(GL_TEXTURE_2D,
                 0,
                 GL_RGBA,
                 surface->w,
                 surface->h,
                 0,
                 GL_RGBA,
                 GL_UNSIGNED_BYTE,
                 intermediate->pixels);

    SDL_UnlockSurface(intermediate);
    SDL_FreeSurface(intermediate);

    glBindTexture(GL_TEXTURE_2D, 0);
}


static void recreate_current_atlas_texture(struct texture_cache *cache) {
    /* TODO: figure out if SDL_UpdateTexture alone is faster than blitting */
    /* TODO: should surfaces be freed after they cannot be referenced in atlas builing? */
    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 = cache->hash[i].value.srcrect.x,
                            .y = cache->hash[i].value.srcrect.y,
                            .w = cache->hash[i].value.srcrect.w,
                            .h = 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(struct texture_cache *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(struct texture_cache *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(struct texture_cache *cache, stbrp_rect *rects) {
    for (size_t i = 0; i < arrlenu(rects); ++i) {
        cache->hash[i].value.srcrect = (t_rect) {
            .x = rects[i].x,
            .y = rects[i].y,
            .w = rects[i].w,
            .h = rects[i].h,
        };
    }
}


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

    cache->node_buffer = cmalloc(sizeof *cache->node_buffer * TEXTURE_ATLAS_SIZE);

    add_new_atlas(cache);
    recreate_current_atlas_texture(cache);
}


void textures_cache_deinit(struct texture_cache *cache) {
    /* free atlas textures */
    for (size_t i = 0; i < arrlenu(cache->atlas_textures); ++i) {
        glDeleteTextures(1, &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) {
        SDL_FreeSurface(cache->hash[i].value.data);
    }
    shfree(cache->hash);

    free(cache->node_buffer);
}


void textures_dump_atlases(struct texture_cache *cache) {
    PHYSFS_mkdir("/dump");

    const char string_template[] = "/dump/atlas%zd.png";
    char buf[2048]; /* larger than will ever be necessary */

    size_t i = 0;
    for (; i < arrlenu(cache->atlas_surfaces); ++i) {
        snprintf(buf, sizeof buf, string_template, i);

        SDL_RWops *handle = PHYSFSRWOPS_openWrite(buf);

        if (handle == NULL) {
            CRY("Texture atlas dump failed.", "File could not be opened");
            return;
        }

        IMG_SavePNG_RW(cache->atlas_surfaces[i], handle, true);
        log_info("Dumped atlas %s", buf);
    }
}


static t_texture_key textures_load(struct texture_cache *cache, const char *path) {
    /* no need to do anything if it was loaded already */
    if (shgeti(cache->hash, path) >= 0)
        return (t_texture_key){0};

    SDL_Surface *surface = image_to_surface(path);
    struct texture new_texture = {0};
    new_texture.data = surface;

    /* it's a "loner texture," it doesn't fit in an atlas so it's not in one */
    if (surface->w >= TEXTURE_ATLAS_SIZE || surface->h >= TEXTURE_ATLAS_SIZE) {
        new_texture.loner_texture = new_gl_texture();
        upload_texture_from_surface(new_texture.loner_texture, surface); 
        new_texture.srcrect = (t_rect) { .w = surface->w, .h = surface->h };
        shput(cache->hash, path, new_texture);
        return (t_texture_key){ (int)shgeti(cache->hash, path) };
    } else {
        new_texture.atlas_index = cache->atlas_index;
        shput(cache->hash, path, new_texture);
        cache->is_dirty = true;
        return (t_texture_key){ (int)shgeti(cache->hash, path) };
    }
}


void textures_update_atlas(struct texture_cache *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,
                      TEXTURE_ATLAS_SIZE,
                      TEXTURE_ATLAS_SIZE,
                      cache->node_buffer,
                      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! */
/* todo: If it's proven to be useful: add runtime checking for .rodata > .data */
#ifdef __unix__ /* use rodata elf section for fast lookups of repeating textures */

extern const char start_rodata_address[];
extern const char stop_rodata_heuristic[];

asm(".set start_rodata_address, .rodata");
asm(".set stop_rodata_heuristic, .data"); /* there's nothing in default linker script to know the size of .rodata */

t_texture_key textures_get_key(struct texture_cache *cache, const char *path) {
    static const char *last_path = NULL;
    static t_texture_key last_texture;

    /* fast path */
    if (path == last_path && path >= start_rodata_address && path < stop_rodata_heuristic)
        return last_texture;

    /* hash tables are assumed to be stable, so we just return indices */
    int texture = (int)shgeti(cache->hash, path);

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

    last_path = path;

    return last_texture;
}

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

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

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

t_rect textures_get_srcrect(const struct texture_cache *cache, t_texture_key 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 (t_rect){ 0, 0, 0, 0 };
    }
}


t_rect textures_get_dims(const struct texture_cache *cache, t_texture_key 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 (t_rect){ .w = TEXTURE_ATLAS_SIZE, .h = TEXTURE_ATLAS_SIZE };
    } else {
        CRY("Texture lookup failed.",
            "Tried to get texture that isn't loaded.");
        return (t_rect){ 0, 0, 0, 0 };
    }
}


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


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