#include <Core/HAL/Display/Win32/Win32Window.h>
#include <Core/Logging/LogManager.h>
#include <Core/Logging/LogTypes.h>
#include <Core/Memory/Allocator.h>
#include <Graphics/D3D12/D3D12CommandList.h>
#include <Graphics/D3D12/D3D12GraphicsDevice.h>
#include <Graphics/D3D12/D3D12Includes.h>
#include <Graphics/D3D12/D3D12SwapChain.h>
#include <Graphics/D3D12/D3D12Synchronization.h>
#include <Graphics/D3D12/D3D12Texture.h>
#include <Graphics/D3D12/D3D12Utils.h>

#include <algorithm>

namespace Juliet::D3D12
{
    namespace
    {
        DXGI_COLOR_SPACE_TYPE SwapchainCompositionToColorSpace[] = {
            DXGI_COLOR_SPACE_RGB_FULL_G22_NONE_P709,   // SDR
            DXGI_COLOR_SPACE_RGB_FULL_G22_NONE_P709,   // SDR_LINEAR
            DXGI_COLOR_SPACE_RGB_FULL_G10_NONE_P709,   // HDR_EXTENDED_LINEAR
            DXGI_COLOR_SPACE_RGB_FULL_G2084_NONE_P2020 // HDR10_ST2084
        };

        DXGI_FORMAT SwapchainCompositionToTextureFormat[] = {
            DXGI_FORMAT_B8G8R8A8_UNORM,     // SDR
            DXGI_FORMAT_B8G8R8A8_UNORM,     // SDR_LINEAR (NOTE: The RTV uses the sRGB format)
            DXGI_FORMAT_R16G16B16A16_FLOAT, // HDR_EXTENDED_LINEAR
            DXGI_FORMAT_R10G10B10A2_UNORM,  // HDR10_ST2084
        };

        TextureFormat SwapchainCompositionToJulietTextureFormat[] = {
            TextureFormat::B8G8R8A8_UNORM,      // SDR
            TextureFormat::B8G8R8A8_UNORM_SRGB, // SDR_LINEAR
            TextureFormat::R16G16B16A16_FLOAT,  // HDR_EXTENDED_LINEAR
            TextureFormat::R10G10B10A2_UNORM,   // HDR10_ST2084
        };

        bool CreateSwapChainTexture(NonNullPtr<D3D12Driver> driver, NonNullPtr<IDXGISwapChain3> swapChain,
                                    SwapChainComposition composition, NonNullPtr<D3D12TextureContainer> textureContainer, uint8 index)
        {
            ID3D12Resource* swapChainTexture = nullptr;
            HRESULT result = swapChain->GetBuffer(index, IID_ID3D12Resource, reinterpret_cast<void**>(&swapChainTexture));
            if (FAILED(result))
            {
                LogError(driver->D3D12Device, "Cannot get buffer from SwapChain", result);
                return false;
            }

            auto texture = static_cast<D3D12Texture*>(Calloc(1, sizeof(D3D12Texture)));
            if (!texture)
            {
                LogError(driver->D3D12Device, "Cannot allocate D3D12Texture (out of memory)", result);
                swapChainTexture->Release();
                return false;
            }

            texture->ReferenceCount += 1;
            texture->SubresourceCount = 1;
            texture->Subresources = static_cast<D3D12TextureSubresource*>(Calloc(1, sizeof(D3D12TextureSubresource)));
            if (!texture->Subresources)
            {
                LogError(driver->D3D12Device, "Cannot allocate D3D12TextureSubresource (out of memory)", result);
                Free(texture);
                swapChainTexture->Release();
                return false;
            }
            texture->Subresources[0].RTVHandles =
                static_cast<D3D12StagingDescriptor*>(Calloc(1, sizeof(D3D12StagingDescriptor)));
            texture->Subresources[0].UAVHandle.Heap = nullptr;
            texture->Subresources[0].UAVHandle.Heap = nullptr;
            texture->Subresources[0].Parent         = texture;
            texture->Subresources[0].Index          = 0;
            texture->Subresources[0].Layer          = 0;
            texture->Subresources[0].Depth          = 1;
            texture->Subresources[0].Level          = 0;

            D3D12_RESOURCE_DESC textureDesc                   = swapChainTexture->GetDesc();
            textureContainer->Header.CreateInfo.Width         = static_cast<uint32>(textureDesc.Width);
            textureContainer->Header.CreateInfo.Height        = static_cast<uint32>(textureDesc.Height);
            textureContainer->Header.CreateInfo.LayerCount    = 1;
            textureContainer->Header.CreateInfo.MipLevelCount = 1;
            textureContainer->Header.CreateInfo.Type          = TextureType::Texture_2D;
            textureContainer->Header.CreateInfo.Flags         = TextureUsageFlag::ColorTarget;
            textureContainer->Header.CreateInfo.SampleCount   = TextureSampleCount::One;
            textureContainer->Header.CreateInfo.Format = SwapchainCompositionToJulietTextureFormat[ToUnderlying(composition)];

            textureContainer->Textures = static_cast<D3D12Texture**>(Calloc(1, sizeof(D3D12Texture*)));
            if (!textureContainer->Textures)
            {
                Free(texture->Subresources);
                Free(texture);
                swapChainTexture->Release();
                return false;
            }

            textureContainer->Capacity      = 1;
            textureContainer->Count         = 1;
            textureContainer->Textures[0]   = texture;
            textureContainer->ActiveTexture = texture;
            textureContainer->CanBeCycled   = false;

            texture->Container        = textureContainer;
            texture->IndexInContainer = 0;

            // Assign RTV to the swapchain texture
            DXGI_FORMAT swapchainFormat = SwapchainCompositionToTextureFormat[ToUnderlying(composition)];
            Internal::AssignStagingDescriptor(driver, D3D12_DESCRIPTOR_HEAP_TYPE_RTV, texture->Subresources[0].RTVHandles[0]);
            D3D12_RENDER_TARGET_VIEW_DESC rtvDesc;
            rtvDesc.Format = (composition == SwapChainComposition::SDR_LINEAR) ? DXGI_FORMAT_B8G8R8A8_UNORM_SRGB : swapchainFormat;
            rtvDesc.ViewDimension        = D3D12_RTV_DIMENSION_TEXTURE2D;
            rtvDesc.Texture2D.MipSlice   = 0;
            rtvDesc.Texture2D.PlaneSlice = 0;

            driver->D3D12Device->CreateRenderTargetView(swapChainTexture, &rtvDesc,
                                                        texture->Subresources[0].RTVHandles[0].CpuHandle);

            swapChainTexture->Release();

            return true;
        }

        bool AcquireSwapChainTexture(bool block, NonNullPtr<CommandList> commandList, NonNullPtr<Window> window, Texture** swapchainTexture)
        {
            auto d3d12CommandList = reinterpret_cast<D3D12CommandList*>(commandList.Get());

            auto* driver = d3d12CommandList->Driver;
            Assert(driver->WindowData);

            // TODO: Find a way to fetch window data more smoothly from the window ptr
            // In the mean time i will just void it or the variable is unused and cause a warning
            (void)window;
            auto* windowData = driver->WindowData;
            Assert(windowData->Window == window.Get());

            if (windowData->InFlightFences[windowData->WindowFrameCounter] != nullptr)
            {
                if (block)
                {
                    // Wait until the fence for the frame is signaled.
                    // In VSYNC this means waiting that the least recent presented frame is done
                    if (!Wait(driver, true, &windowData->InFlightFences[windowData->WindowFrameCounter],
                              1 JULIET_DEBUG_PARAM(ConstString("AcquireSwapChainTexture"))))
                    {
                        return false;
                    }
                }
                else
                {
                    // If work is not done, the least recent fence wont be signaled.
                    // In that case we return true to notify that there is no error, but rendering should be skipped as their will be no swapchainTexture
                    if (!QueryFence(driver, windowData->InFlightFences[windowData->WindowFrameCounter]))
                    {
                        return true;
                    }
                }

                ReleaseFence(driver, windowData->InFlightFences[windowData->WindowFrameCounter] JULIET_DEBUG_PARAM(
                                         ConstString("AcquireSwapChainTexture")));
                windowData->InFlightFences[windowData->WindowFrameCounter] = nullptr;
            }

            uint32  swapchainIndex = windowData->SwapChain->GetCurrentBackBufferIndex();
            HRESULT result         = windowData->SwapChain->GetBuffer(
                swapchainIndex, IID_ID3D12Resource,
                reinterpret_cast<void**>(&windowData->SwapChainTextureContainers[swapchainIndex].ActiveTexture->Resource));
            if (FAILED(result))
            {
                LogError(driver->D3D12Device, "Could not acquire swapchain", result);
                return false;
            }

            // When the swap chain texture is acquired it's time to present
            if (d3d12CommandList->PresentDataCount == d3d12CommandList->PresentDataCapacity)
            {
                d3d12CommandList->PresentDataCapacity += 1;
                d3d12CommandList->PresentDatas = static_cast<D3D12PresentData*>(
                    Realloc(d3d12CommandList->PresentDatas, d3d12CommandList->PresentDataCapacity * sizeof(D3D12PresentData)));
            }
            d3d12CommandList->PresentDatas[d3d12CommandList->PresentDataCount].WindowData          = windowData;
            d3d12CommandList->PresentDatas[d3d12CommandList->PresentDataCount].SwapChainImageIndex = swapchainIndex;
            d3d12CommandList->PresentDataCount += 1;

            // Create the presentation barrier.
            D3D12_RESOURCE_BARRIER barrierDesc;
            barrierDesc.Type                   = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
            barrierDesc.Flags                  = D3D12_RESOURCE_BARRIER_FLAG_NONE;
            barrierDesc.Transition.StateBefore = D3D12_RESOURCE_STATE_PRESENT;
            barrierDesc.Transition.StateAfter  = D3D12_RESOURCE_STATE_RENDER_TARGET;
            barrierDesc.Transition.pResource = windowData->SwapChainTextureContainers[swapchainIndex].ActiveTexture->Resource;
            barrierDesc.Transition.Subresource = 0;

            d3d12CommandList->GraphicsCommandList.CommandList->ResourceBarrier(1, &barrierDesc);

            *swapchainTexture = reinterpret_cast<Texture*>(&windowData->SwapChainTextureContainers[swapchainIndex]);

            return true;
        }
    } // namespace

    bool AcquireSwapChainTexture(NonNullPtr<CommandList> commandList, NonNullPtr<Window> window, Texture** swapChainTexture)
    {
        return AcquireSwapChainTexture(false, commandList, window, swapChainTexture);
    }

    bool WaitAndAcquireSwapChainTexture(NonNullPtr<CommandList> commandList, NonNullPtr<Window> window, Texture** swapChainTexture)
    {
        return AcquireSwapChainTexture(true, commandList, window, swapChainTexture);
    }

    bool WaitForSwapchain(NonNullPtr<GPUDriver> driver, NonNullPtr<Window> /*window*/)
    {
        auto* d3d12Driver = static_cast<D3D12Driver*>(driver.Get());
        auto* windowData  = d3d12Driver->WindowData;
        if (!windowData)
        {
            LogError(LogCategory::Graphics, "Cannot wait for swapchain. Window has no Swapchain");
            return false;
        }

        if (windowData->InFlightFences[windowData->WindowFrameCounter] != nullptr)
        {
            if (!Wait(d3d12Driver, true, &windowData->InFlightFences[windowData->WindowFrameCounter],
                      1 JULIET_DEBUG_PARAM(ConstString("WaitForSwapchain"))))
            {
                return false;
            }
        }

        return true;
    }

    TextureFormat GetSwapChainTextureFormat(NonNullPtr<GPUDriver> driver, [[maybe_unused]] NonNullPtr<Window> window)
    {
        auto* d3d12Driver = static_cast<D3D12Driver*>(driver.Get());

        auto* windowData = d3d12Driver->WindowData;
        if (!windowData)
        {
            LogError(LogCategory::Graphics, "Cannot get swapchain format. Window has no Swapchain");
            return TextureFormat::Invalid;
        }

        Assert(windowData->Window == window.Get());
        return windowData->SwapChainTextureContainers[windowData->WindowFrameCounter].Header.CreateInfo.Format;
    }

    namespace Internal
    {
        bool CreateSwapChain(NonNullPtr<D3D12Driver> driver, NonNullPtr<D3D12WindowData> windowData,
                             SwapChainComposition composition, PresentMode presentMode)
        {
            auto windowWin32State = static_cast<Win32::Window32State*>(windowData->Window->State);
            HWND windowHandle     = windowWin32State->Handle;
            if (!IsWindow(windowHandle))
            {
                Assert(false, "windowWin32State->Handle is not a window handle ???");
                return false;
            }

            // TODO: I have no way to test HDR easily except the steamdeck
            DXGI_FORMAT swapChainFormat = SwapchainCompositionToTextureFormat[ToUnderlying(composition)];

            windowData->SwapChainTextureCount = std::clamp<uint8>(driver->FramesInFlight, 2, 3);

            DXGI_SWAP_CHAIN_DESC1 swapChainDesc = {};
            swapChainDesc.Width                 = 0; // Use the whole width
            swapChainDesc.Height                = 0; // Use the whole height
            swapChainDesc.Format                = swapChainFormat;
            swapChainDesc.Stereo                = 0;
            swapChainDesc.SampleDesc.Count      = 1;
            swapChainDesc.SampleDesc.Quality    = 0;
            swapChainDesc.BufferUsage           = DXGI_USAGE_RENDER_TARGET_OUTPUT;
            swapChainDesc.BufferCount           = windowData->SwapChainTextureCount;
            swapChainDesc.Scaling               = DXGI_SCALING_STRETCH;
            swapChainDesc.SwapEffect            = DXGI_SWAP_EFFECT_FLIP_DISCARD;
            swapChainDesc.AlphaMode             = DXGI_ALPHA_MODE_UNSPECIFIED;
            if (driver->IsTearingSupported)
            {
                swapChainDesc.Flags = DXGI_SWAP_CHAIN_FLAG_ALLOW_TEARING;
            }
            else
            {
                swapChainDesc.Flags = 0;
            }

            DXGI_SWAP_CHAIN_FULLSCREEN_DESC swapChainFullscreenDesc = {};
            swapChainFullscreenDesc.RefreshRate.Numerator           = 0;
            swapChainFullscreenDesc.RefreshRate.Denominator         = 0;
            swapChainFullscreenDesc.ScanlineOrdering                = DXGI_MODE_SCANLINE_ORDER_UNSPECIFIED;
            swapChainFullscreenDesc.Scaling                         = DXGI_MODE_SCALING_UNSPECIFIED;
            swapChainFullscreenDesc.Windowed                        = true;

            IDXGISwapChain1* swapChain = nullptr;
            HRESULT result = driver->DXGIFactory->CreateSwapChainForHwnd(static_cast<IUnknown*>(driver->GraphicsQueue),
                                                                         windowHandle, &swapChainDesc,
                                                                         &swapChainFullscreenDesc, nullptr, &swapChain);
            if (FAILED(result))
            {
                LogError(driver->D3D12Device, "Failed to create SwapChain", result);
                return false;
            }

            IDXGISwapChain3* swapChain3 = nullptr;
            result = swapChain->QueryInterface(IID_IDXGISwapChain3, reinterpret_cast<void**>(&swapChain3));
            swapChain->Release();
            if (FAILED(result))
            {
                LogError(driver->D3D12Device, "Could not query IDXGISwapChain3 interface", result);
                return false;
            }

            if (composition != SwapChainComposition::SDR)
            {
                swapChain3->SetColorSpace1(SwapchainCompositionToColorSpace[ToUnderlying(composition)]);
            }

            IDXGIFactory1* parentFactory = nullptr;
            result = swapChain3->GetParent(IID_IDXGIFactory1, reinterpret_cast<void**>(&parentFactory));
            if (FAILED(result))
            {
                Log(LogLevel::Warning, LogCategory::Graphics, "Cannot get SwapChain Parent! Error Code: " HRESULT_FMT, result);
            }
            else
            {
                // Disable DXGI window crap
                result = parentFactory->MakeWindowAssociation(windowHandle, DXGI_MWA_NO_WINDOW_CHANGES);
                if (FAILED(result))
                {
                    Log(LogLevel::Warning, LogCategory::Graphics, "MakeWindowAssociation failed! Error Code: " HRESULT_FMT, result);
                }
                parentFactory->Release();
            }

            swapChain3->GetDesc1(&swapChainDesc);
            if (FAILED(result))
            {
                LogError(driver->D3D12Device, "Failed to retrieve SwapChain descriptor", result);
                return false;
            }
            windowData->SwapChain            = swapChain3;
            windowData->SwapChainColorSpace  = SwapchainCompositionToColorSpace[ToUnderlying(composition)];
            windowData->SwapChainComposition = composition;
            windowData->WindowFrameCounter   = 0;
            windowData->Width                = swapChainDesc.Width;
            windowData->Height               = swapChainDesc.Height;
            windowData->PresentMode          = presentMode;

            for (uint8 idx = 0; idx < windowData->SwapChainTextureCount; ++idx)
            {
                if (!CreateSwapChainTexture(driver, swapChain3, composition, &windowData->SwapChainTextureContainers[idx], idx))
                {
                    swapChain3->Release();
                    return false;
                }
            }

            return true;
        }

        void DestroySwapChain(NonNullPtr<D3D12Driver> driver, NonNullPtr<D3D12WindowData> windowData)
        {
            for (uint32 idx = 0; idx < windowData->SwapChainTextureCount; ++idx)
            {
                ReleaseStagingDescriptor(driver,
                                         windowData->SwapChainTextureContainers[idx].ActiveTexture->Subresources[0].RTVHandles[0]);

                Free(windowData->SwapChainTextureContainers[idx].ActiveTexture->Subresources[0].RTVHandles);
                Free(windowData->SwapChainTextureContainers[idx].ActiveTexture->Subresources);
                Free(windowData->SwapChainTextureContainers[idx].ActiveTexture);
                Free(windowData->SwapChainTextureContainers[idx].Textures);
            }

            windowData->SwapChain->Release();
            windowData->SwapChain = nullptr;
        }
    } // namespace Internal
} // namespace Juliet::D3D12
