#include "Debug.h"
#include "Rendering/Scene.h"
#include "Algorithm/PathTracing.h"

#define FLIP_Y

bool scene_init(uint64_t triangle_count, uint8_t material_count, uint32_t punctual_light_count, scene_t* scene)
{
    scene_t temp = {0};

    if (!triangle_collection_init(triangle_count, &temp.triangles))
    {
        goto triangle_failed;
    }

    if (!material_collection_init(material_count, &temp.materials))
    {
        goto material_failed;
    }

    if (!light_collection_create(punctual_light_count, 16, &temp.lights)) // NOTE: We just fixed the max directional light count to 16.
    {
        goto light_failed;
    }

    temp.camera = camera_create(
        (vec3s){0.0f, 0.0f, 5.0f},
        (vec3s){0.0f, 0.0f, -1.0f},
        (vec3s){0.0f, 1.0f, 0.0f},
        0.025f,
        0.036f,
        16.0f / 9.0f
    );

    *scene = temp;
    return true;

light_failed:
    material_collection_free(&temp.materials);
material_failed:
    triangle_collection_free(&temp.triangles);
triangle_failed:
    return false;
}

void scene_free(scene_t* scene)
{
    bvh_tree_free(&scene->bvh_tree);
    triangle_collection_free(&scene->triangles);
    material_collection_free(&scene->materials);
    light_collection_free(&scene->lights);
}

bool scene_build_bvh(scene_t* scene)
{
    if (scene == NULL || scene->triangles.count == 0)
    {
        return false;
    }

    bvh_tree_t bvh_tree = {0};
    if (!bvh_tree_init(&bvh_tree, &scene->triangles))
    {
        return false;
    }

    bvh_tree_build(&bvh_tree);
    scene->bvh_tree = bvh_tree;
    return true;
}

static inline void ensure_camera_aspect_ratio(camera_t* camera,  rendering_config_t config)
{
    if (fabsf((float)config.width / config.height - camera->aspect_ratio) > FLT_EPSILON)
    {
        *camera = camera_create(
            camera->position,
            camera->forward,
            camera->up,
            camera->focal_length,
            camera->size_x,
            (float)config.width / (float)config.height
        );
    }
}

static inline vec2s compute_ndc(float x, float y, uint32_t width, uint32_t height)
{
    return (vec2s){
        .x = x / (float)width,
#ifdef FLIP_Y
        .y = 1.0f - y / (float)height
#else
        .y = y / (float)height
#endif
    };
}

static inline uint16_t get_sample_count(uint16_t sample_count, int flag)
{
    switch (flag)
    {
        case DEBUG_BVH:
        case DEBUG_SOBOL:
            return 1;
        default:
            return sample_count;
    }

    return sample_count;
}

static void screen_render_pixel(scene_t* scene, rendering_config_t config, vec3s coord,  uint32_t x,  uint32_t y, int flag, vec4s* pixel_color)
{
    vec4s accumulated_color = glms_vec4_zero();
    *pixel_color = accumulated_color;
    uint32_t pixel_id = y * config.width + x;

    uint16_t sample_count = get_sample_count(config.sample_count, flag);

    for (uint16_t k = 0; k < sample_count; k++)
    {
        vec2s position_ndc = compute_ndc((float)x, (float)y, config.width, config.height);
        float screen_x = position_ndc.x * 2.0f - 1.0f;
        float screen_y = position_ndc.y * 2.0f - 1.0f;
        float sensor_offset_x = screen_x * scene->camera.size_x / 2.0f;
        float sensor_offset_y = screen_y * scene->camera.size_y / 2.0f;

        vec3s image_plane_point = coord;
        image_plane_point = glms_vec3_add(image_plane_point, glms_vec3_scale(scene->camera.right, sensor_offset_x));
        image_plane_point = glms_vec3_add(image_plane_point, glms_vec3_scale(scene->camera.up, sensor_offset_y));

        ray_t ray = {
            .origin = scene->camera.position,
            .direction = glms_vec3_normalize(glms_vec3_sub(image_plane_point, scene->camera.position))
        };

        vec4s out_color = glms_vec4_zero();

        if (flag != 0)
        {
            out_color = render_debug(scene, ray, k, flag);
        }
        else
        {
            // TODO: Hash it
            uint32_t sobol_idx = pixel_id * config.sample_count + (k + 1);
            out_color = path_trace(scene, ray, sobol_idx, config.max_depth);
        }

        accumulated_color = glms_vec4_add(accumulated_color, out_color);
    }

    *pixel_color = glms_vec4_scale(accumulated_color, 1.0f / (float)sample_count);
}


bool scene_render_tile(scene_t* scene, rendering_context_t* ctx, rendering_config_t config,  uint32_t tile_index, int rendering_flag, render_target_t* render_target, tile_t* tile_out)
{
    if (ctx->is_done)
    {
        return false;
    }

    if(!ctx->is_init)
    {
        ensure_camera_aspect_ratio(&scene->camera, config);
        ctx->tile_count_x = (config.width + config.bucket_size - 1) / config.bucket_size;
        ctx->tile_count_y = (config.height + config.bucket_size - 1) / config.bucket_size;
        ctx->coord = glms_vec3_add(scene->camera.position, glms_vec3_scale(scene->camera.forward, scene->camera.focal_length));
        ctx->is_init = true;
    }

    if (tile_index >= ctx->tile_count_x * ctx->tile_count_y)
    {
        return false;
    }

    uint32_t tile_x_0 = tile_index % ctx->tile_count_x * config.bucket_size;
    uint32_t tile_y_0 = tile_index / ctx->tile_count_x * config.bucket_size;
    uint32_t tile_x_1 = (uint32_t)fmin(tile_x_0 + config.bucket_size, config.width);
    uint32_t tile_y_1 = (uint32_t)fmin(tile_y_0 + config.bucket_size, config.height);

    tile_out->x = tile_x_0;
    tile_out->y = tile_y_0;
    tile_out->width = tile_x_1 - tile_x_0;
    tile_out->height = tile_y_1 - tile_y_0;

    int64_t x, y;  // OpenMP requires these to be declared outside the parallel region. Also, they need to be signed integers. To avoid overflow, we need to use int64_t
    rendering_config_t config_copy = config; // Have to copy it, otherwise OpenMP will cause it become invalid. Not sure if this is a bug or not.

#pragma omp parallel for schedule(dynamic, 1) default(none) \
    shared(tile_x_0, tile_x_1, tile_y_0, tile_y_1, config_copy) \
    private(x, y)
    for (y = tile_y_0; y < tile_y_1; y++)
    {
        for (x = tile_x_0; x < tile_x_1; x++)
        {
            vec4s pixel_color;
            screen_render_pixel(scene, config_copy, ctx->coord, (uint32_t)x, (uint32_t)y, rendering_flag, &pixel_color);
            render_target_set_pixel(render_target, (uint32_t)x, (uint32_t)y, pixel_color);
        }
    }

    if (tile_index == ctx->tile_count_x * ctx->tile_count_y - 1)
    {
        ctx->is_done = true;
    }

    return true;
}

bool scene_render(scene_t* scene, rendering_config_t config, int rendering_flag, render_target_t* render_target)
{
    ensure_camera_aspect_ratio(&scene->camera, config);

    if (render_target->buffer != NULL)
    {
        render_target_free(render_target);
    }

    if (render_target_init(config.width, config.height, render_target))
    {
        return false;
    }

    uint32_t tile_count_x = (config.width + config.bucket_size - 1) / config.bucket_size;
    uint32_t tile_count_y = (config.height + config.bucket_size - 1) / config.bucket_size;
    uint32_t tile_count = tile_count_x * tile_count_y;

    float inv_sample = 1.0f / config.sample_count;
    vec3s coord = glms_vec3_add(scene->camera.position, glms_vec3_scale(scene->camera.forward, scene->camera.focal_length));

    int64_t x, y, tile_index; // OpenMP requires these to be declared outside the parallel region. Also, they need to be signed integers. To avoid overflow, we need to use int64_t
    rendering_config_t config_copy = config;

#pragma omp parallel for schedule(dynamic, 1) default(none) \
    shared(tile_count_x, tile_count_y, tile_count, config_copy, coord, inv_sample, render_target) \
    private(x, y, tile_index)
    for (tile_index = 0; tile_index < tile_count; tile_index++)
    {
        uint32_t tile_x_0 = (uint32_t)tile_index % tile_count_x * config.bucket_size;
        uint32_t tile_y_0 = (uint32_t)tile_index / tile_count_x * config.bucket_size;
        uint32_t tile_x_1 = (uint32_t)fmin(tile_x_0 + config.bucket_size, config.width);
        uint32_t tile_y_1 = (uint32_t)fmin(tile_y_0 + config.bucket_size, config.height);

        for (y = tile_y_0; y < tile_y_1; y++)
        {
            for (x = tile_x_0; x < tile_x_1; x++)
            {
                vec4s pixel_color;
                screen_render_pixel(scene, config_copy, coord, (uint32_t)x, (uint32_t)y, rendering_flag, &pixel_color);
                render_target_set_pixel(render_target, (uint32_t)x, (uint32_t)y, pixel_color);
            }
        }
    }

    return true;
}