#include <time.h>
#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/StandardLit.h"
#include "Rendering/AOV.h"
#include "Rendering/PostProcessing.h"
#include "Rendering/Scene.h"
#include "Window.h"

#define TITLE "Path Tracing"
#define SCENE_PATH "./assets/sponza.fbx"
#define HDRI_PATH "./assets/hdri/rogland_sunset_1k.hdr"

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.6f, 1.0f, 0.25f});
    sun_light->color = (vec3s){1.0f, 0.93f, 0.87f};
    sun_light->intensity = 1.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,
    // });
    texture_handle_t hdri = texture_load(HDRI_PATH, false, false, FLOAT_32, &scene->textures);
    scene->textures.buffer[hdri.id].texture.wrap_mode = WM_CLAMP;
    scene->textures.buffer[hdri.id].texture.filter_mode = FM_LINEAR;
    scene->lights.sky_light = sky_create_hdr_sky(&scene->textures, hdri, 1.0f);

    return true;
}

static bool load_assets(scene_t* scene)
{
#if 1
    mesh_load(SCENE_PATH, scene);
#else
    material_handle_t floor_material = material_create_standard_lit_default(&(standard_lit_properties_t)
    {
        .albedo = (vec3s){0.95f, 0.95f, 0.95f},
        .roughness = 0.95f,
        .diffuse_roughness = 0.05f,
        .metallic = 0.0f,
        .albedo_texture = invalid_texture_handle(),
        .metallic_texture = invalid_texture_handle(),
        .roughness_texture = invalid_texture_handle(),
        .normal_texture = invalid_texture_handle(),
    }, &scene->materials);

    material_handle_t quad_material = material_create_standard_lit_default(
        &(standard_lit_properties_t){
            .albedo = (vec3s){0.8f, 0.0f, 0.0f},
            .roughness = 0.05f,
            .diffuse_roughness = 0.05f,
            .metallic = 0.0f,
            .albedo_texture = invalid_texture_handle(),
            .metallic_texture = invalid_texture_handle(),
            .roughness_texture = invalid_texture_handle(),
            .normal_texture = invalid_texture_handle(),
        },
        &scene->materials);

    quad_create((vec3s){0.0f, 1.0f, 0.0f}, (vec3s){0.0f, 1.0f, 0.0f}, (vec3s){1.0f, 0.0f, 0.0f}, 10.0f, floor_material.id, &scene->triangles);
    quad_create((vec3s){0.0f, 1.5f, 0.0f}, (vec3s){1.0f, 0.0f, 0.0f}, (vec3s){0.0f, 1.0f, 0.0f}, 1.0f, quad_material.id, &scene->triangles);
#endif

    return scene_build_bvh(scene);
}

static bool initialize_renderer(const rendering_config_t* config, render_job_t** outJob, scene_t* outScene)
{
    if (!scene_setup(outScene)
        || !load_assets(outScene))
    {
        return false;
    }

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

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

        .rendering_mode = RENDER_TILE_BASED,
        .aov_flags = AOV_BEAUTY,
        .is_done = false,
    };

    if (!renderer_aov_target_init(job, job->aov_flags))
    {
        return false;
    }

    sobol_init();
    srand((unsigned int)time(NULL));

    *outJob = job;
    return true;
}

static void shutdown_renderer(render_job_t* job, scene_t* scene)
{
    render_job_free(job);
    scene_free(scene);

    free(job);
}

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, uint8_t aov_index)
{
    MSG msg;
    bool running = true;
    render_target_t* render_target = job->aov_target[aov_index];

    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(render_target);
        }
    }

    return 0;
}

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

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

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

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

    int result = run_main_loop(job, AOV_BEAUTY_INDEX);

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

    return result;
}