#include "Algorithm/PathTracing.h"
#include "Algorithm/RayIntersection.h"
#include "Lighting/LightEvaluation.h"

// TODO: Split the diffuse and specular into different Monte Carlo, so we can decide the sample count for each one
vec4s path_trace(const scene_t* scene, ray_t ray, uint32_t sample_index, uint16_t max_depth)
{
    vec4s accumulated_color = (vec4s){0.0f, 0.0f, 0.0f, 1.0f};
    vec3s throughput = glms_vec3_one();

    ray_t active_ray = ray;
    float pdf_bsdf = 1.0f;

    uint16_t depth = 0;
    while (depth < max_depth)
    {
        hit_result_t closest_hit = ray_intersect_scene_closest(&active_ray, scene);

        if (!closest_hit.hit)
        {
            // Set bvh to null indicate that the ray is not hit anything
            light_shading_context_t light_context =
            {
                .wo = active_ray.direction,
                .textures = &scene->textures,
            };
            path_output sky_output = evaluate_bsdf_sky(&scene->lights, &light_context, throughput, sample_index);
            accumulated_color = glms_vec4_add(accumulated_color, glms_vec4(sky_output.direct_lighting, 0.0f));
            break;
        }

        const material_t* hit_material = &scene->materials.buffer[scene->triangles.buffer[closest_hit.triangle_id].material_id];
        shading_context_t shading_context =
        {
            .camera_position = scene->camera.position,
            .camera_direction = glms_vec3_normalize(glms_vec3_sub(closest_hit.point, scene->camera.position)),

            .position = closest_hit.point,
            .normal = closest_hit.normal,
            .tangent = closest_hit.tangent,
            .uv = closest_hit.uv,
            .wo = active_ray.direction,
            .throughput = throughput,

            .sample_index = sample_index,
            .bounce_depth = depth,

            .bvh_tree = &scene->bvh_tree,
            .triangles = &scene->triangles,
            .lights = &scene->lights,
            .textures = &scene->textures,
        };

        path_output material_output = render_material(hit_material, &shading_context);
        accumulated_color = glms_vec4_add(accumulated_color, glms_vec4(material_output.direct_lighting, 0.0f));
        pdf_bsdf = material_output.pdf;

        throughput = glms_vec3_mul(throughput, material_output.bsdf);

        // We do Russian roulette to decide whether to continue tracing or terminate the path
        if (depth > 1)
        {
            float q = fminf(glms_vec3_max(throughput), 0.95f);
            float rr = sobol_sample(sample_index, sobol_get_dimension(depth, PRNG_TERMINATE));
            if (rr > q)
            {
                goto end_path_trace;
            }
            // Keep the energy of the path by scaling the throughput
            throughput = glms_vec3_scale(throughput, 1.0f / q);
        }

        switch (material_output.state)
        {
            case PS_TERMINATE:
                goto end_path_trace;
            //case PATH_THROUGH:
            //    active_ray = ray_create(BIAS_RAY_ORIGION(closest_hit.point, glms_vec3_negate(closest_hit.normal)), active_ray.direction);
            //    continue;
            default:
                active_ray = ray_create(offset_ray_origin(closest_hit.point, closest_hit.normal, shading_context.wo), material_output.wi);
                depth++;
                break;
        }
    }

end_path_trace:
    return accumulated_color;
}

// How to handle multi-bounced aov like indirect lighting?
// Maybe we should move aov to path_trace and split accumulated_color into direct/indirect diffuse/specular before returning.
void render_aov(const scene_t* scene, ray_t ray, uint32_t sample_index, uint16_t max_depth, aov_output_t* aov_output)
{
    hit_result_t closest_hit = ray_intersect_scene_closest(&ray, scene);
    if (!closest_hit.hit)
    {
        return;
    }

    const material_t* hit_material = &scene->materials.buffer[scene->triangles.buffer[closest_hit.triangle_id].material_id];
    shading_context_t shading_context =
    {
        .camera_position = scene->camera.position,
        .camera_direction = glms_vec3_normalize(glms_vec3_sub(closest_hit.point, scene->camera.position)),

        .position = closest_hit.point,
        .normal = closest_hit.normal,
        .tangent = closest_hit.tangent,
        .uv = closest_hit.uv,
        .wo = ray.direction,
        .throughput = 1.0f,

        .sample_index = sample_index,
        .bounce_depth = 0,

        .bvh_tree = &scene->bvh_tree,
        .triangles = &scene->triangles,
        .lights = &scene->lights,
        .textures = &scene->textures,
    };

    render_material_aov(hit_material, &shading_context, aov_output);
    aov_output->position = glms_vec4(closest_hit.point, 1.0f);
    aov_output->depth = closest_hit.distance;
}