#include "Rendering/Debug.h"
#include "Algorithm/RayIntersection.h"

static void ray_intersect_bvh_count(ray_t ray, bvh_tree_t bvh_tree, uint64_t node_index, uint32_t* count_out)
{
    const float _MAX_DIST = 1e6f;
    if (bvh_tree.nodes == NULL || bvh_tree.primitive_indices == NULL || count_out == NULL)
    {
        return;
    }

    const bvh_node_t* node = &bvh_tree.nodes[node_index];

    float enter, exit;
    if (!ray_intersect_aabb(&ray, node->bounds, &enter, &exit))
    {
        return;
    }

    if (enter > _MAX_DIST || exit < 0.0f)
    {
        return;
    }

    (*count_out)++;
    if (node->primitive_count == 0)
    {
        // Internal node
        ray_intersect_bvh_count(ray, bvh_tree, node->left_child_offset, count_out);
        ray_intersect_bvh_count(ray, bvh_tree, node->right_child_offset, count_out);
    }
}

vec4s render_debug(scene_t* scene, ray_t ray, uint16_t sample_index, int flag)
{
    if (scene == NULL)
    {
        return glms_vec4_zero();
    }

    switch (flag & 0xFF)
    {
    case DEBUG_BVH:
        uint32_t count = 0;
        ray_intersect_bvh_count(ray, scene->bvh_tree, 0, &count);

        vec4s result = glms_vec4_zero();
        for (uint32_t i = 0; i < count; i++)
        {
            result = glms_vec4_add(result, DEBUG_COLOR_BVH);
        }
        return result;

    case DEBUG_SOBOL:
        float sobol_sample_value = sobol_sample(sample_index, 1); // Assuming dimension 0 for simplicity
        return (vec4s){sobol_sample_value, sobol_sample_value, sobol_sample_value, 1.0f};

    case DEBUG_UV:
        hit_result_t hit_result = ray_intersect_scene_closest(&ray, scene);
        return (vec4s){fmodf(fabsf(hit_result.uv.x), 1.0f), fmodf(fabsf(hit_result.uv.y), 1.0f), 0.0f, 1.0f};

    default:
        return glms_vec4_zero();
    }
}