#ifndef MATERIAL_H
#define MATERIAL_H

#include "Common.h"

typedef struct
{
    vec3s normal;
    vec3s position;
    vec3s wi;
    vec3s wo;
}shading_context_t;

typedef vec3s (*sample_bsdf_f)(const void* data, const vec3s normal, const vec3s wo, sobol_state_t* sobol_state, float* pdf_out);
typedef float (*sample_bsdf_pdf_f)(const void* data, const vec3s normal, const vec3s wo, const vec3s wi);
typedef vec3s (*evaluate_bsdf_f)(const shading_context_t* context, const void* data);

typedef struct
{
    // TODO: alpha, displacement, etc.
    vec3s emission; // We have to have emission outside of data because data is only for bsdf, and emission is not part of bsdf. Maybe another shading properties struct is needed if the number of properties increases.
    sample_bsdf_f sample_bsdf;
    sample_bsdf_pdf_f sample_bsdf_pdf;
    evaluate_bsdf_f evaluate_bsdf;
    void* data;
}material_t;

typedef struct
{
    uint8_t id;
} material_entity_t;


//TODO: Handle material remove, we can use something like sparse set.
typedef struct
{
    uint8_t count;
    uint8_t size;
    material_t* buffer;
} material_collection_t;

inline material_collection_t material_collection_create(size_t size)
{
    if (size > 254)
    {
        size = 254; // Limit the count to 254 to fit in a uint8_t
    }

    material_collection_t collection = {0};
    collection.buffer = (material_t*)malloc(size * sizeof(material_t));
    collection.size = (uint8_t)size;
    return collection;
}

inline void material_collection_resize(material_collection_t* materials, size_t size)
{
    if (size > 254)
    {
        size = 254; // Limit the count to 254 to fit in a uint8_t
    }

    material_t* temp = (material_t*)realloc(materials->buffer, size * sizeof(material_t));
    if (temp != NULL)
    {
        materials->buffer = temp;
        materials->size = (uint8_t)size;
    }
}

inline void material_collection_free(material_collection_t* materials)
{
    if (materials->buffer != NULL)
    {
        free(materials->buffer);
        materials->buffer = NULL;
    }
}

inline material_entity_t material_create(const sample_bsdf_f sample, const sample_bsdf_pdf_f sample_pdf, const evaluate_bsdf_f evaluate, void* data, material_collection_t* collection)
{
    material_t material = {0};

    if (collection->count >= collection->size)
    {
        material_collection_resize(collection, collection->size * 2);
    }

    material.sample_bsdf = sample;
    material.sample_bsdf_pdf = sample_pdf;
    material.evaluate_bsdf = evaluate;
    material.data = data;

    material_entity_t entity = {.id = collection->count};

    collection->buffer[collection->count] = material;
    collection->count++;

    return entity;
}

inline vec3s sample_material_bsdf(const material_t* material, const vec3s normal, const vec3s wo, sobol_state_t* sobol_state, float* pdf_out)
{
    vec3s wi = glms_vec3_zero();
    if (material->sample_bsdf != NULL)
    {
        wi = material->sample_bsdf(material->data, normal, wo, sobol_state, pdf_out);
    }

    return wi;
}

inline float sample_material_bsdf_pdf(const material_t* material, const vec3s normal, const vec3s wo, const vec3s wi)
{
    float pdf = 0.0f;
    if (material->sample_bsdf_pdf != NULL)
    {
        pdf = material->sample_bsdf_pdf(material->data, normal, wo, wi);
    }

    return pdf;
}

inline vec3s evaluate_material_bsdf(const material_t* material, const shading_context_t* context)
{
    vec3s bsdf_color = glms_vec3_zero();
    if (material->evaluate_bsdf != NULL)
    {
        bsdf_color = material->evaluate_bsdf(context, material->data);
    }

    return bsdf_color;
}

#endif // MATERIAL_H