package net.halbear.Terrain4J.EngineCore.Display;

import net.halbear.Terrain4J.EngineCore.Logic.EngineConfig;
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.GUI.GuiRenderer;
import net.halbear.Terrain4J.EngineCore.Vulkan.GUI.GuiTexture;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Pipeline.Images.TextureCache;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.PostProcessing.PostProcess;
import net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.VkModel.*;
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.SceneRenderer;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.SwapChain.SwapChain;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.SwapChain.SwapChainRender;
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.ArrayList;
import java.util.Arrays;
import java.util.List;

import static org.lwjgl.vulkan.VK13.*;

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 SceneRenderer sceneRender;
    private final PostProcess PostProcessor;
    private final SwapChainRender swapChainRender;
    private final GuiRenderer GuiRender;
    private int CurrentFrame;
    private final VulkanContext RendererContext;
    private boolean Resize = false;
    private List<MaterialData> Materials = new ArrayList<>();
    private List<ModelData> Models = new ArrayList<>();
    private final MaterialsCache materialsCache;
    private TextureCache textureCache;
    private EngineConfig.AntiAliasType CurrentAAMode;

    public TextureCache GetTextureCache(){return textureCache;}
    public MaterialsCache GetMaterialsCache(){return materialsCache;}
    public VulkanContext GetVulkanContext(){return RendererContext;}

    public EngineConfig.AntiAliasType GetCurrentAAMode(){return CurrentAAMode;}
    public SceneRender GetSceneRender(){return (SceneRender) sceneRender;}

    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);
        PostProcessor = new PostProcess(RendererContext,sceneRender.GetAttachmentColour());
        GuiRender = new GuiRenderer(engineInstance, RendererContext, GraphicsQueue, PostProcessor.GetAttachment());
        swapChainRender = new SwapChainRender(RendererContext,PostProcessor.GetAttachment());
        textureCache = new TextureCache();
        materialsCache = new MaterialsCache();
        modelsCache = new ModelsCache();
        Resize = false;
        CurrentAAMode = EngineConfig.getInstance().RenderAAType();
    }

    public void Initialise(InitData initData){

        Materials.addAll(initData.Materials());
        Logger.debug("Loading {} Materials", Materials.size());
        //materialsCache.CleanUp(RendererContext);
        materialsCache.LoadMaterials(RendererContext, Materials,textureCache, CommandPools[0], GraphicsQueue);
        Logger.debug("Loaded {} Materials", Materials.size());

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

        Models.addAll(initData.Models());
        Logger.debug("Loading {} models", Models.size());
        //modelsCache.CleanUp(RendererContext);
        modelsCache.loadModels(RendererContext, Models, CommandPools[0], GraphicsQueue);
        Logger.debug("Loaded {} models", Models.size());

        List<GuiTexture> guiTextures = initData.GuiTextures();
        if(guiTextures != null){
            initData.GuiTextures().forEach(texture -> textureCache.AddTexture(RendererContext, texture.TexturePath(), texture.TexturePath(), VK_FORMAT_R8G8B8A8_SRGB));
        }

        sceneRender.LoadMaterials(RendererContext,materialsCache,textureCache);
        GuiRender.LoadTextures(RendererContext,initData.GuiTextures(),textureCache);
    }

    public GuiRenderer GetGUIRenderer(){return GuiRender;}
    public long GetGuiTexture(long ID){return GuiRender.GetVkDescriptorImage(ID);}

    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);
        PostProcessor.CleanUp(RendererContext);
        swapChainRender.CleanUp(RendererContext);
        GuiRender.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);

        sceneRender.Render(engineInstance,RendererContext,CommandBuffer, modelsCache,materialsCache,CurrentFrame);
        PostProcessor.Render(RendererContext,CommandBuffer,sceneRender.GetAttachmentColour());
        GuiRender.Render(RendererContext,CommandBuffer,CurrentFrame,PostProcessor.GetAttachment());

        int ImageIndex;
        if (Resize || (ImageIndex = swapChain.FetchNextImage(RendererContext.GetDevice(), PresentCompleteSemaphores[CurrentFrame])) < 0){
            resize(engineInstance);
            return;
        }

        swapChainRender.Render(RendererContext,CommandBuffer,PostProcessor.GetAttachment(),ImageIndex);

        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);
        PostProcessor.Resize(RendererContext,sceneRender.GetAttachmentColour());
        GuiRender.Resize(RendererContext,PostProcessor.GetAttachment());
        swapChainRender.Resize(RendererContext,PostProcessor.GetAttachment());
    }

    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);
        }
    }
}
