#include "scene.hpp"
#include "mesh.hpp"

// void scene::random_scene() {
//     material ground_material = { { 0.5, 0.5, 0.5 } };
//     spheres.emplace_back(vec3 { 0, -1000, 0 }, ground_material, 1000);
// 
//     int32_t spheres_count = 11;
//     for (int a = -spheres_count; a < spheres_count; a++) {
//         for (int b = -spheres_count; b < spheres_count; b++) {
//             auto choose_mat = randd();
//             point3 center(a + 0.9 * randd(), 0.2, b + 0.9 * randd());
// 
//             if ((center - point3{ 4, 0.2, 0 }).length() > 0.9) {
//                 material sphere_material = {};
// 
//                 if (choose_mat < 0.8) {
//                     // diffuse
//                     sphere_material.base_color = color::random() * color::random();
//                 } else if (choose_mat < 0.95) {
//                     // metal
//                     sphere_material.base_color = color::random(0.5, 1);
//                     sphere_material.metalness = randd(0, 0.5);
//                     sphere_material.roughness = 0.f;
//                 } else {
//                     // glass
//                     // sphere_material.ior = 1.5;
//                     sphere_material.metalness = 1.f;
//                     sphere_material.roughness = 0.f;
//                 }
// 
//                 spheres.emplace_back(center, sphere_material, 0.2);
//             }
//         }
//     }
// 
//     material material1 = {};
//     // material1.ior = 1.5;
//     spheres.emplace_back(vec3 { 0, 1, 0 }, material1, 1.0);
// 
//     material material2 = {};
//     material2.base_color = color{ 0.4, 0.2, 0.1 };
//     spheres.emplace_back(vec3 { -4, 1, 0 }, material2, 1.0);
// 
//     material material3 = {};
//     material3.base_color = color{ 0.7, 0.6, 0.5 };
//     material3.roughness = 0.0;
//     spheres.emplace_back(vec3 { 4, 1, 0 }, material3, 1.0);
// 
//     spheres_count = spheres.size();
// }


metadata::metadata(const camera &cam, u16 width, u16 height)
    : cam(cam), width(width), height(height) {}

scene::scene(const camera& cam, u16 width, u16 height)
    // :model("../models/sponza/Sponza.gltf"), meta(cam, width, height) {
    // :model("../models/BistroInterior/BistroInterior.gltf"), meta(cam, width, height) {
    // :model("../models/BistroExterior/BistroExterior.gltf"), meta(cam, width, height) {
    // :model("../models/main_sponza/NewSponza_Main_glTF_003.gltf"), meta(cam, width, height) {
    // :model("../models/pkg_a_curtains/NewSponza_Curtains_glTF.gltf"), meta(cam, width, height) {
    :model("../models/sponza.glb", true), meta(cam, width, height) {

    // Geometry & BVH
    std::vector<uint32_t>       indices;
    std::vector<float>          positions;
    std::vector<float>          normals;
    std::vector<float>          uvs;
    std::vector<gpu_material>   materials;
    std::vector<packed_bvh_node> packed_nodes;

    std::vector<triangle> triangles;
    u32 triangles_offset = 0U;
    u32 vertices_offset = 0U;
    u32 positions_count = 0U;
    u32 normals_count = 0U;
    u32 uvs_count = 0U;
    std::vector<u32> nodes_to_load = model.root_nodes;

    for (u32 cur_node_index = 0U; cur_node_index < nodes_to_load.size(); cur_node_index++) {
        struct node &node = model.nodes[nodes_to_load[cur_node_index]];

        if (node.mesh != UINT32_MAX) {
            const Mesh& mesh = model.meshes[node.mesh];
            const auto& mesh_positions = mesh.get_attribute(ATTRIBUTE_TYPE::POSITION);
            const auto& mesh_normals = mesh.get_attribute(ATTRIBUTE_TYPE::NORMAL);
            const auto& mesh_uvs_0 = mesh.get_attribute(ATTRIBUTE_TYPE::UV_0);
            const auto& mesh_indices = mesh.get_indices();

            indices.resize(indices.size() + mesh_indices.size());
            positions.resize(positions_count + mesh_positions.size());
            normals.resize(normals_count + mesh_normals.size());
            uvs.resize(uvs_count + mesh_uvs_0.size());

            std::memcpy(
               positions.data() + positions_count,
               mesh_positions.data(),
               mesh_positions.size() * sizeof(mesh_positions[0U])
            );

            std::memcpy(
               normals.data() + normals_count,
               mesh_normals.data(),
               mesh_normals.size() * sizeof(mesh_normals[0U])
            );

            std::memcpy(
                uvs.data() + uvs_count,
                mesh_uvs_0.data(),
                mesh_uvs_0.size() * sizeof(mesh_uvs_0[0U])
            );

            auto index_offset { 0U };
            // Add offset to indices in order to have absolute indices
            for (const auto& submesh: mesh.get_submeshes()) {
                auto vertex_offset = vertices_offset + submesh.vertex_offset;
                const auto& material = submesh.material;

                for (size_t element_index { 0U }; element_index < submesh.element_count; element_index++, index_offset += 3) {
                    size_t submesh_level_index_1 = mesh_indices[index_offset];
                    size_t submesh_level_index_2 = mesh_indices[index_offset + 1];
                    size_t submesh_level_index_3 = mesh_indices[index_offset + 2];

                    size_t mesh_level_index_1 = (submesh.vertex_offset + submesh_level_index_1) * 3;
                    size_t mesh_level_index_2 = (submesh.vertex_offset + submesh_level_index_2) * 3;
                    size_t mesh_level_index_3 = (submesh.vertex_offset + submesh_level_index_3) * 3;

                    triangles.emplace_back(
                        vec3 { mesh_positions[mesh_level_index_1], mesh_positions[mesh_level_index_1 + 1], mesh_positions[mesh_level_index_1 + 2] },
                        vec3 { mesh_positions[mesh_level_index_2], mesh_positions[mesh_level_index_2 + 1], mesh_positions[mesh_level_index_2 + 2] },
                        vec3 { mesh_positions[mesh_level_index_3], mesh_positions[mesh_level_index_3 + 1], mesh_positions[mesh_level_index_3 + 2] }
                    );

                    auto triangle_offset = triangles_offset * 3 + index_offset;

                    indices[triangle_offset]        = static_cast<uint32_t>((submesh_level_index_1 + vertex_offset) | (0xff000000 & (materials.size() << 8)));
                    indices[triangle_offset + 1]    = static_cast<uint32_t>((submesh_level_index_2 + vertex_offset) | (0xff000000 & (materials.size() << 16)));
                    indices[triangle_offset + 2]    = static_cast<uint32_t>((submesh_level_index_3 + vertex_offset) | (0xff000000 & (materials.size() << 24)));
                }

                auto& mat = materials.emplace_back(gpu_material {
                    .base_color = material.base_color,
//                     .albedo_texture_id = albedo_image.bindless_sampled_index,
//                     .albedo_texture_sampler_id = albedo_sampler.bindless_index,
                    .metalness = material.metalness,
                    .roughness = material.roughness
                });

                if (material.base_color_texture.id != handle<device_texture>::invalid_id) {
                    mat.albedo_texture_id = vkdevice::get().get_texture(material.base_color_texture)->get_sampled_index();
                    // mat.albedo_texture_sampler_id = device.get_sampler(0);
                }

                if (material.metallic_roughness_texture.id != handle<device_texture>::invalid_id) {
                    mat.metallic_roughness_texture_id = vkdevice::get().get_texture(material.metallic_roughness_texture)->get_sampled_index();
                    // mat.metallic_roughness_texture_sampler_id = device.get_sampler(0);
                }
            }

            triangles_offset += mesh.triangle_count();
            vertices_offset += mesh.vertices_count();
            positions_count = vertices_offset * Mesh::attribute_components(ATTRIBUTE_TYPE::POSITION);
            normals_count = vertices_offset * Mesh::attribute_components(ATTRIBUTE_TYPE::NORMAL);
            uvs_count = vertices_offset * Mesh::attribute_components(ATTRIBUTE_TYPE::UV_0);
        }

        for (u32 child_index: node.children) {
            nodes_to_load.push_back(child_index);
        }
    }

    // random_scene();

    bvh builder(triangles, packed_nodes);
    // bvh builder(spheres, packed_nodes);

    {
        indices_buffer_handle = vkdevice::get().create_buffer({
            .size = indices.size() * sizeof(indices[0]),
            .usages = VK_BUFFER_USAGE_2_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_2_SHADER_DEVICE_ADDRESS_BIT,
            .memory_properties = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
            .memory_usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
            .name = "indices",
        });
        const auto& indices_buffer = vkdevice::get().get_buffer(indices_buffer_handle);
        void* ptr = indices_buffer->mapped_ptr;
        memcpy(ptr, indices.data(), indices.size() * sizeof(indices[0]));
    }

    {
        positions_buffer_handle = vkdevice::get().create_buffer({
            .size = positions.size() * sizeof(positions[0]),
            .usages = VK_BUFFER_USAGE_2_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_2_SHADER_DEVICE_ADDRESS_BIT,
            .memory_properties = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
            .memory_usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
            .name = "positions",
        });
        const auto& positions_buffer = vkdevice::get().get_buffer(positions_buffer_handle);
        void* ptr = positions_buffer->mapped_ptr;
        memcpy(ptr, positions.data(), positions.size() * sizeof(positions[0]));
    }

    {
        normals_buffer_handle = vkdevice::get().create_buffer({
            .size = normals.size() * sizeof(normals[0]),
            .usages = VK_BUFFER_USAGE_2_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_2_SHADER_DEVICE_ADDRESS_BIT,
            .memory_properties = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
            .memory_usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
            .name = "normals",
        });
        const auto& normals_buffer = vkdevice::get().get_buffer(normals_buffer_handle);
        void* ptr = normals_buffer->mapped_ptr;
        memcpy(ptr, normals.data(), normals.size() * sizeof(normals[0]));
    }

    {
        uvs_buffer_handle = vkdevice::get().create_buffer({
            .size = uvs.size() * sizeof(uvs[0]),
            .usages = VK_BUFFER_USAGE_2_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_2_SHADER_DEVICE_ADDRESS_BIT,
            .memory_properties = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
            .memory_usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
            .name = "uvs",
        });
        const auto& uvs_buffer = vkdevice::get().get_buffer(uvs_buffer_handle);
        void* ptr = uvs_buffer->mapped_ptr;
        memcpy(ptr, uvs.data(), uvs.size() * sizeof(uvs[0]));
    }

    {
        bvh_buffer_handle = vkdevice::get().create_buffer({
            .size = packed_nodes.size() * sizeof(packed_nodes[0]),
            .usages = VK_BUFFER_USAGE_2_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_2_SHADER_DEVICE_ADDRESS_BIT,
            .memory_properties = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
            .memory_usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
            .name = "bvh",
        });
        const auto& bvh_buffer = vkdevice::get().get_buffer(bvh_buffer_handle);
        void* ptr = bvh_buffer->mapped_ptr;
        memcpy(ptr, packed_nodes.data(), packed_nodes.size() * sizeof(packed_nodes[0]));
    }

    {
        materials_buffer_handle = vkdevice::get().create_buffer({
            .size = materials.size() * sizeof(materials[0]),
            .usages = VK_BUFFER_USAGE_2_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_2_SHADER_DEVICE_ADDRESS_BIT,
            .memory_properties = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
            .memory_usage = VMA_MEMORY_USAGE_CPU_TO_GPU,
            .name = "materials",
        });
        const auto& materials_buffer = vkdevice::get().get_buffer(materials_buffer_handle);
        void* ptr = materials_buffer->mapped_ptr;
        memcpy(ptr, materials.data(), materials.size() * sizeof(materials[0]));
    }
}

void scene::deinit() {
    vkdevice::get().destroy_buffer(indices_buffer_handle);
    vkdevice::get().destroy_buffer(positions_buffer_handle);
    vkdevice::get().destroy_buffer(normals_buffer_handle);
    vkdevice::get().destroy_buffer(uvs_buffer_handle);
    vkdevice::get().destroy_buffer(bvh_buffer_handle);
    vkdevice::get().destroy_buffer(materials_buffer_handle);

    model.deinit();
}