#include "Win32DxWindowInterface.h"

#include "DirectXInterface.h"
#include "Win32Window.h"
#include "Window.h"
#include "Widget.h"
#include "FileLogger.h"

#include "DirectXDescriptors.h"
#include "DirectX2dIntegration.h"
#include "DirectXBuffers.h"
#include "DirectXCommandList.h"

#include "DrawingContext.h"
#include "DirectXPainter.h"
#include "DirectXMeshPainter.h"

#include <dxgi.h>
#include <dxgi1_6.h>
#include <d3d12sdklayers.h>
#include <D3Dcommon.h>

Win32DxWindowInterface::Win32DxWindowInterface(mt::Window* window, DirectXInterface* dxInterface)
	:   mWindow(window),
        mDxInterface(dxInterface)
{

}

Win32DxWindowInterface::~Win32DxWindowInterface()
{

}

void Win32DxWindowInterface::destroyWindow()
{
    flushCommandQueue();
    ::CloseHandle(mFenceEvent);
}

void Win32DxWindowInterface::onRender()
{
    if (!mIsValid)
    {
        MLOG_INFO("Invalid state - returning");
        return;
    }

    if (!mBuffers->isInitialized())
    {
        MLOG_INFO("Buffer not initialized - returning");
        return;
    }

    updateCommandList();

    mCommandList->execute(mDxInterface->getCommandQueue());

    mD2dIntegration->render();

    mSwapChain->Present(1, 0);

    flushCommandQueue();
}

void Win32DxWindowInterface::updateCommandList()
{
    // Command list allocators can only be reset when the associated 
    // command lists have finished execution on the GPU; apps should use 
    // fences to determine GPU execution progress.
    mCommandList->resetAllocator();

    // However, when ExecuteCommandList() is called on a particular command 
    // list, that command list can then be reset at any time and must be before 
    // re-recording.
    mCommandList->reset(getPipelineState());

    // Set necessary state.
    auto painter = dynamic_cast<DirectXPainter*>(mWindow->getDrawingContent()->getPainter());
    mCommandList->get()->SetGraphicsRootSignature(painter->getMeshPainter()->getRootSignature());
    mCommandList->get()->RSSetViewports(1, &mViewport);
    mCommandList->get()->RSSetScissorRects(1, &mScissorRect);

    // Indicate that the back buffer will be used as a render target.
    auto barrier = CD3DX12_RESOURCE_BARRIER::Transition(mBuffers->getCurrentBackBuffer(), D3D12_RESOURCE_STATE_PRESENT, D3D12_RESOURCE_STATE_RENDER_TARGET);
    mCommandList->get()->ResourceBarrier(1, &barrier);

    auto rtv_handle = mDescriptors->getRtvHandle(static_cast<int>(mBuffers->getCurrentBackBufferIndex()));
    mCommandList->get()->OMSetRenderTargets(1, &rtv_handle, FALSE, nullptr);

    painter->paintBackground(rtv_handle, mCommandList->get());
    painter->paintMesh(mCommandList->get());

    // Indicate that the back buffer will now be used to present. (Exclude as handled by D2D integration later)
    //barrier = CD3DX12_RESOURCE_BARRIER::Transition(mBuffers->getCurrentBackBuffer(), D3D12_RESOURCE_STATE_RENDER_TARGET, D3D12_RESOURCE_STATE_PRESENT);
    //mCommandList->get()->ResourceBarrier(1, &barrier);

    mCommandList->close();
}

bool Win32DxWindowInterface::initialize()
{
	if (mInitialized)
	{
		return true;
	}
    mInitialized = true;
    MLOG_INFO("D3D Win init");

    mD2dIntegration = std::make_unique<DirectX2dIntegration>(mWindow, mDxInterface, FrameCount);
    mBuffers = std::make_unique<DirectXBuffers>(FrameCount);

    if (!createSwapChain())
    {
        return false;
    }

    mCommandList = std::make_unique<DirectXCommandList>();
    mCommandList->create(mDxInterface->getD3dDevice(), getPipelineState());

    createSyncObjects();

    mDescriptors = std::make_unique<DirectXDescriptors>();
    mDescriptors->setupDefaultDescriptorHeaps(mDxInterface->getD3dDevice());

    onResize();

    if (!loadAssets())
    {
        return false;
    }

    flushCommandQueue();

    mIsValid = true;
    return true;
}

void Win32DxWindowInterface::onResize()
{
    MLOG_INFO("On resize");
    mIsValid = false;

    flushCommandQueue();

    mCommandList->reset(getPipelineState());

    clearBuffers();

    resizeSwapChain();

    populateBuffers();

    mCommandList->close();
    
    mCommandList->execute(mDxInterface->getCommandQueue());

    flushCommandQueue();

    updateViewport();

    mIsValid = true;
}

void Win32DxWindowInterface::resizeSwapChain()
{
    const auto width = mWindow->getWidth();
    const auto height = mWindow->getHeight();
    MLOG_INFO("Resizing swap chain to: " << width << " | " << height);
    mSwapChain->ResizeBuffers(FrameCount, width, height, DXGI_FORMAT_R8G8B8A8_UNORM, DXGI_SWAP_CHAIN_FLAG_ALLOW_MODE_SWITCH);
    updateBackbufferIndex(0);
}

void Win32DxWindowInterface::updateViewport()
{
    const auto width = mWindow->getWidth();
    const auto height = mWindow->getHeight();
    mViewport = CD3DX12_VIEWPORT(0.0f, 0.0f, static_cast<float>(width), static_cast<float>(height));
    mScissorRect = CD3DX12_RECT(0, 0, width, height);
}

void Win32DxWindowInterface::clearBuffers()
{
    mD2dIntegration->clearBuffers();
    mBuffers->clear();  
}

ID3D12PipelineState* Win32DxWindowInterface::getPipelineState() const
{
    auto painter = dynamic_cast<DirectXPainter*>(mWindow->getDrawingContent()->getPainter());
    return painter->getMeshPainter()->getPipelineState();
}

void Win32DxWindowInterface::populateBuffers()
{
    MLOG_INFO("Populating buffers");
    auto d3d_device = mDxInterface->getD3dDevice();
    mBuffers->setupBackBuffers(d3d_device, mSwapChain.Get(), mDescriptors.get());
    for (UINT idx = 0; idx < FrameCount; idx++)
    {
        mD2dIntegration->wrapBuffer(mBuffers->getBackBuffer(idx));
    }

    const auto width = mWindow->getWidth();
    const auto height = mWindow->getHeight();
    if (width > 0 and height > 0)
    {
        mBuffers->setupDepthStencilHeap(mDxInterface->getD3dDevice(), width, height);
        mBuffers->setupDepthStencilView(mDxInterface->getD3dDevice(), mDescriptors.get());

        auto barrier = CD3DX12_RESOURCE_BARRIER::Transition(mBuffers->getDepthStencilBuffer(), D3D12_RESOURCE_STATE_COMMON, D3D12_RESOURCE_STATE_DEPTH_WRITE);
        mCommandList->get()->ResourceBarrier(1, &barrier);
    }
}

bool Win32DxWindowInterface::createSwapChain()
{
    mSwapChain.Reset();

    const auto width = mWindow->getWidth();
    const auto height = mWindow->getHeight();
    const auto hwnd = dynamic_cast<Win32Window*>(mWindow->getPlatformWindow())->getHandle();

    DXGI_SWAP_CHAIN_DESC1 swapChainDesc = {};
    swapChainDesc.BufferCount = FrameCount;
    swapChainDesc.Width = width;
    swapChainDesc.Height = height;
    swapChainDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    swapChainDesc.BufferUsage = DXGI_USAGE_RENDER_TARGET_OUTPUT;
    swapChainDesc.SwapEffect = DXGI_SWAP_EFFECT_FLIP_DISCARD;
    swapChainDesc.SampleDesc.Count = 1;

    Microsoft::WRL::ComPtr<IDXGISwapChain1> swapChain;
    if (FAILED(mDxInterface->getDxgiFactory()->CreateSwapChainForHwnd(mDxInterface->getCommandQueue(), hwnd, &swapChainDesc, nullptr, nullptr, &swapChain)))
    {
        MLOG_ERROR("Failed to create swap chain for hwnd: " << hwnd);
        return false;
    }
    swapChain.As(&mSwapChain);

    // Prevent fullscreen toggle
    mDxInterface->getDxgiFactory()->MakeWindowAssociation(hwnd, DXGI_MWA_NO_ALT_ENTER);
    
    // Update index for backbuffer
    updateBackbufferIndex(mSwapChain->GetCurrentBackBufferIndex());
    return true;
}

void Win32DxWindowInterface::updateBackbufferIndex(UINT index)
{
    mBuffers->updateBackbufferIndex(index);
    mD2dIntegration->updateBackbufferIndex(index);
}

bool Win32DxWindowInterface::loadAssets()
{
    auto painter = dynamic_cast<DirectXPainter*>(mWindow->getDrawingContent()->getPainter());
    painter->initializeMesh();
    painter->updateMesh();

    return true;
}

void Win32DxWindowInterface::createSyncObjects()
{
    mDxInterface->getD3dDevice()->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&mFence));
    mFenceValue = 1;

    // Create an event handle to use for frame synchronization.
    mFenceEvent = ::CreateEvent(nullptr, FALSE, FALSE, nullptr);
    if (mFenceEvent == nullptr)
    {
        //ThrowIfFailed(HRESULT_FROM_WIN32(GetLastError()));
    }

    // Wait for the command list to execute; we are reusing the same command 
    // list in our main loop but for now, we just want to wait for setup to 
    // complete before continuing.
    flushCommandQueue();
}

void Win32DxWindowInterface::flushCommandQueue()
{
    // Signal and increment the fence value.
    const UINT64 fence = mFenceValue;
    mDxInterface->getCommandQueue()->Signal(mFence.Get(), fence);
    mFenceValue++;

    // Wait until the previous frame is finished.
    if (mFence->GetCompletedValue() < fence)
    {
        mFence->SetEventOnCompletion(fence, mFenceEvent);
        ::WaitForSingleObject(mFenceEvent, INFINITE);
    }

    updateBackbufferIndex(mSwapChain->GetCurrentBackBufferIndex());
}