#include "Rendering/Texture.h"
#include "Common/String.h"

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

#define GET_CHANNEL_DATA(pixel, channel, channel_count, default, max) (channel < channel_count ? pixel[channel] : default) / max

bool texture_collection_init(uint16_t size, texture_collection_t* textures)
{
    texture_collection_t temp = {0};
    temp.buffer = (texture_asset_t*)malloc(size * sizeof(texture_asset_t));
    if (temp.buffer == NULL)
    {
        return false;
    }

    temp.size = size;
    temp.count = 0;
    *textures = temp;

    return true;
}

void texture_collection_resize(texture_collection_t* textures, uint16_t size)
{
    if (size == INVALID_TEXTURE_ID)
    {
        size = INVALID_TEXTURE_ID - 1;
    }

    if (size == textures->size)
    {
        return;
    }

    texture_asset_t* temp = (texture_asset_t*)realloc(textures->buffer, size * sizeof(texture_asset_t));
    if (temp != NULL)
    {
        textures->buffer = temp;
        textures->size = size;
    }
}

void texture_collection_free(texture_collection_t* textures)
{
    if (textures == NULL)
    {
        return;
    }

    for (uint16_t i = 0; i < textures->count; i++)
    {
        texture_free(&textures->buffer[i].texture);

        char* full_name = textures->buffer[i].full_name;
        if (full_name != NULL)
        {
            free(full_name);
        }
    }

    free(textures->buffer);
    textures->buffer = NULL;
}

texture_entity_t texture_load(const char* filename, bool srgb, stride_t stride, texture_collection_t* textures)
{
    // TODO: This hurts performance, consider using a hash map or similar structure for faster lookups

    // for (uint16_t i = 0; i < textures->count; i++)
    // {
    //     if (strcmp(textures->buffer[i].full_name, filename) == 0)
    //     {
    //         return (texture_entity_t){.id = i};
    //     }
    // }

    int width, height, channels;
    char* data = NULL;

    switch (stride)
    {
        case UINT_8:
            uint8_t* img_data = stbi_load(filename, &width, &height, &channels, 0);
            if (srgb)
            {
                // Convert to linear space if the texture is in sRGB format
                for (int i = 0; i < width * height * channels; i++)
                {
                    img_data[i] = (uint8_t)(powf(img_data[i] / 255.0f, 2.2f) * 255.0f);
                }
            }
            data = (char*)img_data;
            break;

        case UINT_16:
            data = (char*)stbi_load_16(filename, &width, &height, &channels, 0);
            break;
        case FLOAT_32:
            data = (char*)stbi_loadf(filename, &width, &height, &channels, 0);
            break;
    }

    if (data == NULL)
    {
        return invalid_texture_entity();
    }

    texture_t texture = {0};
    texture.width = (uint32_t)width;
    texture.height = (uint32_t)height;
    texture.channel_count = (uint8_t)channels;
    texture.stride = stride;
    texture.data = data;

    texture.wrap_mode = REPEAT;
    texture.filter_mode = LINEAR;

    if (textures->count >= textures->size)
    {
        texture_collection_resize(textures, textures->size * 2);
    }

    texture_entity_t entity = {.id = textures->count};

    textures->buffer[textures->count] = (texture_asset_t){.full_name = string_copy(filename), .texture = texture};
    textures->count++;

    return entity;
}

static inline void warp_uv(wrap_mode_t mode, float* u, float* v)
{
    switch (mode)
    {
        case REPEAT:
            *u = fmodf(fabsf(*u), 1.0f);
            *v = fmodf(fabsf(*v), 1.0f);
            break;
        case CLAMP:
            *u = fminf(fmaxf(*u, 0.0f), 1.0f);
            *v = fminf(fmaxf(*v, 0.0f), 1.0f);
            break;
    }
}

vec4s texture_get_pixel(const texture_t* texture, uint32_t x, uint32_t y)
{
    if (x >= texture->width || y >= texture->height)
    {
        return (vec4s){0.0f, 0.0f, 0.0f, 1.0f};
    }

    uint32_t pixel_index = y * texture->width + x;
    size_t pixel_offset = (size_t)pixel_index * texture->channel_count * texture->stride;

    uint8_t* base8 = (uint8_t*)texture->data + pixel_offset;
    uint16_t* base16 = (uint16_t*)(texture->data) + ((y * texture->width + x) * texture->channel_count);
    float* base32 = (float*)(texture->data) + ((y * texture->width + x) * texture->channel_count);

    vec4s out = {0, 0, 0, 1}; // default alpha = 1

    for (int c = 0; c < texture->channel_count && c < 4; ++c)
    {
        float value = 0.0f;
        switch (texture->stride)
        {
            case UINT_8:
                value = base8[c] / 255.0f;
                break;
            case UINT_16:
                value = base16[c] / 65535.0f;
                break;
            case FLOAT_32:
                value = base32[c];
                break;
            default:
                value = (c == 3) ? 1.0f : 0.0f;
                break;
        }

        out.raw[c] = value;
    }

    return out;
}

static vec4s nearest_filter(const texture_t* texture, float u, float v)
{
    uint32_t x = (uint32_t)floorf(u * (texture->width  - 1));
    uint32_t y = (uint32_t)floorf(v * (texture->height - 1));

    x = x < texture->width  ? x : texture->width  - 1;
    y = y < texture->height ? y : texture->height - 1;

    return texture_get_pixel(texture, x, y);
}

static vec4s linear_filter(const texture_t* texture, float u, float v)
{
    float x = u * (texture->width  - 1);
    float y = v * (texture->height - 1);

    uint32_t x0 = (uint32_t)floorf(x);
    uint32_t x1 = x0 + 1;
    uint32_t y0 = (uint32_t)floorf(y);
    uint32_t y1 = y0 + 1;

    float sx = x - (float)x0;
    float sy = y - (float)y0;

    // Clamp to edges
    x0 = x0 < texture->width  ? x0 : texture->width  - 1;
    x1 = x1 < texture->width  ? x1 : texture->width  - 1;
    y0 = y0 < texture->height ? y0 : texture->height - 1;
    y1 = y1 < texture->height ? y1 : texture->height - 1;

    // Sample 4 texels
    vec4s c00 = texture_get_pixel(texture, x0, y0);
    vec4s c10 = texture_get_pixel(texture, x1, y0);
    vec4s c01 = texture_get_pixel(texture, x0, y1);
    vec4s c11 = texture_get_pixel(texture, x1, y1);

    // Interpolate along x
    vec4s c0 = glms_vec4_lerp(c00, c10, sx);
    vec4s c1 = glms_vec4_lerp(c01, c11, sx);

    // Interpolate along y
    vec4s result = glms_vec4_lerp(c0, c1, sy);

    return result;
}

static inline vec4s filter_texture(const texture_t* texture, float u, float v)
{
    switch (texture->filter_mode)
    {
        case NEAREST:
            return nearest_filter(texture, u, v);
        case LINEAR:
            return linear_filter(texture, u, v);
        default:
            return (vec4s){0.0f, 0.0f, 0.0f, 1.0f};
    }
}

vec4s texture_sample(const texture_t* texture, float u, float v)
{
    warp_uv(texture->wrap_mode, &u, &v);
    return filter_texture(texture, u, v);
}

void texture_free(texture_t* texture)
{
    if (texture != NULL && texture->data != NULL)
    {
        stbi_image_free(texture->data);
    }
}