#include "Algorithm/PathTracing.h"
#include "Algorithm/RayIntersection.h"
#include "Lighting/LightEvaluation.h"
#include "Algorithm/BSDF.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;

    // PDF of the direction that generated the current ray segment (used for MIS on env hits).
    float last_bsdf_pdf = 0.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,
                .spread_angle = active_ray.spread_angle,
            };
            path_output sky_output = evaluate_bsdf_sky(&scene->lights, &light_context, throughput, sample_index);

            // MIS for BSDF-sampled environment hit (for depth==0 camera ray, use weight 1).
            float w = 1.0f;
            if (depth > 0)
            {
                float pdf_env = sky_output.pdf;
                float pdf_bsdf = last_bsdf_pdf;
                if (pdf_env > 0.0f && pdf_bsdf > 0.0f)
                {
                    w = power_heuristic(pdf_bsdf, pdf_env);
                }
            }

            vec3s env_contrib = glms_vec3_scale(sky_output.direct_lighting, w);
            accumulated_color = glms_vec4_add(accumulated_color, glms_vec4(env_contrib, 0.0f));
            break;
        }

        uint8_t material_id = scene->triangles.buffer[closest_hit.triangle_id].material_id;
        const material_t* hit_material = &scene->materials.buffer[material_id];
        
        // Calculate ray cone width at the hit point
        float current_cone_width = active_ray.width + closest_hit.distance * active_ray.spread_angle;

        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,

            .triangle_id = closest_hit.triangle_id,
            .cone_width = current_cone_width,
            .spread_angle = active_ray.spread_angle,
        };

        path_output material_output = render_material(hit_material, &shading_context);
        if (glms_vec3_isinf(material_output.direct_lighting) || glms_vec3_isnan(material_output.direct_lighting))
        {
            goto end_path_trace;
        }

        accumulated_color = glms_vec4_add(accumulated_color, glms_vec4(material_output.direct_lighting, 0.0f));

        if (material_output.pdf < FLT_EPSILON)
        {
            goto end_path_trace;
        }

        last_bsdf_pdf = material_output.pdf;

        throughput = glms_vec3_mul(throughput, material_output.bsdf);
        if (glms_vec3_isinf(throughput) || glms_vec3_isnan(throughput))
        {
            goto end_path_trace;
        }

        // 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 PATH_THROUGH:
            //    active_ray = ray_create(BIAS_RAY_ORIGION(closest_hit.point, glms_vec3_negate(closest_hit.normal)), active_ray.direction);
            //    continue;
            case PS_SUCCESS:
                vec3s origin = offset_ray_origin(closest_hit.point, closest_hit.normal, shading_context.wo);
                active_ray = ray_create(origin, material_output.wi, current_cone_width, material_output.spread_angle);
                depth++;
                break;
            default:
                goto end_path_trace;
        }
    }

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,

        .triangle_id = closest_hit.triangle_id,
        .cone_width = ray.width + closest_hit.distance * ray.spread_angle,
        .spread_angle = ray.spread_angle,
    };

    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;
}