#pragma once

#include <vulkan/vulkan.h>
#include "SceneObject.h"
#include "TextureLoader.h"


class StartScreen : public SceneObject{
    public:
        std::vector<TextureData> m_textureData;
        std::vector<VkDescriptorSet> m_descriptorSets;

        StartScreen(VkDevice dev, VkPhysicalDevice phyDev, VkDescriptorPool descPool, VkDescriptorSetLayout descSetLayout, VkQueue queue, ShadowMap& shadowMap, glm::vec2 screenRes, TextureLoader& textureLoader)
        : SceneObject(dev, phyDev, descPool, descSetLayout, queue,  { {{-0.2f,  0.2f, 0.0f}}, {{-0.2f, -0.2f, 0.0f}}, {{ 0.2f,  0.2f, 0.0f}}, {{ 0.2f, -0.2f, 0.0f}} }, {0, 1, 2, 1, 3, 2}, 
            sampler, cmdPool, sync, pfn, shadowMap)
        {
            // Create Geometry
            GeometryData data;
            data.positions = { {-1.0f,  1.0f, 0.0f}, {-1.0f, -1.0f, 0.0f}, { 1.0f,  1.0f, 0.0f}, {1.0f, -1.0f, 0.0f} };
            data.textureCoordinates = { {0,1}, {0,0}, {1,1}, {1,0} };
            data.indices = { 0, 1, 2, 1, 3, 2 };
            data.normals = {
                {0.0f, 0.0f, 1.0f}, // Normal for all vertices (facing forward)
                {0.0f, 0.0f, 1.0f},
                {0.0f, 0.0f, 1.0f},
                {0.0f, 0.0f, 1.0f}
            };
            data.tangents = {
                {0.0f, 0.0f, 0.0f}, // Red for top-left corner
                {0.0f, 0.0f, 0.0f}, // Green for bottom-left corner
                {0.0f, 0.0f, 1.0f}, // Blue for top-right corner
                {0.0f, 0.0f, 0.0f}  // Yellow for bottom-right corner
            };
            sceneObjectGeometry = createAndUploadIntoGpuMemory(data);
            worldTransform = glm::mat4(1.0f); // identity
            localTransform = glm::mat4(1.0f);

            // Create one ubo for both descriptor sets
            uniformBuffer = vklCreateHostCoherentBufferWithBackingMemory(
                sizeof(UniformBuffer), VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT
            );

            // Create a 2 descriptorsets for 2 start screen textures
            for (int i = 0; i <= 3; i++)
            {
                // Load texture data into provided structs
                m_textureData.push_back(TextureData(dev));
                std::string filePath = "assets/textures/screen" + std::to_string(i) + ".png";
                textureLoader.load(filePath, m_textureData.back());

                // Create own descriptorset for current texture
                m_descriptorSets.push_back(allocDescriptorSet(dev, descPool, descSetLayout));
                writeDescriptorSet(dev, m_descriptorSets.back(), uniformBuffer);
                writeDescriptorSet(dev, m_descriptorSets.back(), m_textureData.back().textureImageView, m_textureData.back().textureSampler, 1, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
            }
            descriptorSet = m_descriptorSets[0];
        }

    /**
     * Updates the uniform buffer with transformation matrices and other data.
     * This function should be called every frame to update shader data based on transformations.
     */
    virtual void updateUniformBuffer(Student::CameraForGLFW& camera, const glm::vec4& windowSize) 
    {
        ubo.modelMatrix = glm::mat4(1.0f);  // This trafo controls how big the trnasparent quad is. it should be small .
        ubo.modelViewProjMatrix = glm::mat4(1.0f);  // Combine with the camera matrix
        ubo.userInput = glm::ivec4(0, windowSize.x, windowSize.y, 0 ); // last index tells that this renders textured quad
        vklCopyDataIntoHostCoherentBuffer(this->uniformBuffer, &ubo, sizeof(ubo));
    }

    void render(VkPipeline hudPipeline){
    
        // Wait until we get an image from the swapchain to render into:
        vklWaitForNextSwapchainImage();

        // Render scene to current swap chain image 
        vklStartRecordingCommands();
        drawGeometryWithMaterial(hudPipeline, this->sceneObjectGeometry, descriptorSet);
        vklEndRecordingCommands();

        // Present rendered image to the screen:
        vklPresentCurrentSwapchainImage();
    }

    void updateTextureIndex(const uint32_t& index)
    {
        // index should be 0 or 1 since we only have 2 start screens
        descriptorSet = m_descriptorSets[index];
    }
};



class ShadowScreen : public SceneObject {
public:
    std::vector<TextureData> m_textureData;
    std::vector<VkDescriptorSet> m_descriptorSets;

    ShadowScreen(VkDevice dev, VkPhysicalDevice phyDev, VkDescriptorPool descPool, VkDescriptorSetLayout descSetLayout, VkQueue queue, ShadowMap& shadowMap, glm::vec2 screenRes, TextureLoader& textureLoader)
        : SceneObject(dev, phyDev, descPool, descSetLayout, queue, { {{-0.2f,  0.2f, 0.0f}}, {{-0.2f, -0.2f, 0.0f}}, {{ 0.2f,  0.2f, 0.0f}}, {{ 0.2f, -0.2f, 0.0f}} }, { 0, 1, 2, 1, 3, 2 },
            sampler, cmdPool, sync, pfn, shadowMap)
    {
        // Create Geometry
        GeometryData data;
        data.positions = { {-1.0f,  1.0f, 0.0f}, {-1.0f, -1.0f, 0.0f}, { 1.0f,  1.0f, 0.0f}, {1.0f, -1.0f, 0.0f} };
        data.textureCoordinates = { {0,1}, {0,0}, {1,1}, {1,0} };
        data.indices = { 0, 1, 2, 1, 3, 2 };
        data.normals = {
            {0.0f, 0.0f, 1.0f}, // Normal for all vertices (facing forward)
            {0.0f, 0.0f, 1.0f},
            {0.0f, 0.0f, 1.0f},
            {0.0f, 0.0f, 1.0f}
        };
        data.tangents = {
            {0.0f, 0.0f, 0.0f}, // Red for top-left corner
            {0.0f, 0.0f, 0.0f}, // Green for bottom-left corner
            {0.0f, 0.0f, 1.0f}, // Blue for top-right corner
            {0.0f, 0.0f, 0.0f}  // Yellow for bottom-right corner
        };
        sceneObjectGeometry = createAndUploadIntoGpuMemory(data);
        worldTransform = glm::mat4(1.0f); // identity
        localTransform = glm::mat4(1.0f);

        // Create one ubo for both descriptor sets
        uniformBuffer = vklCreateHostCoherentBufferWithBackingMemory(
            sizeof(UniformBuffer), VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT
        );

        // Create a 2 descriptorsets for 2 start screen textures
        for (int i = 0; i <=0; i++)
        {
            m_descriptorSets.push_back(allocDescriptorSet(dev, descPool, descSetLayout));
            writeDescriptorSet(dev, m_descriptorSets.back(), uniformBuffer);
            // Bind shadowmap to binding 1
            writeDescriptorSet(dev, m_descriptorSets.back(),shadowMap.m_image_view, shadowMap.m_image_sampler,1);
        }
        descriptorSet = m_descriptorSets[0];
    }

    /**
     * Updates the uniform buffer with transformation matrices and other data.
     * This function should be called every frame to update shader data based on transformations.
     */
    virtual void updateUniformBuffer(Student::CameraForGLFW& camera, const glm::vec4& windowSize)
    {
        ubo.modelMatrix = glm::mat4(1.0f);  // This trafo controls how big the trnasparent quad is. it should be small .
        ubo.modelViewProjMatrix = glm::mat4(1.0f);  // Combine with the camera matrix
        ubo.userInput = glm::ivec4(0, windowSize.x, windowSize.y, 0); // last index tells that this renders textured quad
        vklCopyDataIntoHostCoherentBuffer(this->uniformBuffer, &ubo, sizeof(ubo));
    }

    void render(VkPipeline hudPipeline) {

        // Wait until we get an image from the swapchain to render into:
        vklWaitForNextSwapchainImage();

        // Render scene to current swap chain image 
        vklStartRecordingCommands();
        drawGeometryWithMaterial(hudPipeline, this->sceneObjectGeometry, descriptorSet);
        vklEndRecordingCommands();

        // Present rendered image to the screen:
        vklPresentCurrentSwapchainImage();
    }

    void updateTextureIndex(const uint32_t& index)
    {
        // index should be 0 or 1 since we only have 2 start screens
        descriptorSet = m_descriptorSets[index];
    }
};



