#include "Algorithm/RayIntersection.h"
#include "Common.h"
#include "Geometry/Triangle.h"
#include "cglm/struct/vec3.h"

ray_t ray_create(vec3s origin, vec3s direction)
{
    return (ray_t)
    {
        .origin = origin,
        .direction = direction,
        .inverse_direction = glms_vec3_div(glms_vec3_one(), direction),
        .sign =
        ((direction.x < 0.0f) ? 1 : 0) |
        ((direction.y < 0.0f) ? 2 : 0) |
        ((direction.z < 0.0f) ? 4 : 0)
    };
}

static inline float next_float_up(float value)
{
    if (isnan(value) || (isfinite(value) && value > 0))
    {
        return value;
    }

    return nextafterf(value, INFINITY);
}

static inline float next_float_down(float value)
{
    if (isnan(value) || (isfinite(value) && value < 0))
    {
        return value;
    }

    return nextafterf(value, -INFINITY);
}

vec3s offset_ray_origin(vec3s point, vec3s normal, vec3s wo)
{
    vec3s abs_normal = glms_vec3_abs(normal);
    float c = glms_vec3_max(point) * gamma(3);
    float d = glms_vec3_dot(abs_normal, (vec3s){c, c, c});

    vec3s offset = glms_vec3_scale(abs_normal, d);
    if (glms_vec3_dot(wo, normal) < 0.0f)
    {
        offset = glms_vec3_negate(offset);
    }

    vec3s position = glms_vec3_add(point, offset);
    for (int i = 0; i < 3; i++)
    {
        if (offset.raw[i] > 0.0f)
        {
            position.raw[i] = next_float_up(position.raw[i]);
        }
        else if (offset.raw[i] < 0.0f)
        {
            position.raw[i] = next_float_down(position.raw[i]);
        }
    }

    return position;
    //return point;
}

#if 0
// TODO: We still have small amount of block dots in current implementation. It's because of floating point precision. May need to fall back to double or handle it in a different way.
hit_result_t ray_intersect_triangle(const ray_t* ray, const triangle_t* triangle)
{
    hit_result_t result = {0};

    vec3s edge1 = glms_vec3_sub(triangle->vertices[1].position, triangle->vertices[0].position);
    vec3s edge2 = glms_vec3_sub(triangle->vertices[2].position, triangle->vertices[1].position);
    vec3s edge3 = glms_vec3_sub(triangle->vertices[0].position, triangle->vertices[2].position);

    vec3s normal = triangle->face_normal;
    float n_dot_r = glms_vec3_dot(normal, ray->direction);
    if (n_dot_r > 0.0f)
    {
        normal = glms_vec3_negate(normal);
    }

    // triangle is parallel to the ray
    if (fabsf(n_dot_r) < FLT_EPSILON)
    {
        return result;
    }

    // Get distance from ray origin to triangle plane
    float distance = (glms_vec3_dot(normal, triangle->vertices[0].position) - glms_vec3_dot(normal, ray->origin)) / glms_vec3_dot(normal, ray->direction);

    float eps = gamma(3) * glms_vec3_max(glms_vec3_abs(ray->origin));
    if (distance <= eps)
    {
        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 vp = glms_vec3_sub(intersection_point, triangle->vertices[0].position);
    vec3s vp2 = glms_vec3_sub(intersection_point, triangle->vertices[1].position);
    vec3s vp3 = glms_vec3_sub(intersection_point, triangle->vertices[2].position);

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

    float n_dot_c1 = glms_vec3_dot(normal, c1);
    float n_dot_c2 = glms_vec3_dot(normal, c2);
    float n_dot_c3 = glms_vec3_dot(normal, c3);

    bool r1 = n_dot_c1 > 0.0f && n_dot_c2 > 0.0f && n_dot_c3 > 0.0f;
    bool r2 = n_dot_c1 < 0.0f && n_dot_c2 < 0.0f && n_dot_c3 < 0.0f;

    if (!r1 && !r2)
    {
        return result;
    }

    // TODO: Normal interpolation
    vec3s v0v1 = glms_vec3_sub(triangle->vertices[1].position, triangle->vertices[0].position);
    vec3s v0v2 = glms_vec3_sub(triangle->vertices[2].position, triangle->vertices[0].position);
    vec3s v0p = glms_vec3_sub(intersection_point, triangle->vertices[0].position);
    float d00 = glms_vec3_dot(v0v1, v0v1); // dot(v0v1, v0v1)
    float d01 = glms_vec3_dot(v0v1, v0v2); // dot(v0v1, v0v2)
    float d11 = glms_vec3_dot(v0v2, v0v2); // dot(v0v2, v0v2)
    float d20 = glms_vec3_dot(v0p,  v0v1); // dot(v0p, v0v1)
    float d21 = glms_vec3_dot(v0p,  v0v2); // dot(v0p, v0v2)
    float denom = d00 * d11 - d01 * d01;
    //if (denom < FLT_EPSILON)
    //{
    //    return result;
    //}

    float invDenom = 1.0f / denom;
    float u = (d11 * d20 - d01 * d21) * invDenom; // This is b1 (weight for V1)
    float v = (d00 * d21 - d01 * d20) * invDenom; // This is b2 (weight for V2)

    float w = 1.0f - u - v;
  
    //vec3s smooth_normal = glms_vec3_add(glms_vec3_scale(triangle->vertices[0].normal, u), glms_vec3_add(glms_vec3_scale(triangle->vertices[1].normal, v), glms_vec3_scale(triangle->vertices[2].normal, w)));

    result.hit = true;
    result.point = intersection_point;
    result.normal = normal;
    result.uv = (vec2s)
    {
        .x = w * triangle->vertices[0].uv.x + u * triangle->vertices[1].uv.x + v * triangle->vertices[2].uv.x,
        .y = w * triangle->vertices[0].uv.y + u * triangle->vertices[1].uv.y + v * triangle->vertices[2].uv.y,
    };
    result.distance = distance;

    return result;
}
#else
hit_result_t ray_intersect_triangle(const ray_t* ray, const triangle_t* triangle)
{
    hit_result_t result = {0};

    vec3s origin = ray->origin;
    vec3s direction = ray->direction;
    vec3s v0 = triangle->vertices[0].position;
    vec3s v1 = triangle->vertices[1].position;
    vec3s v2 = triangle->vertices[2].position;

    vec3s e1 = glms_vec3_sub(v1, v0);
    vec3s e2 = glms_vec3_sub(v2, v0);

    // Begin Möller–Trumbore
    vec3s P = glms_vec3_cross(direction, e2);
    float det = glms_vec3_dot(e1, P);

    if (fabsf(det) < FLT_EPSILON)
    {
        return result;
    }

    float invDet = 1.0f / det;

    // Calculate barycentric u
    vec3s T = glms_vec3_sub(origin, v0);
    float u = glms_vec3_dot(T, P) * invDet;
    if (u < 0.0f || u > 1.0f)
    {
        return result;
    }

    // Calculate barycentric v
    vec3s Q = glms_vec3_cross(T, e1);
    float v = glms_vec3_dot(direction, Q) * invDet;
    if (v < 0.0f || u + v > 1.0f)
    {
        return result;
    }

    // Distance along the ray
    float t = glms_vec3_dot(e2, Q) * invDet;
    if (t < RAY_EPSILON)
    {
        return result;
    }

    float w = 1.0f - u - v;

    result.hit = true;
    result.distance = t;
    result.point = glms_vec3_add(origin, glms_vec3_scale(direction, t));

    vec3s normal = glms_vec3_scale(triangle->vertices[0].normal, w);
    normal = glms_vec3_add(normal, glms_vec3_scale(triangle->vertices[1].normal, u));
    normal = glms_vec3_add(normal, glms_vec3_scale(triangle->vertices[2].normal, v));
    normal = glms_vec3_dot(normal, direction) < 0.0f ? normal : glms_vec3_negate(normal);
    result.normal = glms_vec3_normalize(normal);

    vec3s tangent = glms_vec3_scale(triangle->vertices[0].tangent, w);
    tangent = glms_vec3_add(tangent, glms_vec3_scale(triangle->vertices[1].tangent, u));
    tangent = glms_vec3_add(tangent, glms_vec3_scale(triangle->vertices[2].tangent, v));
    result.tangent = glms_vec3_normalize(tangent);

    result.uv.x = w * triangle->vertices[0].uv.x
             + u * triangle->vertices[1].uv.x
             + v * triangle->vertices[2].uv.x;
    result.uv.y = w * triangle->vertices[0].uv.y
             + u * triangle->vertices[1].uv.y
             + v * triangle->vertices[2].uv.y;

    return result;
}
#endif

bool ray_intersect_aabb(const ray_t* ray, aabb_t aabb, float* enter_out, float* exit_out)
{
    // select slab min/max per axis based on sign:
    float tx_min = ((SIGN_BIT(ray->sign, 0) ? aabb.max.x : aabb.min.x) - ray->origin.x ) * ray->inverse_direction.x;
    float tx_max = ((SIGN_BIT(ray->sign, 0) ? aabb.min.x : aabb.max.x) - ray->origin.x ) * ray->inverse_direction.x;

    float ty_min = ((SIGN_BIT(ray->sign, 1) ? aabb.max.y : aabb.min.y) - ray->origin.y ) * ray->inverse_direction.y;
    float ty_max = ((SIGN_BIT(ray->sign, 1) ? aabb.min.y : aabb.max.y) - ray->origin.y ) * ray->inverse_direction.y;

    // early exit if slabs miss
    if ((tx_min > ty_max) || (ty_min > tx_max))
    {
        return false;
    }

    // merge X and Y
    float t0 = tx_min > ty_min ? tx_min : ty_min;
    float t1 = tx_max < ty_max ? tx_max : ty_max;

    float tz_min = ( (SIGN_BIT(ray->sign, 2) ? aabb.max.z : aabb.min.z) - ray->origin.z ) * ray->inverse_direction.z;
    float tz_max = ( (SIGN_BIT(ray->sign, 2) ? aabb.min.z : aabb.max.z) - ray->origin.z ) * ray->inverse_direction.z;

    // final overlap test
    if ((t0 > tz_max) || (tz_min > t1))
    {
        return false;
    }

    // update entry/exit
    if (enter_out != NULL)
    {
        *enter_out = t0 > tz_min ? t0 : tz_min;
    }
    if (exit_out != NULL)
    {
        *exit_out  = t1 < tz_max ? t1 : tz_max;
    }

    return true;
}

static inline float distance_to_aabb(vec3s point, aabb_t aabb)
{
    float dx = fmaxf(aabb.min.x - point.x, 0.0f) + fmaxf(point.x - aabb.max.x, 0.0f);
    float dy = fmaxf(aabb.min.y - point.y, 0.0f) + fmaxf(point.y - aabb.max.y, 0.0f);
    float dz = fmaxf(aabb.min.z - point.z, 0.0f) + fmaxf(point.z - aabb.max.z, 0.0f);
    return sqrtf(dx * dx + dy * dy + dz * dz);
}

// TODO: Use a stack to avoid recursion.
void ray_intersect_bvh(const ray_t* ray, const bvh_node_t* bvh_nodes, const uint64_t* primitive_indices, const triangle_collection_t* all_triangles, uint64_t node_index, float* closest_out,
                       hit_result_t* best_hit_out)
{
    const bvh_node_t* node = &bvh_nodes[node_index];

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

    if (enter > *closest_out)
    {
        return;
    }

    // If primitive_count > 0 implies leaf:
    if (node->primitive_count > 0)
    {
        for (uint32_t i = 0; i < node->primitive_count; i++)
        {
            uint64_t triangle_index = primitive_indices[node->start_index + i];
            hit_result_t hit_result = ray_intersect_triangle(ray, &all_triangles->buffer[triangle_index]);
            if (hit_result.hit && hit_result.distance < *closest_out)
            {
                hit_result.triangle_id = triangle_index;
                *best_hit_out = hit_result;
                *closest_out = hit_result.distance;
            }
        }
    }
    else // Internal node (primitive_count == 0 implies internal)
    {
        uint64_t left_child_index = node->left_child_offset;
        uint64_t right_child_index = node->right_child_offset;

        const bvh_node_t* left_child = &bvh_nodes[left_child_index];
        const bvh_node_t* right_child = &bvh_nodes[right_child_index];

        float left_enter, left_exit, right_enter, right_exit;
        bool hit_left_aabb = ray_intersect_aabb(ray, left_child->bounds, &left_enter, &left_exit);
        bool hit_right_aabb = ray_intersect_aabb(ray, right_child->bounds, &right_enter, &right_exit);

        // Traverse children based on closest AABB and current best hit distance (*closest_t)
        if (hit_left_aabb && hit_right_aabb)
        {
            if (left_enter < right_enter)
            {
                ray_intersect_bvh(ray, bvh_nodes, primitive_indices, all_triangles, node->left_child_offset, closest_out, best_hit_out);
                if (right_enter < *closest_out)
                {
                    ray_intersect_bvh(ray, bvh_nodes, primitive_indices, all_triangles, node->right_child_offset, closest_out, best_hit_out);
                }
            }
            else
            {
                ray_intersect_bvh(ray, bvh_nodes, primitive_indices, all_triangles, node->right_child_offset, closest_out, best_hit_out);
                if (left_enter < *closest_out)
                {
                    ray_intersect_bvh(ray, bvh_nodes, primitive_indices, all_triangles, node->left_child_offset, closest_out, best_hit_out);
                }
            }
        }
        else if (hit_left_aabb)
        {
            ray_intersect_bvh(ray, bvh_nodes, primitive_indices, all_triangles, node->left_child_offset, closest_out, best_hit_out);
        }
        else if (hit_right_aabb)
        {
            ray_intersect_bvh(ray, bvh_nodes, primitive_indices, all_triangles, node->right_child_offset, closest_out, best_hit_out);
        }
    }
}

hit_result_t ray_intersect_scene(const ray_t* ray, const scene_t* scene)
{
    hit_result_t result = {0};
    float closest = FLT_MAX;

    if (scene == NULL || scene->bvh_tree.nodes == NULL || scene->triangles.count == 0 || scene->bvh_tree.node_count == 0 || scene->bvh_tree.primitive_count == 0)
    {
        return result;
    }
    ray_intersect_bvh(ray, scene->bvh_tree.nodes, scene->bvh_tree.primitive_indices, &scene->triangles, 0, &closest, &result);

    return result;
}