#include <omp.h>
#include <stdint.h>
#include <svpng.inc>

#include "Algorithm/Sobol.h"
// #include "Geometry/GeometryUtilities.h"
#include "Geometry/Mesh.h"
#include "Lighting/SkyLight.h"
// #include "Material/SimpleLit.h"
#include "Rendering/PostProcessing.h"
#include "Rendering/Scene.h"
#include "Window.h"

#define TITLE "Path Tracing"
#define SPONZA_PATH "./assets/sponza.fbx"

static bool scene_setup(scene_t* scene)
{
    if (!scene_init(scene, 167000, 3, 8, 1))
    {
        return false;
    }

    scene->camera.position = (vec3s){-7.5f, 2.0f, 0.0f};
    scene->camera.rotation = euler_to_quat(10.0f, -90.0f, 0.0f);

    // TODO: Standardize light unit
    light_entity_t sun = light_create_directional_light(&scene->lights);
    directional_light_t* sun_light = &scene->lights.directional_lights[sun.id];
    sun_light->direction = glms_vec3_normalize((vec3s){-0.5f, 1.0f, 0.25f});
    sun_light->color = (vec3s){1.0f, 0.93f, 0.87f};
    sun_light->intensity = 2.0f;
    sun_light->angular_diameter = 0.53f;

    scene->lights.sky_light = sky_create_constant_sky(&(constant_sky_data_t)
    {
        .color = (vec3s){0.73f, 0.82f, 1.0f},
        .intensity = 1.0f,
    });

    return true;
}

static bool load_assets(scene_t* scene)
{
    mesh_load(SPONZA_PATH, scene);
    // quad_create((vec3s){0.0f, 0.0f, 0.0f}, (vec3s){0.0f, 1.0f, 0.0f}, (vec3s){1.0f, 0.0f, 0.0f}, 10.0f, floor_material.id, &scene->triangles);

    return scene_build_bvh(scene);
}

static bool initialize_renderer(const rendering_config_t* config, render_job_t** outJob, render_target_t* outImg, scene_t* outScene)
{
    if (!scene_setup(outScene)
        || !load_assets(outScene)
        || !scene_build_bvh(outScene)
        || !render_target_init(config->width, config->height, outImg))
    {
        return false;
    }

    render_job_t* job = malloc(sizeof(render_job_t));
    if (job == NULL)
    {
        return false;
    }

    *job = (render_job_t){
        .scene = outScene,
        .render_target = outImg,
        .config = config,

        .rendering_type = TILE_BASED,
        .debug_flag = DEBUG_NONE,
        .is_done = false,
    };

    sobol_init();

    *outJob = job;
    return true;
}

static void shutdown_renderer(render_job_t* job, render_target_t* img, scene_t* scene)
{
    if (job != NULL)
    {
        free(job);
    }

    render_target_free(img);
    scene_free(scene);
}

static void update_pixel_buffer(render_target_t* render_target)
{
    if (render_target == NULL || render_target->buffer == NULL)
    {
        return;
    }

    for (uint32_t y = 0; y < render_target->height; y++)
    {
        for (uint32_t x = 0; x < render_target->width; x++)
        {
            vec4s pixel = render_target_get_pixel(render_target, x, y);
            pixel = gamma_correct(pixel, 2.2f);
            pixel = aces_tone_map(pixel);
            window_update_pixel(pixel, x, y);
        }
    }

    window_refresh_region(0, 0, render_target->width, render_target->height);
}

static int run_main_loop(render_job_t* job, render_target_t* img)
{
    MSG msg;
    bool running = true;

    while (running)
    {
        while (PeekMessage(&msg, NULL, 0, 0, PM_REMOVE))
        {
            switch (msg.message)
            {
                case WM_QUIT:
                case WM_CLOSE:
                {
                    running = false;
                    job->is_done = true;
                }
            }

            TranslateMessage(&msg);
            DispatchMessage(&msg);
        }

        if (!job->is_done && running)
        {
            update_pixel_buffer(img);
        }
    }

    return 0;
}

// int main()
int WINAPI wWinMain(_In_ HINSTANCE hInstance, _In_opt_ HINSTANCE hPrevInstance, _In_ PWSTR pCmdLine, _In_ int nCmdShow)
{
    omp_set_num_threads(16);

    scene_t scene = {0};
    render_target_t img = {0};
    render_job_t* job = NULL;

    rendering_config_t config = {
        .width = 1920 / 2,
        .height = 1080 / 2,
        .sample_count = 64 * 1,
        .max_depth = 4,
        .bucket_size = 64,
    };

    if (!initialize_renderer(&config, &job, &img, &scene)
        || !window_create(TITLE, hInstance, config.width, config.height, job))
    {
        shutdown_renderer(job, &img, &scene);
        return -1;
    }

    int result = run_main_loop(job, &img);

    window_close();
    shutdown_renderer(job, &img, &scene);

    return result;
}