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

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

static void save_img(render_target_t* source,  uint32_t width,  uint32_t height, const char* filename)
{
    FILE* file_stream;
    fopen_s(&file_stream, filename, "wb");
    if (file_stream == NULL)
    {
        fprintf(stderr, "Failed to open file for writing: %s\n", filename);
        return;
    }

    unsigned char* img_buffer = render_target_to_char(source);
    svpng(file_stream, width, height, img_buffer, 1);

    fclose(file_stream);
    free(img_buffer);
}

// int main()
int WINAPI wWinMain(_In_ HINSTANCE hInstance, _In_opt_ HINSTANCE hPrevInstance, _In_ PWSTR pCmdLine, _In_ int nCmdShow)
{
    rendering_config_t config = {
        .width = 1920 / 1,
        .height = 1080 / 1,
        .sample_count = 64,
        .max_depth = 4,
        .bucket_size = 64,
    };

    int result = window_create("Path Tracing", hInstance, config.width, config.height);
    if (result != 0)
    {
        fprintf(stderr, "Failed to create window\n");
        return -1;
    }

#pragma region SceneSetup
    omp_set_num_threads(16);
    sobol_init();

    scene_t scene;
    if (!scene_init(67000, 8, 1, &scene))
    {
        return -1;
    }

    scene.camera.position = (vec3s){-7.5f, 2.0f, 0.0f};
    scene.camera.forward = glms_vec3_normalize((vec3s){1.0f, 0.0f, 0.0f});

    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.15f});
    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,
    });

    simple_lit_data_t gray_lit_data =
    {
        .albedo = (vec3s){0.73f, 0.73f, 0.73f},
        .roughness = 0.5f,
        .metallic = 0.0f,
    };
    simple_lit_data_t blue_lit_data =
    {
        .albedo = (vec3s){0.0f, 0.0f, 1.0f},
        .roughness = 0.5f,
        .metallic = 0.0f,
    };
    simple_lit_data_t red_lit_data =
    {
        .albedo = (vec3s){1.0f, 0.0f, 0.0f},
        .roughness = 0.5f,
        .metallic = 0.0f,
    };
    simple_lit_data_t green_lit_data =
    {
        .albedo = (vec3s){0.0f, 1.0f, 0.0f},
        .roughness = 0.5f,
        .metallic = 0.0f,
    };
    simple_lit_data_t floor_lit_data =
    {
        .albedo = (vec3s){1.0f, 1.0f, 1.0f},
        .roughness = 0.95f,
        .metallic = 0.0f,
    };
    material_entity_t gray_material = material_create_simple_lit(&gray_lit_data, &scene.materials);
    material_entity_t gray_light_material = material_create_simple_lit(&gray_lit_data, &scene.materials);
    material_entity_t blue_material = material_create_simple_lit(&blue_lit_data, &scene.materials);
    material_entity_t red_material = material_create_simple_lit(&red_lit_data, &scene.materials);
    material_entity_t green_material = material_create_simple_lit(&green_lit_data, &scene.materials);
    material_entity_t floor_material = material_create_simple_lit(&floor_lit_data, &scene.materials);

    //scene.materials.buffer[gray_light_material.id].emission = (vec3s){10.0f, 10.0f, 10.0f};

    mesh_load("F:/c/SimpleRayTracer/assets/sponza.obj", floor_material.id, &scene.triangles, &scene.materials);

    // quad_create(
    //     (vec3s){0.0f, 3.95f, 0.0f},
    //     (vec3s){0.0f, -1.0f, 0.0f},
    //     (vec3s){0.0f, 0.0f, 1.0f},
    //     1.0f, gray_light_material.id, &scene.triangles
    // );
    // quad_create(
    //     (vec3s){0.0f, 0.0f, 0.0f},
    //     (vec3s){0.0f, 1.0f, 0.0f},
    //     (vec3s){0.0f, 0.0f, 1.0f},
    //     4.0f, floor_material.id, &scene.triangles
    // );
    // quad_create(
    //     (vec3s){0.0f, 4.0f, 0.0f},
    //     (vec3s){0.0f, -1.0f, 0.0f},
    //     (vec3s){0.0f, 0.0f, 1.0f},
    //     4.0f, gray_material.id, &scene.triangles
    // );
    // quad_create(
    //     (vec3s){0.0f, 2.0f, -2.0f},
    //     (vec3s){0.0f, 0.0f, 1.0f},
    //     (vec3s){0.0f, 1.0f, 0.0f},
    //     4.0f, green_material.id, &scene.triangles
    // );
    // quad_create(
    //     (vec3s){-2.0f, 2.0f, 0.0f},
    //     (vec3s){1.0f, 0.0f, 0.0f},
    //     (vec3s){0.0f, 1.0f, 0.0f},
    //     4.0f, blue_material.id, &scene.triangles
    // );
    // quad_create(
    //     (vec3s){2.0f, 2.0f, 0.0f},
    //     (vec3s){-1.0f, 0.0f, 0.0f},
    //     (vec3s){0.0f, 1.0f, 0.0f},
    //     4.0f, red_material.id, &scene.triangles
    // );

#pragma endregion

    render_target_t img;
    render_target_init(config.width, config.height, &img);

    if (!scene_build_bvh(&scene))
    {
        fprintf(stderr, "Failed to build BVH\n");
        goto exit;
    }

    // BVH Debug
    for (uint32_t i = 0; i < scene.bvh_tree.node_count; i++)
    {
        bvh_node_t* node = &scene.bvh_tree.nodes[i];
        aabb_t bounds = node->bounds;
        vec3s min = bounds.min;
        vec3s max = bounds.max;
        printf("Node %u: Min: (%f, %f, %f), Max: (%f, %f, %f)\n", i, min.x, min.y, min.z, max.x, max.y, max.z);
    }

    MSG msg;
    uint16_t tile_index = 0;
    tile_t current_tile = {0};
    rendering_context_t ctx = {0};

    while (true)
    {
        while (PeekMessage(&msg, NULL, 0, 0, PM_REMOVE))
        {
            if (msg.message == WM_QUIT)
            {
                goto exit;
            }

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

        if (scene_render_tile(&scene, &ctx, config, tile_index, DEBUG_BVH, &img, &current_tile))
        {
            for (uint32_t y = current_tile.y; y < current_tile.y + current_tile.height; y++)
            {
                for (uint32_t x = current_tile.x; x < current_tile.x + current_tile.width; x++)
                {
                    vec4s pixel = render_target_get_pixel(&img, x, y);
                    pixel = gamma_correct(pixel, 2.2f);
                    pixel = aces_tone_map(pixel);
                    window_update_pixels(pixel, x, y);
                }
            }

            window_refresh_region(current_tile.x, current_tile.y, current_tile.width, current_tile.height);
            tile_index++;
        }
    }

exit:
    render_target_free(&img);
    scene_free(&scene);

    return 0;
}