package net.halbear.Terrain4J.EngineCore.Display;

import net.halbear.Terrain4J.EngineCore.Logic.EngineInstance;
import net.halbear.Terrain4J.EngineCore.Logic.InitData;
import net.halbear.Terrain4J.EngineCore.Logic.Rendering.VulkanContext;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Pipeline.Images.TextureCache;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.VkModel.MaterialData;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.VkModel.MaterialsCache;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.VkModel.ModelData;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.VkModel.ModelsCache;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.CommandBuffer;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.CommandPool;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.GPUSynchronisation.Fence;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.GPUSynchronisation.Semaphore;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.Queues.Queue;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DisplayToScreen.SceneRender;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DisplayToScreen.SwapChain;
import net.halbear.Terrain4J.EngineCore.Vulkan.Util.VulkanUtils;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.vulkan.VkCommandBufferSubmitInfo;
import org.lwjgl.vulkan.VkExtent2D;
import org.lwjgl.vulkan.VkSemaphoreSubmitInfo;
import org.tinylog.Logger;

import java.util.Arrays;
import java.util.List;

import static org.lwjgl.vulkan.VK13.VK_PIPELINE_STAGE_2_BOTTOM_OF_PIPE_BIT;
import static org.lwjgl.vulkan.VK13.VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT;

public class Render {

    private final ModelsCache modelsCache;

    private final CommandBuffer[] CommandBuffers;
    private final CommandPool[] CommandPools;
    private final Fence[] Fences;
    private final Queue.GraphicsQueue GraphicsQueue;
    private final Semaphore[] PresentCompleteSemaphores;
    private final Queue.PresentQueue PresentQueue;
    private final Semaphore[] RenderCompleteSemaphores;
    private final SceneRender sceneRender;
    private int CurrentFrame;
    private final VulkanContext RendererContext;
    private boolean Resize = false;

    private final MaterialsCache materialsCache;
    private TextureCache textureCache;

    public Render(EngineInstance engineInstance) {
        RendererContext = new VulkanContext(engineInstance.window());
        CurrentFrame = 0;
        GraphicsQueue = new Queue.GraphicsQueue(RendererContext, 0);
        PresentQueue = new Queue.PresentQueue(RendererContext, 0);
        CommandPools= new CommandPool[VulkanUtils.MAX_IN_FLIGHT];
        CommandBuffers= new CommandBuffer[VulkanUtils.MAX_IN_FLIGHT];
        Fences= new Fence[VulkanUtils.MAX_IN_FLIGHT];
        PresentCompleteSemaphores = new Semaphore[VulkanUtils.MAX_IN_FLIGHT];
        int SwapChainImageCount = RendererContext.GetSwapChain().GetImageCount();
        RenderCompleteSemaphores = new Semaphore[SwapChainImageCount];
        for(int i = 0; i < VulkanUtils.MAX_IN_FLIGHT; i++){
            CommandPools[i] = new CommandPool(RendererContext, GraphicsQueue.GetQueueFamilyIndex(), false);
            CommandBuffers[i] = new CommandBuffer(RendererContext, CommandPools[i],true, true);
            PresentCompleteSemaphores[i] = new Semaphore(RendererContext);
            Fences[i] = new Fence(RendererContext, true);
        }
        for(int i = 0; i < SwapChainImageCount; i++){
            RenderCompleteSemaphores[i] = new Semaphore(RendererContext);
        }
        sceneRender = new SceneRender(RendererContext);
        textureCache = new TextureCache();
        materialsCache = new MaterialsCache();
        modelsCache = new ModelsCache();
        Resize = false;
    }

    public void Initialise(InitData initData){

        List<MaterialData> Materials = initData.Materials();
        Logger.debug("Loading {} Materials", Materials.size());
        materialsCache.LoadMaterials(RendererContext, Materials,textureCache, CommandPools[0], GraphicsQueue);
        Logger.debug("Loaded {} Materials", Materials.size());

        Logger.debug("Transitioning Textures");
        textureCache.TransitionTexts(RendererContext, CommandPools[0], GraphicsQueue);
        Logger.debug("Transitioned Textures");

        List<ModelData> Models = initData.Models();
        Logger.debug("Loading {} models", Models.size());
        modelsCache.loadModels(RendererContext, Models, CommandPools[0], GraphicsQueue);
        Logger.debug("Loaded {} models", Models.size());
        sceneRender.LoadMaterials(RendererContext,materialsCache,textureCache);
    }

    private void RecordingStart(CommandPool commandPool, CommandBuffer commandBuffer){
        commandPool.Reset(RendererContext);
        commandBuffer.BeginRecording();
    }

    private void RecordingStop(CommandBuffer commandBuffer){
        commandBuffer.EndRecording();
    }

    public void cleanup() {
        RendererContext.GetDevice().waitIdle();
        Logger.debug("Waiting Vulkan Context");
        sceneRender.cleanup(RendererContext);
        Logger.debug("Dynamic Renderer Cleaned up");
        Arrays.asList(RenderCompleteSemaphores).forEach(i->i.cleanup(RendererContext));
        Logger.debug("Render Semaphores cleaned up");
        Arrays.asList(PresentCompleteSemaphores).forEach(i->i.cleanup(RendererContext));
        Logger.debug("Present Semaphores cleaned up");
        Arrays.asList(Fences).forEach(i->i.cleanup(RendererContext));
        Logger.debug("Fences cleaned up");
        for(int i = 0; i < CommandPools.length; i++){
            CommandBuffers[i].cleanup(RendererContext, CommandPools[i]);
            CommandPools[i].cleanup(RendererContext);
        }
        Logger.debug("Command pools cleaned up");
        modelsCache.CleanUp(RendererContext);
        materialsCache.CleanUp(RendererContext);
        textureCache.CleanUp(RendererContext);
        RendererContext.cleanup();
    }
    public void render(EngineInstance engineInstance){
        SwapChain swapChain = RendererContext.GetSwapChain();
        WaitForFence(CurrentFrame); //problem child found
        var CommandPool = CommandPools[CurrentFrame];
        var CommandBuffer = CommandBuffers[CurrentFrame];
        RecordingStart(CommandPool, CommandBuffer);
        int ImageIndex;// = swapChain.FetchNextImage(RendererContext.GetDevice(), PresentCompleteSemaphores[CurrentFrame]);
        if (Resize || (ImageIndex = swapChain.FetchNextImage(RendererContext.GetDevice(), PresentCompleteSemaphores[CurrentFrame])) < 0){
            resize(engineInstance);
            return;
        }
        sceneRender.Render(engineInstance,RendererContext,CommandBuffer, modelsCache,materialsCache,ImageIndex,CurrentFrame);
        RecordingStop(CommandBuffer);
        Submit(CommandBuffer, CurrentFrame, ImageIndex);
        Resize = swapChain.PresentImage(PresentQueue, RenderCompleteSemaphores[ImageIndex],ImageIndex);
        CurrentFrame = (CurrentFrame + 1) % VulkanUtils.MAX_IN_FLIGHT;
    }

    private void resize(EngineInstance engineInstance){
        Window window = engineInstance.window();
        if(window.getWidth() == 0 && window.getHeight() == 0){
            return;
        }
        Resize = false;
        RendererContext.GetDevice().waitIdle();
        RendererContext.Resize(window);

        Arrays.asList(RenderCompleteSemaphores).forEach(i->i.cleanup(RendererContext));
        Arrays.asList(PresentCompleteSemaphores).forEach(i->i.cleanup(RendererContext));
        for(int i = 0; i < VulkanUtils.MAX_IN_FLIGHT; i++){
                PresentCompleteSemaphores[i] = new Semaphore(RendererContext);
        }
        for(int i = 0; i < RendererContext.GetSwapChain().GetImageCount(); i++){
            RenderCompleteSemaphores[i] = new Semaphore(RendererContext);
        }

        VkExtent2D extend = RendererContext.GetSwapChain().GetSwapChainExtent();
        engineInstance.scene().GetProjection().Resize(extend.width(),extend.height());
        sceneRender.Resize(engineInstance,RendererContext);
    }

    private void ResetFence(int currentFrame){
        var fence = Fences[currentFrame];
        fence.Reset(RendererContext);
    }
    private void WaitForFence(int currentFrame){
        var fence = Fences[currentFrame];
        fence.FenceWait(RendererContext);
    }

    public void Submit(CommandBuffer commandBuffer, int currentFrame, int ImageIndex){
        try(var MemStack = MemoryStack.stackPush()){
            var fence = Fences[currentFrame];
            fence.Reset(RendererContext);
            var Commands = VkCommandBufferSubmitInfo.calloc(1,MemStack)
                    .sType$Default()
                    .commandBuffer(commandBuffer.GetVulkanCommandBuffer());
            VkSemaphoreSubmitInfo.Buffer SemaphoreWait = VkSemaphoreSubmitInfo.calloc(1,MemStack)
                    .sType$Default()
                    .stageMask(VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT)
                    .semaphore(PresentCompleteSemaphores[currentFrame].GetSemaphore());
            VkSemaphoreSubmitInfo.Buffer SemaphoreSignals = VkSemaphoreSubmitInfo.calloc(1, MemStack)
                    .sType$Default()
                    .stageMask(VK_PIPELINE_STAGE_2_BOTTOM_OF_PIPE_BIT)
                    .semaphore(RenderCompleteSemaphores[ImageIndex].GetSemaphore());
            GraphicsQueue.Submit(Commands, SemaphoreWait, SemaphoreSignals, fence);
        }
    }
}
