#include "Backend.hpp"

#if BUILD_CORE_WITH_DX12_BACKEND
#	include <backend/imgui_impl_dx12.h>
#	include <backend/imgui_impl_win32.h>
#	include <d3d12.h>
#	include <dxgi1_4.h>
#	include <tchar.h>
#	include <backend/imgui_impl_dx12.cpp>
#	include <stdexcept>

constexpr int kNumFramesInFlight = 3;
constexpr int kNumBackBuffers = 3;

// Forward declare message handler from imgui_impl_win32.cpp
extern IMGUI_IMPL_API LRESULT ImGui_ImplWin32_WndProcHandler(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam);

class DirectX12Backend : public RenderingBackend
{
private:
	struct FrameContext
	{
		ID3D12CommandAllocator* CommandAllocator;
		UINT64 FenceValue;
	};

	HWND hWnd;
	WNDCLASSEX wc;

	FrameContext mFrameContext[kNumFramesInFlight] = {};
	UINT mFrameIndex = 0;

	ID3D12Device* mD3dDevice = nullptr;
	ID3D12DescriptorHeap* mD3dRtvDescHeap = nullptr;
	ID3D12DescriptorHeap* mD3dSrvDescHeap = nullptr;
	ID3D12CommandQueue* mD3dCommandQueue = nullptr;
	ID3D12GraphicsCommandList* mD3dCommandList = nullptr;
	ID3D12Fence* mFence = nullptr;
	HANDLE mFenceEvent = nullptr;
	UINT64 mFenceLastSignaledValue = 0;
	IDXGISwapChain3* mSwapChain = nullptr;
	HANDLE mSwapChainWaitableObject = nullptr;
	ID3D12Resource* mMainRenderTargetResource[kNumBackBuffers] = {};
	D3D12_CPU_DESCRIPTOR_HANDLE mMainRenderTargetDescriptor[kNumBackBuffers] = {};

public:
	DirectX12Backend()
	{
		ImGui_ImplWin32_EnableDpiAwareness();

		wc.cbSize = sizeof(WNDCLASSEX);
		wc.style = CS_CLASSDC;
		wc.lpfnWndProc = &StaticWndProc;
		wc.cbClsExtra = 0L;
		wc.cbWndExtra = 0L;
		wc.hInstance = GetModuleHandle(nullptr);
		wc.hIcon = nullptr;
		wc.hCursor = nullptr;
		wc.hbrBackground = nullptr;
		wc.lpszMenuName = nullptr;
		wc.lpszClassName = _T("Cplt");
		wc.hIconSm = nullptr;
		::RegisterClassEx(&wc);

		hWnd = ::CreateWindow(
			wc.lpszClassName,
			_T("Cplt main window"),
			WS_OVERLAPPEDWINDOW,
			/* x */ 100,
			/* y */ 100,
			/* window width */ 1280,
			/* window height */ 800,
			nullptr,
			nullptr,
			wc.hInstance,
			this);

		if (!CreateDeviceD3D()) {
			CleanupDeviceD3D();
			::UnregisterClass(wc.lpszClassName, wc.hInstance);
			throw std::runtime_error("Failed to create d3d device.");
		}

		::ShowWindow(hWnd, SW_SHOWDEFAULT);
		::UpdateWindow(hWnd);

		IMGUI_CHECKVERSION();
		ImGui::CreateContext();

		ImGui_ImplWin32_Init(hWnd);
		ImGui_ImplDX12_Init(mD3dDevice, kNumFramesInFlight, DXGI_FORMAT_R8G8B8A8_UNORM, mD3dSrvDescHeap, mD3dSrvDescHeap->GetCPUDescriptorHandleForHeapStart(), mD3dSrvDescHeap->GetGPUDescriptorHandleForHeapStart());
	}

	virtual ~DirectX12Backend()
	{
		WaitForLastSubmittedFrame();

		// Cleanup
		ImGui_ImplDX12_Shutdown();
		ImGui_ImplWin32_Shutdown();
		ImGui::DestroyContext();

		CleanupDeviceD3D();
		::DestroyWindow(hWnd);
		::UnregisterClass(wc.lpszClassName, wc.hInstance);
	}

	virtual void RunUntilWindowClose(void (*windowContent)())
	{
		while (true) {
			MSG msg;
			bool done = false;
			while (::PeekMessage(&msg, nullptr, 0U, 0U, PM_REMOVE))
			{
				::TranslateMessage(&msg);
				::DispatchMessage(&msg);
				if (msg.message == WM_QUIT) {
					done = true;
				}
			}
			if (done) break;

			// Start the Dear ImGui frame
			ImGui_ImplDX12_NewFrame();
			ImGui_ImplWin32_NewFrame();
			ImGui::NewFrame();

			windowContent();

			ImGui::Render();

			FrameContext* frameCtx = WaitForNextFrameResources();
			UINT backBufferIdx = mSwapChain->GetCurrentBackBufferIndex();
			frameCtx->CommandAllocator->Reset();

			D3D12_RESOURCE_BARRIER barrier = {};
			barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
			barrier.Flags = D3D12_RESOURCE_BARRIER_FLAG_NONE;
			barrier.Transition.pResource = mMainRenderTargetResource[backBufferIdx];
			barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
			barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_PRESENT;
			barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_RENDER_TARGET;
			mD3dCommandList->Reset(frameCtx->CommandAllocator, nullptr);
			mD3dCommandList->ResourceBarrier(1, &barrier);

			const ImVec4 kClearColor = ImVec4(0.45f, 0.55f, 0.60f, 1.00f);
			const float kClearColorWithAlpha[4] = { kClearColor.x * kClearColor.w, kClearColor.y * kClearColor.w, kClearColor.z * kClearColor.w, kClearColor.w };
			mD3dCommandList->ClearRenderTargetView(mMainRenderTargetDescriptor[backBufferIdx], kClearColorWithAlpha, 0, nullptr);
			mD3dCommandList->OMSetRenderTargets(1, &mMainRenderTargetDescriptor[backBufferIdx], FALSE, nullptr);
			mD3dCommandList->SetDescriptorHeaps(1, &mD3dSrvDescHeap);
			ImGui_ImplDX12_RenderDrawData(ImGui::GetDrawData(), mD3dCommandList);
			barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_RENDER_TARGET;
			barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_PRESENT;
			mD3dCommandList->ResourceBarrier(1, &barrier);
			mD3dCommandList->Close();

			mD3dCommandQueue->ExecuteCommandLists(1, (ID3D12CommandList* const*)&mD3dCommandList);

			mSwapChain->Present(1, 0); // Present with vsync

			UINT64 fenceValue = mFenceLastSignaledValue + 1;
			mD3dCommandQueue->Signal(mFence, fenceValue);
			mFenceLastSignaledValue = fenceValue;
			frameCtx->FenceValue = fenceValue;
		}
	}

private:
	bool CreateDeviceD3D()
	{
		// Setup swap chain
		DXGI_SWAP_CHAIN_DESC1 sd;
		{
			ZeroMemory(&sd, sizeof(sd));
			sd.BufferCount = kNumBackBuffers;
			sd.Width = 0;
			sd.Height = 0;
			sd.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
			sd.Flags = DXGI_SWAP_CHAIN_FLAG_FRAME_LATENCY_WAITABLE_OBJECT;
			sd.BufferUsage = DXGI_USAGE_RENDER_TARGET_OUTPUT;
			sd.SampleDesc.Count = 1;
			sd.SampleDesc.Quality = 0;
			sd.SwapEffect = DXGI_SWAP_EFFECT_FLIP_DISCARD;
			sd.AlphaMode = DXGI_ALPHA_MODE_UNSPECIFIED;
			sd.Scaling = DXGI_SCALING_STRETCH;
			sd.Stereo = FALSE;
		}

		// Create device
		D3D_FEATURE_LEVEL featureLevel = D3D_FEATURE_LEVEL_11_0;
		if (D3D12CreateDevice(nullptr, featureLevel, IID_PPV_ARGS(&mD3dDevice)) != S_OK) {
			return false;
		}

		{
			D3D12_DESCRIPTOR_HEAP_DESC desc = {};
			desc.Type = D3D12_DESCRIPTOR_HEAP_TYPE_RTV;
			desc.NumDescriptors = kNumBackBuffers;
			desc.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_NONE;
			desc.NodeMask = 1;
			if (mD3dDevice->CreateDescriptorHeap(&desc, IID_PPV_ARGS(&mD3dRtvDescHeap)) != S_OK) {
				return false;
			}

			SIZE_T rtvDescriptorSize = mD3dDevice->GetDescriptorHandleIncrementSize(D3D12_DESCRIPTOR_HEAP_TYPE_RTV);
			D3D12_CPU_DESCRIPTOR_HANDLE rtvHandle = mD3dRtvDescHeap->GetCPUDescriptorHandleForHeapStart();
			for (UINT i = 0; i < kNumBackBuffers; i++) {
				mMainRenderTargetDescriptor[i] = rtvHandle;
				rtvHandle.ptr += rtvDescriptorSize;
			}
		}

		{
			D3D12_DESCRIPTOR_HEAP_DESC desc = {};
			desc.Type = D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV;
			desc.NumDescriptors = 1;
			desc.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE;
			if (mD3dDevice->CreateDescriptorHeap(&desc, IID_PPV_ARGS(&mD3dSrvDescHeap)) != S_OK) {
				return false;
			}
		}

		{
			D3D12_COMMAND_QUEUE_DESC desc = {};
			desc.Type = D3D12_COMMAND_LIST_TYPE_DIRECT;
			desc.Flags = D3D12_COMMAND_QUEUE_FLAG_NONE;
			desc.NodeMask = 1;
			if (mD3dDevice->CreateCommandQueue(&desc, IID_PPV_ARGS(&mD3dCommandQueue)) != S_OK) {
				return false;
			}
		}

		for (UINT i = 0; i < kNumFramesInFlight; i++) {
			if (mD3dDevice->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_DIRECT, IID_PPV_ARGS(&mFrameContext[i].CommandAllocator)) != S_OK) {
				return false;
			}
		}

		if (mD3dDevice->CreateCommandList(0, D3D12_COMMAND_LIST_TYPE_DIRECT, mFrameContext[0].CommandAllocator, nullptr, IID_PPV_ARGS(&mD3dCommandList)) != S_OK ||
			mD3dCommandList->Close() != S_OK)
		{
			return false;
		}

		if (mD3dDevice->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&mFence)) != S_OK) return false;

		mFenceEvent = CreateEvent(nullptr, FALSE, FALSE, nullptr);
		if (mFenceEvent == nullptr) return false;

		{
			IDXGIFactory4* dxgiFactory = nullptr;
			IDXGISwapChain1* swapChain1 = nullptr;
			if (CreateDXGIFactory1(IID_PPV_ARGS(&dxgiFactory)) != S_OK)
				return false;
			if (dxgiFactory->CreateSwapChainForHwnd(mD3dCommandQueue, hWnd, &sd, nullptr, nullptr, &swapChain1) != S_OK)
				return false;
			if (swapChain1->QueryInterface(IID_PPV_ARGS(&mSwapChain)) != S_OK)
				return false;
			swapChain1->Release();
			dxgiFactory->Release();
			mSwapChain->SetMaximumFrameLatency(kNumBackBuffers);
			mSwapChainWaitableObject = mSwapChain->GetFrameLatencyWaitableObject();
		}

		CreateRenderTarget();
		return true;
	}

	void CleanupDeviceD3D()
	{
		CleanupRenderTarget();
		if (mSwapChain) {
			mSwapChain->Release();
			mSwapChain = nullptr;
		}
		if (mSwapChainWaitableObject != nullptr) {
			CloseHandle(mSwapChainWaitableObject);
		}
		for (UINT i = 0; i < kNumFramesInFlight; i++)
			if (mFrameContext[i].CommandAllocator) {
				mFrameContext[i].CommandAllocator->Release();
				mFrameContext[i].CommandAllocator = nullptr;
			}
		if (mD3dCommandQueue) {
			mD3dCommandQueue->Release();
			mD3dCommandQueue = nullptr;
		}
		if (mD3dCommandList) {
			mD3dCommandList->Release();
			mD3dCommandList = nullptr;
		}
		if (mD3dRtvDescHeap) {
			mD3dRtvDescHeap->Release();
			mD3dRtvDescHeap = nullptr;
		}
		if (mD3dSrvDescHeap) {
			mD3dSrvDescHeap->Release();
			mD3dSrvDescHeap = nullptr;
		}
		if (mFence) {
			mFence->Release();
			mFence = nullptr;
		}
		if (mFenceEvent) {
			CloseHandle(mFenceEvent);
			mFenceEvent = nullptr;
		}
		if (mD3dDevice) {
			mD3dDevice->Release();
			mD3dDevice = nullptr;
		}
	}

	void CreateRenderTarget()
	{
		for (UINT i = 0; i < kNumBackBuffers; i++)
		{
			ID3D12Resource* pBackBuffer = nullptr;
			mSwapChain->GetBuffer(i, IID_PPV_ARGS(&pBackBuffer));
			mD3dDevice->CreateRenderTargetView(pBackBuffer, nullptr, mMainRenderTargetDescriptor[i]);
			mMainRenderTargetResource[i] = pBackBuffer;
		}
	}

	void CleanupRenderTarget()
	{
		WaitForLastSubmittedFrame();

		for (UINT i = 0; i < kNumBackBuffers; i++)
			if (mMainRenderTargetResource[i]) {
				mMainRenderTargetResource[i]->Release();
				mMainRenderTargetResource[i] = nullptr;
			}
	}

	void WaitForLastSubmittedFrame()
	{
		FrameContext* frameCtx = &mFrameContext[mFrameIndex % kNumFramesInFlight];

		UINT64 fenceValue = frameCtx->FenceValue;
		if (fenceValue == 0)
			return; // No fence was signaled

		frameCtx->FenceValue = 0;
		if (mFence->GetCompletedValue() >= fenceValue)
			return;

		mFence->SetEventOnCompletion(fenceValue, mFenceEvent);
		WaitForSingleObject(mFenceEvent, INFINITE);
	}

	FrameContext* WaitForNextFrameResources()
	{
		UINT nextFrameIndex = mFrameIndex + 1;
		mFrameIndex = nextFrameIndex;

		HANDLE waitableObjects[] = { mSwapChainWaitableObject, nullptr };
		DWORD numWaitableObjects = 1;

		FrameContext* frameCtx = &mFrameContext[nextFrameIndex % kNumFramesInFlight];
		UINT64 fenceValue = frameCtx->FenceValue;
		if (fenceValue != 0) // means no fence was signaled
		{
			frameCtx->FenceValue = 0;
			mFence->SetEventOnCompletion(fenceValue, mFenceEvent);
			waitableObjects[1] = mFenceEvent;
			numWaitableObjects = 2;
		}

		WaitForMultipleObjects(numWaitableObjects, waitableObjects, TRUE, INFINITE);

		return frameCtx;
	}

	void ResizeSwapChain(int width, int height)
	{
		DXGI_SWAP_CHAIN_DESC1 sd;
		mSwapChain->GetDesc1(&sd);
		sd.Width = width;
		sd.Height = height;

		IDXGIFactory4* dxgiFactory = nullptr;
		mSwapChain->GetParent(IID_PPV_ARGS(&dxgiFactory));

		mSwapChain->Release();
		CloseHandle(mSwapChainWaitableObject);

		IDXGISwapChain1* swapChain1 = nullptr;
		dxgiFactory->CreateSwapChainForHwnd(mD3dCommandQueue, hWnd, &sd, nullptr, nullptr, &swapChain1);
		swapChain1->QueryInterface(IID_PPV_ARGS(&mSwapChain));
		swapChain1->Release();
		dxgiFactory->Release();

		mSwapChain->SetMaximumFrameLatency(kNumBackBuffers);

		mSwapChainWaitableObject = mSwapChain->GetFrameLatencyWaitableObject();
		assert(mSwapChainWaitableObject != nullptr);
	}

	static LRESULT CALLBACK StaticWndProc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
	{
		DirectX12Backend* self;
		if (uMsg == WM_NCCREATE) {
			auto lpcs = reinterpret_cast<LPCREATESTRUCT>(lParam);
			self = static_cast<DirectX12Backend*>(lpcs->lpCreateParams);
			self->hWnd = hWnd;
			SetWindowLongPtr(hWnd, GWLP_USERDATA, reinterpret_cast<LONG_PTR>(self));
		} else {
			self = reinterpret_cast<DirectX12Backend*>(GetWindowLongPtr(hWnd, GWLP_USERDATA));
		}

		if (self) {
			return self->WndProc(uMsg, wParam, lParam);
		} else {
			return DefWindowProc(hWnd, uMsg, wParam, lParam);
		}
	}

	LRESULT WndProc(UINT msg, WPARAM wParam, LPARAM lParam)
	{
		if (ImGui_ImplWin32_WndProcHandler(hWnd, msg, wParam, lParam)) {
			return true;
		}

		switch (msg) {
			case WM_SIZE: {
				if (mD3dDevice != nullptr && wParam != SIZE_MINIMIZED) {
					WaitForLastSubmittedFrame();
					ImGui_ImplDX12_InvalidateDeviceObjects();
					CleanupRenderTarget();
					ResizeSwapChain((UINT)LOWORD(lParam), (UINT)HIWORD(lParam));
					CreateRenderTarget();
					ImGui_ImplDX12_CreateDeviceObjects();
				}
				return 0;
			}

			case WM_SYSCOMMAND: {
				// Disable ALT application menu
				if ((wParam & 0xfff0) == SC_KEYMENU) {
					return 0;
				}
			} break;

			case WM_DESTROY: {
				::PostQuitMessage(0);
				return 0;
			}
		}
		return ::DefWindowProc(hWnd, msg, wParam, lParam);
	}
};

std::unique_ptr<RenderingBackend> RenderingBackend::CreateDx12Backend()
{
	try {
		return std::make_unique<DirectX12Backend>();
	} catch (std::exception& e) {
		return nullptr;
	}
}

#else // ^^ BUILD_CORE_WITH_DX12_BACKEND | BUILD_CORE_WITH_DX12_BACKEND vv

std::unique_ptr<RenderingBackend> RenderingBackend::CreateDx12Backend()
{
	return nullptr;
}

#endif