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

ray_t ray_create(vec3s origin, vec3s direction, float cone_width, float spread_angle)
{
    ray_t ray =
    {
        .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),
        .width = cone_width,
        .spread_angle = spread_angle
    };

    ray.esp = glms_vec3_max(glms_vec3_abs(ray.origin)) * gamma(15);
    return ray;
}

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 w)
{
    float c = glms_vec3_max(glms_vec3_abs(point)) * gamma(15);
    vec3s error = (vec3s){c, c, c};
    // float g = gamma(10);
    // vec3s error = {fabsf(point.x) * g, fabsf(point.y) * g, fabsf(point.z) * g};
    float d = glms_vec3_dot(glms_vec3_abs(normal), error);

    vec3s offset = glms_vec3_scale(normal, d);
    if (glms_vec3_dot(glms_vec3_negate(w), 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;
}

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->esp)
    {
        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));

    // Should we output u, v, w instead of normal, tangent, and uv?
    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;
}

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 && exit_out != NULL)
    {
        *enter_out = t0 > tz_min ? t0 : tz_min;
        *exit_out  = t1 < tz_max ? t1 : tz_max;

        if (fmaxf(*enter_out, ray->esp) > *exit_out)
        {
            return false;
        }
    }

    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_closest(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_closest(ray, bvh_nodes, primitive_indices, all_triangles, node->left_child_offset, closest_out, best_hit_out);
                if (right_enter < *closest_out)
                {
                    ray_intersect_bvh_closest(ray, bvh_nodes, primitive_indices, all_triangles, node->right_child_offset, closest_out, best_hit_out);
                }
            }
            else
            {
                ray_intersect_bvh_closest(ray, bvh_nodes, primitive_indices, all_triangles, node->right_child_offset, closest_out, best_hit_out);
                if (left_enter < *closest_out)
                {
                    ray_intersect_bvh_closest(ray, bvh_nodes, primitive_indices, all_triangles, node->left_child_offset, closest_out, best_hit_out);
                }
            }
        }
        else if (hit_left_aabb)
        {
            ray_intersect_bvh_closest(ray, bvh_nodes, primitive_indices, all_triangles, node->left_child_offset, closest_out, best_hit_out);
        }
        else if (hit_right_aabb)
        {
            ray_intersect_bvh_closest(ray, bvh_nodes, primitive_indices, all_triangles, node->right_child_offset, closest_out, best_hit_out);
        }
    }
}

void ray_intersect_bvh_any(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, hit_result_t* any_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 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)
            {
                *any_hit_out = hit_result;
                any_hit_out->triangle_id = triangle_index;

                return;
            }
        }
    }
    else
    {
        // Internal node: traverse children in near‐first order
        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);

        if (hit_left_aabb && hit_right_aabb)
        {
            if (left_enter < right_enter)
            {
                ray_intersect_bvh_any(ray, bvh_nodes, primitive_indices, all_triangles, node->left_child_offset, any_hit_out);
                if (!any_hit_out->hit)
                {
                    ray_intersect_bvh_any(ray, bvh_nodes, primitive_indices, all_triangles, node->right_child_offset, any_hit_out);
                }
            }
            else
            {
                ray_intersect_bvh_any(ray, bvh_nodes, primitive_indices, all_triangles, node->right_child_offset, any_hit_out);
                if (!any_hit_out->hit)
                {
                    ray_intersect_bvh_any(ray, bvh_nodes, primitive_indices, all_triangles, node->left_child_offset, any_hit_out);
                }
            }
        }
        else if (hit_left_aabb)
        {
            ray_intersect_bvh_any(ray, bvh_nodes, primitive_indices, all_triangles, node->left_child_offset, any_hit_out);
        }
        else if (hit_right_aabb)
        {
            ray_intersect_bvh_any(ray, bvh_nodes, primitive_indices, all_triangles, node->right_child_offset, any_hit_out);
        }
    }
}

hit_result_t ray_intersect_scene_closest(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_closest(ray, scene->bvh_tree.nodes, scene->bvh_tree.primitive_indices, &scene->triangles, 0, &closest, &result);
    return result;
}

hit_result_t ray_intersect_scene_any(const ray_t* ray, const scene_t* scene)
{
    hit_result_t result = {0};
    result.distance = 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_any(ray, scene->bvh_tree.nodes, scene->bvh_tree.primitive_indices, &scene->triangles, 0, &result);
    return result;
}