#include "Rendering/Texture.h"

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

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

bool texture_collection_init(uint16_t size, texture_collection_t* textures)
{
    texture_collection_t temp = {0};
    temp.buffer = (texture_t*)malloc(size * sizeof(texture_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)
{
    texture_t* temp = (texture_t*)realloc(textures->buffer, size * sizeof(texture_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]);
    }

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

texture_entity_t texture_load(const char* filename, bool srgb, texture_collection_t* textures)
{
    int width, height, channels;
    uint8_t* data = stbi_load(filename, &width, &height, &channels, 0);
    if (data == NULL)
    {
        return invalid_texture_entity();
    }

    if (srgb)
    {
        // Convert to linear space if the texture is in sRGB format
        for (int i = 0; i < width * height * channels; i++)
        {
            data[i] = (uint8_t)(powf(data[i] / 255.0f, 2.2f) * 255.0f);
        }
    }

    texture_t texture = {0};
    texture.width = (uint32_t)width;
    texture.height = (uint32_t)height;
    texture.channel_count = (uint8_t)channels;
    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;
    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;
    }
}

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

    uint8_t* pixel = &texture->data[pixel_index * texture->channel_count];
    return (vec4s)
    {
        GET_CHANNEL_DATA(pixel, 0, texture->channel_count, 0),
        GET_CHANNEL_DATA(pixel, 1, texture->channel_count, 0),
        GET_CHANNEL_DATA(pixel, 2, texture->channel_count, 0),
        GET_CHANNEL_DATA(pixel, 3, texture->channel_count, 1)
    };
}

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 get_pixel_color(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 = get_pixel_color(texture, x0, y0);
    vec4s c10 = get_pixel_color(texture, x1, y0);
    vec4s c01 = get_pixel_color(texture, x0, y1);
    vec4s c11 = get_pixel_color(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);
    }
}