#include "Rendering/Scene.h"

#include <stdlib.h>
#include <string.h>

static inline void mesh_model_collection_init(mesh_model_collection_t* models, uint32_t capacity)
{
    models->count = capacity;
    models->capacity = capacity;
    models->buffer = (mesh_model_t*)calloc(capacity, sizeof(mesh_model_t));
}

static inline void mesh_instance_collection_init(mesh_instance_collection_t* instances, uint32_t capacity)
{
    instances->count = capacity;
    instances->capacity = capacity;
    instances->buffer = (mesh_instance_t*)calloc(capacity, sizeof(mesh_instance_t));
}

static inline void mesh_model_collection_free(mesh_model_collection_t* models)
{
    if (models == NULL || models->buffer == NULL)
    {
        return;
    }

    for (uint32_t i = 0; i < models->capacity; ++i)
    {
        mesh_model_t* m = &models->buffer[i];
        if (m->active)
        {
            bvh_tree_free(&m->blas);
            triangle_collection_free(&m->triangles);
        }
    }

    free(models->buffer);
    models->buffer = NULL;
    models->count = 0;
    models->capacity = 0;
}

static inline void mesh_instance_collection_free(mesh_instance_collection_t* instances)
{
    if (instances == NULL)
    {
        return;
    }

    free(instances->buffer);
    instances->buffer = NULL;
    instances->count = 0;
    instances->capacity = 0;
}

static inline vec3s mat4_mul_point(mat4s m, vec3s p)
{
    return glms_mat4_mulv3(m, p, 1.0f);
}

static inline aabb_t aabb_transform(mat4s m, aabb_t aabb)
{
    // Transform 8 corners and compute bounds.
    vec3s c000 = mat4_mul_point(m, (vec3s){aabb.min.x, aabb.min.y, aabb.min.z});
    vec3s c001 = mat4_mul_point(m, (vec3s){aabb.min.x, aabb.min.y, aabb.max.z});
    vec3s c010 = mat4_mul_point(m, (vec3s){aabb.min.x, aabb.max.y, aabb.min.z});
    vec3s c011 = mat4_mul_point(m, (vec3s){aabb.min.x, aabb.max.y, aabb.max.z});
    vec3s c100 = mat4_mul_point(m, (vec3s){aabb.max.x, aabb.min.y, aabb.min.z});
    vec3s c101 = mat4_mul_point(m, (vec3s){aabb.max.x, aabb.min.y, aabb.max.z});
    vec3s c110 = mat4_mul_point(m, (vec3s){aabb.max.x, aabb.max.y, aabb.min.z});
    vec3s c111 = mat4_mul_point(m, (vec3s){aabb.max.x, aabb.max.y, aabb.max.z});

    aabb_t out = invalid_aabb();
    aabb_growth(&out, c000);
    aabb_growth(&out, c001);
    aabb_growth(&out, c010);
    aabb_growth(&out, c011);
    aabb_growth(&out, c100);
    aabb_growth(&out, c101);
    aabb_growth(&out, c110);
    aabb_growth(&out, c111);
    return out;
}

static inline mat3s compute_normal_matrix(mat4s local_to_world)
{
    // normalMatrix = transpose(inverse(mat3(local_to_world)))
    mat4s inv = glms_mat4_inv(local_to_world);
    mat3s m3 = glms_mat4_pick3(inv);
    return glms_mat3_transpose(m3);
}

static bool scene_rebuild_tlas(scene_t* scene)
{
    if (scene == NULL)
    {
        return false;
    }

    // Build list of active instances.
    uint64_t active_count = 0;
    for (uint32_t i = 0; i < scene->mesh_instances.capacity; ++i)
    {
        if (scene->mesh_instances.buffer[i].active)
        {
            active_count++;
        }
    }

    if (active_count == 0)
    {
        tlas_tree_free(&scene->tlas);
        scene->tlas_dirty = false;
        return true;
    }

    uint64_t* indices = (uint64_t*)malloc(sizeof(uint64_t) * active_count);
    if (indices == NULL)
    {
        return false;
    }

    uint64_t cursor = 0;
    for (uint32_t i = 0; i < scene->mesh_instances.capacity; ++i)
    {
        if (scene->mesh_instances.buffer[i].active)
        {
            indices[cursor++] = i;
        }
    }

    // Build an array of bounds for all instances (indexed by instance_id).
    // TLAS references this via primitive indices.
    // We can pass the backing array directly.
    aabb_t* bounds = (aabb_t*)malloc(sizeof(aabb_t) * scene->mesh_instances.capacity);
    if (bounds == NULL)
    {
        free(indices);
        return false;
    }
    for (uint32_t i = 0; i < scene->mesh_instances.capacity; ++i)
    {
        bounds[i] = scene->mesh_instances.buffer[i].world_bounds;
    }

    // Store bounds pointer via tlas->instance_bounds; Scene owns this allocation.
    // For simplicity, reuse the buffer by freeing previous and storing on scene.
    // (We attach it to tlas.instance_bounds and free it in scene_free via tlas_tree_free doesn't free it.)
    // We'll keep it alive by storing in a static on scene via tlas.instance_bounds.
    // TLAS builder does not take ownership; we manage it here.
    const aabb_t* old_bounds = scene->tlas.instance_bounds;

    bool ok = tlas_tree_build(&scene->tlas, indices, active_count, bounds);
    free(indices);
    if (!ok)
    {
        free(bounds);
        return false;
    }

    // Free old bounds after successful rebuild.
    free((void*)old_bounds);

    scene->tlas_dirty = false;
    return true;
}

bool scene_init(scene_t* scene, uint64_t triangle_count, uint16_t texture_count, uint8_t material_count, uint32_t punctual_light_count)
{
    scene_t temp = {0};

    if (!triangle_collection_init(triangle_count, &temp.triangles))
    {
        goto triangle_failed;
    }

    if (!texture_collection_init(texture_count, &temp.textures))
    {
        goto texture_failed;
    }

    if (!material_collection_init(material_count, &temp.materials))
    {
        goto material_failed;
    }

    if (!light_collection_create(punctual_light_count, 16, &temp.lights)) // NOTE: We just fixed the max directional light count to 16.
    {
        goto light_failed;
    }

    temp.camera = camera_create(
        (vec3s){0.0f, 0.0f, 5.0f},
        glms_quat_identity(),
        0.025f,
        0.036f,
        16.0f / 9.0f
    );

    // New mesh system: start with small default capacities (simple first).
    (void)triangle_count;
    mesh_model_collection_init(&temp.mesh_models, 64);
    mesh_instance_collection_init(&temp.mesh_instances, 128);
    temp.tlas_dirty = true;

    *scene = temp;
    return true;

light_failed:
    material_collection_free(&temp.materials);
material_failed:
    texture_collection_free(&temp.textures);
texture_failed:
    triangle_collection_free(&temp.triangles);
triangle_failed:
    return false;
}

bool scene_build_bvh(scene_t* scene)
{
    if (scene == NULL)
    {
        return false;
    }

    // Prefer TLAS if any mesh instances exist.
    if (scene->tlas_dirty)
    {
        if (!scene_rebuild_tlas(scene))
        {
            return false;
        }
    }

    // Legacy BVH build if triangles are present.
    if (scene->triangles.count > 0)
    {
        bvh_tree_free(&scene->bvh_tree);

        bvh_tree_t bvh_tree = {0};
        if (!bvh_tree_init(&bvh_tree, &scene->triangles))
        {
            return false;
        }

        bvh_tree_build(&bvh_tree);
        scene->bvh_tree = bvh_tree;
    }

    return true;
}

void scene_free(scene_t* scene)
{
    if (scene == NULL)
    {
        return;
    }

    bvh_tree_free(&scene->bvh_tree);
    triangle_collection_free(&scene->triangles);

    // Mesh system
    tlas_tree_free(&scene->tlas);
    free((void*)scene->tlas.instance_bounds);
    scene->tlas.instance_bounds = NULL;
    mesh_instance_collection_free(&scene->mesh_instances);
    mesh_model_collection_free(&scene->mesh_models);
    texture_collection_free(&scene->textures);
    material_collection_free(&scene->materials);
    light_collection_free(&scene->lights);
}

static uint32_t find_free_mesh_model_slot(scene_t* scene)
{
    for (uint32_t i = 0; i < scene->mesh_models.capacity; ++i)
    {
        if (!scene->mesh_models.buffer[i].active)
        {
            return i;
        }
    }

    uint32_t old_capacity = scene->mesh_models.capacity;
    uint32_t new_capacity = old_capacity == 0 ? 64 : old_capacity * 2;
    mesh_model_t* resized = (mesh_model_t*)realloc(scene->mesh_models.buffer, sizeof(mesh_model_t) * new_capacity);
    if (resized == NULL)
    {
        return UINT32_MAX;
    }
    memset(resized + old_capacity, 0, sizeof(mesh_model_t) * (new_capacity - old_capacity));
    scene->mesh_models.buffer = resized;
    scene->mesh_models.capacity = new_capacity;
    scene->mesh_models.count = new_capacity;

    return old_capacity;
}

static uint32_t find_free_mesh_instance_slot(scene_t* scene)
{
    for (uint32_t i = 0; i < scene->mesh_instances.capacity; ++i)
    {
        if (!scene->mesh_instances.buffer[i].active)
        {
            return i;
        }
    }

    uint32_t old_capacity = scene->mesh_instances.capacity;
    uint32_t new_capacity = old_capacity == 0 ? 128 : old_capacity * 2;
    mesh_instance_t* resized = (mesh_instance_t*)realloc(scene->mesh_instances.buffer, sizeof(mesh_instance_t) * new_capacity);
    if (resized == NULL)
    {
        return UINT32_MAX;
    }
    memset(resized + old_capacity, 0, sizeof(mesh_instance_t) * (new_capacity - old_capacity));
    scene->mesh_instances.buffer = resized;
    scene->mesh_instances.capacity = new_capacity;
    scene->mesh_instances.count = new_capacity;

    return old_capacity;
}

uint32_t scene_add_mesh_model(scene_t* scene, uint64_t triangle_reserve)
{
    if (scene == NULL)
    {
        return UINT32_MAX;
    }

    uint32_t slot = find_free_mesh_model_slot(scene);
    if (slot == UINT32_MAX)
    {
        return UINT32_MAX;
    }

    mesh_model_t* model = &scene->mesh_models.buffer[slot];
    *model = (mesh_model_t){0};
    model->active = true;
    model->local_bounds = invalid_aabb();

    if (!triangle_collection_init((size_t)(triangle_reserve > 0 ? triangle_reserve : 1), &model->triangles))
    {
        model->active = false;
        return UINT32_MAX;
    }

    return slot;
}

uint32_t scene_add_mesh_instance(scene_t* scene, uint32_t model_id, mat4s local_to_world)
{
    if (scene == NULL || model_id >= scene->mesh_models.capacity || !scene->mesh_models.buffer[model_id].active)
    {
        return UINT32_MAX;
    }

    uint32_t slot = find_free_mesh_instance_slot(scene);
    if (slot == UINT32_MAX)
    {
        return UINT32_MAX;
    }

    mesh_instance_t* inst = &scene->mesh_instances.buffer[slot];
    *inst = (mesh_instance_t){0};
    inst->active = true;
    inst->model_id = model_id;
    inst->local_to_world = local_to_world;
    inst->world_to_local = glms_mat4_inv(local_to_world);
    inst->normal_matrix = compute_normal_matrix(local_to_world);
    inst->world_bounds = aabb_transform(local_to_world, scene->mesh_models.buffer[model_id].local_bounds);

    scene->tlas_dirty = true;
    return slot;
}

void scene_remove_mesh_instance(scene_t* scene, uint32_t instance_id)
{
    if (scene == NULL || instance_id >= scene->mesh_instances.capacity)
    {
        return;
    }

    if (!scene->mesh_instances.buffer[instance_id].active)
    {
        return;
    }

    scene->mesh_instances.buffer[instance_id].active = false;
    scene->tlas_dirty = true;
}

void scene_set_mesh_instance_transform(scene_t* scene, uint32_t instance_id, mat4s local_to_world)
{
    if (scene == NULL || instance_id >= scene->mesh_instances.capacity)
    {
        return;
    }

    mesh_instance_t* inst = &scene->mesh_instances.buffer[instance_id];
    if (!inst->active)
    {
        return;
    }

    inst->local_to_world = local_to_world;
    inst->world_to_local = glms_mat4_inv(local_to_world);
    inst->normal_matrix = compute_normal_matrix(local_to_world);
    if (inst->model_id < scene->mesh_models.capacity && scene->mesh_models.buffer[inst->model_id].active)
    {
        inst->world_bounds = aabb_transform(local_to_world, scene->mesh_models.buffer[inst->model_id].local_bounds);
    }

    scene->tlas_dirty = true;
}