#include "Algorithm/PathTracing.h"
#include "Common.h"
#include "Material.h"

static hit_result_t ray_intersect(const triangle_t triangle, const ray_t ray)
{
    hit_result_t result = {0};

    vec3s normal = triangle.normal;
    float n_dot_r = glms_vec3_dot(normal, ray.direction);
    if (n_dot_r > 0.0f)
    {
        normal = glms_vec3_scale(normal, -1.0f);
    }
    
    // triangle is parallel to the ray
    if (fabsf(n_dot_r) < FLT_EPSILON)
    {
        result.hit = false;
        return result;
    }

    // Get distance from ray origin to triangle plane
    float distance = (glms_vec3_dot(normal, triangle.point1) - glms_vec3_dot(normal, ray.origin)) / glms_vec3_dot(normal, ray.direction);

    if (distance < FLT_EPSILON)
    {
        result.hit = false;
        return result;
    }

    vec3s intersection_point = glms_vec3_add(ray.origin, glms_vec3_scale(ray.direction, distance));
    
    // Check if the intersection point is inside the triangle using barycentric coordinates
    vec3s edge1 = glms_vec3_sub(triangle.point2, triangle.point1);
    vec3s edge2 = glms_vec3_sub(triangle.point3, triangle.point2);
    vec3s edge3 = glms_vec3_sub(triangle.point1, triangle.point3);
    vec3s vp = glms_vec3_sub(intersection_point, triangle.point1);
    vec3s vp2 = glms_vec3_sub(intersection_point, triangle.point2);
    vec3s vp3 = glms_vec3_sub(intersection_point, triangle.point3);

    vec3s c1 = glms_vec3_cross(edge1, vp);
    vec3s c2 = glms_vec3_cross(edge2, vp2);
    vec3s c3 = glms_vec3_cross(edge3, vp3);

    if (glms_vec3_dot(triangle.normal, c1) < 0.0f
        || glms_vec3_dot(triangle.normal, c2) < 0.0f 
        || glms_vec3_dot(triangle.normal, c3) < 0.0f)
    {
        result.hit = false;
        return result;
    }

    result.hit = true;
    result.point = intersection_point;
    result.normal = normal;
    result.distance = distance;

    return result;
}

// TODO: Implement faster methods like BVH, KD-Tree or uniform grid acceleration
vec3s path_trace(const triangle_collection_t* triangles, const material_collection_t* materials, ray_t ray, int max_depth)
{
    vec3s accumulated_color = glms_vec3_zero();
    vec3s throughput = glms_vec3_one();

    int depth = 0;
    while (depth < max_depth)
    {
        uint8_t material_id = 255;
        hit_result_t closest_hit = {0};
        closest_hit.distance = 1145141919.810f;

        for (uint64_t i = 0; i < triangles->count; i++)
        {
            hit_result_t hit_result = ray_intersect(triangles->buffer[i], ray);
            if (hit_result.hit && hit_result.distance < closest_hit.distance)
            {
                closest_hit = hit_result;
                material_id = triangles->buffer[i].material_id;
            }
        }

        if (!closest_hit.hit)
        {
            // accumulated_color = glms_vec3_add(accumulated_color, throughput); // TODO: Skybox
            break;
        }

        material_t* hit_material = &materials->buffer[material_id];
        vec3s emission = hit_material->emission;
        accumulated_color = glms_vec3_add(accumulated_color, glms_vec3_mul(throughput, emission));

        float pdf;
        vec3s wo = glms_vec3_negate(ray.direction); // We need to negate the direction of the incoming ray
        vec3s wi = sample_material_bsdf(hit_material, closest_hit.normal, wo, &pdf);
        //vec3s wi = random_cosine_direction(closest_hit.normal);
        //float cos_theta_s = fmaxf(0.0f, glms_vec3_dot(wi, closest_hit.normal));
        //float pdf = cos_theta_s / (float)M_PI; 
        if (pdf < 0.0f)
        {
            break;
        }

        shading_context_t shading_context = {
            .normal = closest_hit.normal,
            .wi = wi,
            .wo = wo
        };
        vec3s bsdf = evaluate_material_bsdf(hit_material, &shading_context);
        float cos_theta = fmaxf(0.0f, glms_vec3_dot(wi, closest_hit.normal));
        throughput = glms_vec3_mul(throughput, glms_vec3_scale(bsdf, cos_theta / pdf));

        // We do Russian roulette to decide whether to continue tracing or terminate the path
        if (depth > 2)
        {
            float q = fminf(glms_vec3_max(throughput), 0.95f);
            if (random_float() > q)
            {
                break; // Terminate the path
            }
        }

        ray.origin = glms_vec3_add(closest_hit.point, glms_vec3_scale(closest_hit.normal, FLT_EPSILON));
        ray.direction = wi;

        depth++;
    }

    return accumulated_color;
}