/* Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
 *
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 the "License";
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*
 * Timestamp queries (based on the HDR sample), using vulkan.hpp
 */

#pragma once

#include <hpp_api_vulkan_sample.h>

class HPPTimestampQueries : public HPPApiVulkanSample
{
  public:
	HPPTimestampQueries();
	~HPPTimestampQueries();

  private:
	struct Bloom
	{
		bool                    enabled               = true;
		vk::DescriptorSetLayout descriptor_set_layout = {};
		vk::DescriptorSet       descriptor_set;
		vk::PipelineLayout      pipeline_layout;
		vk::Pipeline            pipelines[2];

		void destroy(vk::Device device)
		{
			device.destroyPipeline(pipelines[0]);
			device.destroyPipeline(pipelines[1]);
			device.destroyPipelineLayout(pipeline_layout);
			device.destroyDescriptorSetLayout(descriptor_set_layout);
			// no need to free the descriptor_set, as it's implicitly free'd with the descriptor_pool
		}
	};

	struct Composition
	{
		vk::DescriptorSetLayout descriptor_set_layout = {};
		vk::DescriptorSet       descriptor_set        = {};
		vk::PipelineLayout      pipeline_layout       = {};
		vk::Pipeline            pipeline              = {};

		void destroy(vk::Device device)
		{
			device.destroyPipeline(pipeline);
			device.destroyPipelineLayout(pipeline_layout);
			device.destroyDescriptorSetLayout(descriptor_set_layout);
			// no need to free the descriptor_set, as it's implicitly free'd with the descriptor_pool
		}
	};

	// Framebuffer for offscreen rendering
	struct FramebufferAttachment
	{
		vk::Format       format = {};
		vk::Image        image  = {};
		vk::DeviceMemory mem    = {};
		vk::ImageView    view   = {};

		void destroy(vk::Device device)
		{
			device.destroyImageView(view);
			device.destroyImage(image);
			device.freeMemory(mem);
		}
	};

	struct FilterPass
	{
		vk::Extent2D          extent      = {};
		vk::Framebuffer       framebuffer = {};
		FramebufferAttachment color       = {};
		vk::RenderPass        render_pass = {};
		vk::Sampler           sampler     = {};

		void destroy(vk::Device device)
		{
			device.destroySampler(sampler);
			device.destroyFramebuffer(framebuffer);
			device.destroyRenderPass(render_pass);
			color.destroy(device);
		}
	};

	struct Geometry
	{
		vk::DescriptorSet                                                      descriptor_set = {};
		vk::Pipeline                                                           pipeline       = {};
		std::vector<std::unique_ptr<vkb::scene_graph::components::HPPSubMesh>> meshes         = {};

		void destroy(vk::Device device, vk::DescriptorPool descriptor_pool)
		{
			// no need to free the descriptor_set, as it's implicitly free'd with the descriptor_pool
			device.destroyPipeline(pipeline);
		}
	};

	struct Models
	{
		vk::DescriptorSetLayout descriptor_set_layout = {};
		vk::PipelineLayout      pipeline_layout       = {};
		Geometry                objects               = {};
		Geometry                skybox                = {};
		std::vector<glm::mat4>  transforms            = {};
		int32_t                 object_index          = 0;

		void destroy(vk::Device device, vk::DescriptorPool descriptor_pool)
		{
			objects.destroy(device, descriptor_pool);
			skybox.destroy(device, descriptor_pool);
			device.destroyPipelineLayout(pipeline_layout);
			device.destroyDescriptorSetLayout(descriptor_set_layout);
		}
	};

	struct Offscreen
	{
		vk::Extent2D          extent      = {};
		vk::Framebuffer       framebuffer = {};
		FramebufferAttachment color[2]    = {};
		FramebufferAttachment depth       = {};
		vk::RenderPass        render_pass = {};
		vk::Sampler           sampler     = {};

		void destroy(vk::Device device)
		{
			device.destroySampler(sampler);
			device.destroyFramebuffer(framebuffer);
			device.destroyRenderPass(render_pass);
			color[0].destroy(device);
			color[1].destroy(device);
			depth.destroy(device);
		}
	};

	struct Textures
	{
		HPPTexture envmap;

		void destroy(vk::Device device)
		{
			device.destroySampler(envmap.sampler);
		}
	};

	struct TimeStamps
	{
		std::array<uint64_t, 6> values;            // GPU time stamps will be stored in an array
		vk::QueryPool           query_pool;        // A query pool is required to use GPU time stamps

		void destroy(vk::Device device)
		{
			device.destroyQueryPool(query_pool);
		}
	};

	struct UBOMatrices
	{
		glm::mat4 projection;
		glm::mat4 modelview;
		glm::mat4 skybox_modelview;
		glm::mat4 inverse_modelview;
		float     modelscale = 0.05f;
	};

	struct UBOParams
	{
		float exposure = 1.0f;
	};

	struct UniformBuffers
	{
		std::unique_ptr<vkb::core::BufferCpp> matrices;
		std::unique_ptr<vkb::core::BufferCpp> params;
	};

  private:
	// from vkb::Application
	virtual bool prepare(const vkb::ApplicationOptions &options) override;
	virtual bool resize(const uint32_t width, const uint32_t height) override;

	// from vkb::VulkanSample
	virtual void request_gpu_features(vkb::core::HPPPhysicalDevice &gpu) override;

	// from HPPApiVulkanSample
	virtual void build_command_buffers() override;
	virtual void on_update_ui_overlay(vkb::Drawer &drawer) override;
	virtual void render(float delta_time) override;

	vk::DeviceMemory      allocate_memory(vk::Image image);
	FramebufferAttachment create_attachment(vk::Format format, vk::ImageUsageFlagBits usage);
	vk::DescriptorPool    create_descriptor_pool();
	vk::Pipeline          create_bloom_pipeline(uint32_t direction);
	vk::Pipeline          create_composition_pipeline();
	vk::RenderPass        create_filter_render_pass();
	vk::Image             create_image(vk::Format format, vk::ImageUsageFlagBits usage);
	vk::Pipeline          create_models_pipeline(uint32_t shaderType, vk::CullModeFlagBits cullMode, bool depthTestAndWrite);
	vk::RenderPass        create_offscreen_render_pass();
	vk::RenderPass        create_render_pass(std::vector<vk::AttachmentDescription> const &attachment_descriptions, vk::SubpassDescription const &subpass_description);
	void                  draw();
	void                  get_time_stamp_results();
	void                  load_assets();
	void                  prepare_bloom();
	void                  prepare_camera();
	void                  prepare_composition();
	void                  prepare_models();
	void                  prepare_offscreen_buffer();
	void                  prepare_time_stamps();
	void                  prepare_uniform_buffers();
	void                  update_composition_descriptor_set();
	void                  update_bloom_descriptor_set();
	void                  update_model_descriptor_set(vk::DescriptorSet descriptor_set);
	void                  update_params();
	void                  update_uniform_buffers();

  private:
	Bloom                    bloom;
	Composition              composition;
	bool                     display_skybox = true;
	FilterPass               filter_pass;
	Models                   models;
	std::vector<std::string> object_names;
	Offscreen                offscreen;
	Textures                 textures;
	TimeStamps               time_stamps;
	UBOMatrices              ubo_matrices;
	UBOParams                ubo_params;
	UniformBuffers           uniform_buffers;
};

std::unique_ptr<vkb::Application> create_hpp_timestamp_queries();