#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "BattleFireDirect.h"
#include "StaticMeshComponent.h"
#include "stbi/stb_image.h"
#include "Utils.h"
#include "NanoVDBLoader.h"

#pragma comment(lib,"d3d12.lib")
#pragma comment(lib,"dxgi.lib")
#pragma comment(lib,"d3dcompiler.lib")
#pragma comment(lib,"winmm.lib")

LPCWSTR gWindowClassName = L"BattleFire";

LRESULT CALLBACK WindowProc(HWND inHWND, UINT inMSG, WPARAM inWParam, LPARAM inLParam) {
	switch (inMSG) {
	case WM_CLOSE:
		PostQuitMessage(0);
		break;
	}
	return DefWindowProc(inHWND, inMSG, inWParam, inLParam);
}
int WINAPI wWinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPWSTR lpCmdLine, int inShowCmd) {
	AttachConsole(ATTACH_PARENT_PROCESS);
	freopen("CONOUT$", "w", stdout);
	freopen("CONOUT$", "w", stderr);
	printf("XCVolumeRenderer started\n");
	printf("Initializing D3D12...\n");
	//register
	WNDCLASSEX wndClassEx;
	wndClassEx.cbSize = sizeof(WNDCLASSEX);
	wndClassEx.style = CS_HREDRAW | CS_VREDRAW;
	wndClassEx.cbClsExtra = NULL;//class
	wndClassEx.cbWndExtra = NULL;//instance
	wndClassEx.hInstance = hInstance;
	wndClassEx.hIcon = LoadIcon(NULL, IDI_APPLICATION);
	wndClassEx.hIconSm = LoadIcon(NULL, IDI_APPLICATION);
	wndClassEx.hCursor = LoadCursor(NULL, IDC_ARROW);
	wndClassEx.hbrBackground = NULL;
	wndClassEx.lpszMenuName = NULL;
	wndClassEx.lpszClassName = gWindowClassName;
	wndClassEx.lpfnWndProc = WindowProc;
	if (!RegisterClassEx(&wndClassEx)) {
		MessageBox(NULL, L"Register Class Failed!", L"Error", MB_OK | MB_ICONERROR);
		return -1;
	}
	//create
	int viewportWidth = 1280;
	int viewportHeight = 720;
	RECT rect;
	rect.left = 0;
	rect.top = 0;
	rect.right = viewportWidth;
	rect.bottom = viewportHeight;
	AdjustWindowRect(&rect, WS_OVERLAPPEDWINDOW, FALSE);
	int windowWidth = rect.right - rect.left;
	int windowHeight = rect.bottom - rect.top;
	HWND hwnd = CreateWindowEx(NULL,
		gWindowClassName,
		L"My Render Window",
		WS_OVERLAPPEDWINDOW,
		CW_USEDEFAULT, CW_USEDEFAULT,
		windowWidth, windowHeight,
		NULL,
		NULL,
		hInstance,
		NULL);
	if (!hwnd) {
		MessageBox(NULL, L"Create Window Failed!", L"Error", MB_OK | MB_ICONERROR);
		return -1;
	}
	//show
	InitD3D12(hwnd, 1280, 720);
	printf("D3D12 initialized\n");
	ID3D12GraphicsCommandList* commandList = GetCommandList();
	ID3D12CommandAllocator* commandAllocator = GetCommandAllocator();

	NanoVDBData nanoVDBData = {};
	LoadNanoVDB("Res/NanoVDB/cloud.nvdb", nanoVDBData, commandList, commandAllocator);
	printf("NanoVDB loaded: %llu bytes\n", (unsigned long long)nanoVDBData.byteSize);
	printf("NanoVDB BBox: [%.2f, %.2f, %.2f] - [%.2f, %.2f, %.2f]\n",
		nanoVDBData.worldBBox[0], nanoVDBData.worldBBox[1], nanoVDBData.worldBBox[2],
		nanoVDBData.worldBBox[3], nanoVDBData.worldBBox[4], nanoVDBData.worldBBox[5]);

	float bboxMin[3] = { (float)nanoVDBData.worldBBox[0], (float)nanoVDBData.worldBBox[1], (float)nanoVDBData.worldBBox[2] };
	float bboxMax[3] = { (float)nanoVDBData.worldBBox[3], (float)nanoVDBData.worldBBox[4], (float)nanoVDBData.worldBBox[5] };
	float bboxCenter[3] = { (bboxMin[0] + bboxMax[0]) / 2, (bboxMin[1] + bboxMax[1]) / 2, (bboxMin[2] + bboxMax[2]) / 2 };
	float bboxSize[3] = { bboxMax[0] - bboxMin[0], bboxMax[1] - bboxMin[1], bboxMax[2] - bboxMin[2] };
	printf("BBox center: [%.2f, %.2f, %.2f], size: [%.2f, %.2f, %.2f]\n",
		bboxCenter[0], bboxCenter[1], bboxCenter[2],
		bboxSize[0], bboxSize[1], bboxSize[2]);

	float quadVertices[] = {
		-1.0f, -1.0f, 0.0f, 0.0f,
		 1.0f, -1.0f, 1.0f, 0.0f,
		 1.0f,  1.0f, 1.0f, 1.0f,
		-1.0f,  1.0f, 0.0f, 1.0f,
	};
	uint32_t quadIndices[] = { 0, 1, 2, 0, 2, 3 };
	ID3D12Resource* quadVBO = CreateBufferObject(commandList, quadVertices, sizeof(quadVertices), D3D12_RESOURCE_STATE_VERTEX_AND_CONSTANT_BUFFER);
	ID3D12Resource* quadIBO = CreateBufferObject(commandList, quadIndices, sizeof(quadIndices), D3D12_RESOURCE_STATE_INDEX_BUFFER);
	printf("Quad mesh created\n");

	StaticMeshComponent staticMeshComponent;
	staticMeshComponent.InitFromFile(commandList, "Res/Model/Sphere.lhsm");
	printf("Mesh loaded\n");

	ID3D12RootSignature* rootSignature = InitRootSignature();
	printf("Root signature created\n");
	D3D12_SHADER_BYTECODE vs,gs,ps;
	CreateShaderFromFile(L"Res/Shader/gs.hlsl", "MainVS", "vs_5_1", &vs);
	CreateShaderFromFile(L"Res/Shader/gs.hlsl", "MainGS", "gs_5_1", &gs);
	CreateShaderFromFile(L"Res/Shader/gs.hlsl", "MainPS", "ps_5_1", &ps);
	ID3D12PipelineState*pso=CreatePSO(rootSignature, vs, ps, gs);
	printf("PSO created\n");

	ID3D12RootSignature* volumeRootSignature = InitVolumeRootSignature();
	printf("Volume root signature created\n");
	
	D3D12_SHADER_BYTECODE volumeVS, volumePS;
	memset(&volumeVS, 0, sizeof(volumeVS));
	memset(&volumePS, 0, sizeof(volumePS));
	CreateShaderFromFile(L"Res/Shader/volume.hlsl", "MainVS", "vs_5_1", &volumeVS);
	CreateShaderFromFile(L"Res/Shader/volume.hlsl", "MainPS", "ps_5_1", &volumePS);
	printf("Volume VS: ptr=%p, size=%zu\n", volumeVS.pShaderBytecode, volumeVS.BytecodeLength);
	printf("Volume PS: ptr=%p, size=%zu\n", volumePS.pShaderBytecode, volumePS.BytecodeLength);
	
	ID3D12PipelineState* quadPSO = CreateQuadPSO(volumeRootSignature, volumeVS, volumePS);
	if (!quadPSO) {
		printf("Quad PSO creation failed!\n");
	} else {
		printf("Quad PSO created: %p\n", quadPSO);
	}

	ID3D12Resource* volumeCB = CreateConstantBufferObject(65536);

	ID3D12Resource* cb = CreateConstantBufferObject(65536);//1024x64(4x4)
	DirectX::XMMATRIX projectionMatrix=DirectX::XMMatrixPerspectiveFovLH(
		(45.0f*3.141592f)/180.0f,1280.0f/720.0f,0.1f,1000.0f);
	DirectX::XMMATRIX viewMatrix = DirectX::XMMatrixLookAtLH(
		DirectX::XMVectorSet(-10.0f, 300.0f, -1200.0f, 1.0f),
		DirectX::XMVectorSet(-10.0f, 73.0f, 0.0f, 1.0f),
		DirectX::XMVectorSet(0.0f, 1.0f, 0.0f, 1.0f));
	DirectX::XMMATRIX modelMatrix = DirectX::XMMatrixTranslation(0.0f,0.0f,5.0f);
	//modelMatrix *= DirectX::XMMatrixRotationZ(90.0f*3.141592f/180.0f);
	DirectX::XMFLOAT4X4 tempMatrix;
	float matrices[64];

	DirectX::XMStoreFloat4x4(&tempMatrix, projectionMatrix);
	memcpy(matrices, &tempMatrix, sizeof(float) * 16);
	DirectX::XMStoreFloat4x4(&tempMatrix, viewMatrix);
	memcpy(matrices+16, &tempMatrix, sizeof(float) * 16);
	DirectX::XMStoreFloat4x4(&tempMatrix, modelMatrix);
	memcpy(matrices + 32, &tempMatrix, sizeof(float) * 16);;
	DirectX::XMVECTOR determinant;
	DirectX::XMMATRIX inverseModelMatrix = DirectX::XMMatrixInverse(&determinant, modelMatrix);
	if (DirectX::XMVectorGetX(determinant) != 0.0f) {
		DirectX::XMMATRIX normalMatrix = DirectX::XMMatrixTranspose(inverseModelMatrix);
		DirectX::XMStoreFloat4x4(&tempMatrix, modelMatrix);
		memcpy(matrices + 48, &tempMatrix, sizeof(float) * 16);;
	}
	UpdateConstantBuffer(cb, matrices, sizeof(float) * 64);


	ID3D12Resource* sb = CreateConstantBufferObject(65536);//1024x64(4x4)
	struct MaterialData {
		float r;
	};
	MaterialData* materialDatas = new MaterialData[3000];
	for (int i=0;i<3000;i++){
		materialDatas[i].r = srandom() * 0.1f + 0.1f;//0.0~1.0
	}
	UpdateConstantBuffer(sb, materialDatas, sizeof(MaterialData) * 3000);

	int imageWidth, imageHeight,imageChannel;
	stbi_uc* pixels = stbi_load("Res/Image/earth_d.jpg", &imageWidth, &imageHeight, &imageChannel, 4);
	printf("Texture loaded: %dx%d\n", imageWidth, imageHeight);
	ID3D12Resource* texture = CreateTexture2D(commandList, pixels,
		imageWidth * imageHeight * imageChannel, imageWidth, imageHeight,DXGI_FORMAT_R8G8B8A8_UNORM);
	delete[]pixels;
	ID3D12Device* d3dDevice = GetD3DDevice();

	ID3D12DescriptorHeap* srvHeap = nullptr;
	D3D12_DESCRIPTOR_HEAP_DESC d3dDescriptorHeapDescSRV = {};
	d3dDescriptorHeapDescSRV.NumDescriptors = 3;
	d3dDescriptorHeapDescSRV.Type = D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV;
	d3dDescriptorHeapDescSRV.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE;
	d3dDevice->CreateDescriptorHeap(&d3dDescriptorHeapDescSRV, IID_PPV_ARGS(&srvHeap));

	ID3D12DescriptorHeap* descriptorHeaps[] = {srvHeap};

	D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc = {};
	srvDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
	srvDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
	srvDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
	srvDesc.Texture2D.MipLevels = 1;

	D3D12_CPU_DESCRIPTOR_HANDLE srvHeapPtr = srvHeap->GetCPUDescriptorHandleForHeapStart();
	d3dDevice->CreateShaderResourceView(texture, &srvDesc, srvHeapPtr);
	srvHeapPtr.ptr += d3dDevice->GetDescriptorHandleIncrementSize(D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);

	EndCommandList();
	WaitForCompletionOfCommandList();
	printf("All resources uploaded to GPU\n");
	
	ShowWindow(hwnd, inShowCmd);
	UpdateWindow(hwnd);
	printf("Window shown, entering render loop\n");
	float color[] = {0.5f,0.5f,0.5f,1.0f};
	MSG msg;
	DWORD last_time = timeGetTime();
	DWORD appStartTime = last_time;
	int frameCount = 0;
	while (true){
		ZeroMemory(&msg, sizeof(MSG));
		if (PeekMessage(&msg,NULL,0,0,PM_REMOVE)) {
			if (msg.message == WM_QUIT) {
				break;
			}
			TranslateMessage(&msg);
			DispatchMessage(&msg);
		} else {
			//rendering
			WaitForCompletionOfCommandList();
			DWORD current_time = timeGetTime();//ms
			DWORD frameTime = current_time - last_time;
			DWORD timeSinceAppStartInMS = current_time - appStartTime;
			last_time = current_time;
			float frameTimeInSecond = float(frameTime) / 1000.0f;//second
			float timeSinceAppStartInSecond = float(timeSinceAppStartInMS) / 1000.0f;
			color[0] = timeSinceAppStartInSecond;
			commandAllocator->Reset();
			commandList->Reset(commandAllocator, nullptr);
			BeginRenderToSwapChain(commandList);
			//draw
			commandList->SetPipelineState(pso);
			commandList->SetGraphicsRootSignature(rootSignature);
			commandList->SetDescriptorHeaps(_countof(descriptorHeaps),descriptorHeaps);
			commandList->SetGraphicsRootConstantBufferView(0, cb->GetGPUVirtualAddress());
			commandList->SetGraphicsRoot32BitConstants(1, 4, color, 0);
			commandList->SetGraphicsRootDescriptorTable(2, srvHeap->GetGPUDescriptorHandleForHeapStart());
			commandList->SetGraphicsRootShaderResourceView(3, sb->GetGPUVirtualAddress());
			commandList->IASetPrimitiveTopology(D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
			if (frameCount == 1) {
				printf("Drawing mesh with %d vertices, VBO=%p, subMeshes=%zu\n", 
					staticMeshComponent.mVertexCount, staticMeshComponent.mVBO, staticMeshComponent.mSubMeshes.size());
			}
			staticMeshComponent.Render(commandList);

			DirectX::XMMATRIX viewProj = viewMatrix * projectionMatrix;
			DirectX::XMMATRIX invViewProj = DirectX::XMMatrixInverse(nullptr, viewProj);
			DirectX::XMFLOAT4X4 invViewProjMat;
			DirectX::XMStoreFloat4x4(&invViewProjMat, invViewProj);

			float volumeCBData[36];
			memcpy(volumeCBData, &invViewProjMat, sizeof(float) * 16);
			volumeCBData[16] = -10.0f;
			volumeCBData[17] = 300.0f;
			volumeCBData[18] = -1200.0f;
			volumeCBData[19] = 0.2f;  // DensityScale
			volumeCBData[20] = (float)nanoVDBData.worldBBox[0];
			volumeCBData[21] = (float)nanoVDBData.worldBBox[1];
			volumeCBData[22] = (float)nanoVDBData.worldBBox[2];
			volumeCBData[23] = 1.0f;  // StepSize
			volumeCBData[24] = (float)nanoVDBData.worldBBox[3];
			volumeCBData[25] = (float)nanoVDBData.worldBBox[4];
			volumeCBData[26] = (float)nanoVDBData.worldBBox[5];
			volumeCBData[27] = 2000.0f;  // MaxSteps as float
			volumeCBData[28] = timeSinceAppStartInSecond * 0.3f;  // RotationY
			volumeCBData[29] = 0.0f;  // Pad0
			volumeCBData[30] = 0.0f;  // Pad1
			volumeCBData[31] = 0.0f;  // Pad2
			volumeCBData[32] = 0.5f;  // LightDir X
			volumeCBData[33] = 0.8f;  // LightDir Y
			volumeCBData[34] = 0.3f;  // LightDir Z
			volumeCBData[35] = 8.0f;  // LightSamples as float
			UpdateConstantBuffer(volumeCB, volumeCBData, sizeof(float) * 36);
			if (frameCount == 1) {
				printf("Volume BBox: [%.2f, %.2f, %.2f] - [%.2f, %.2f, %.2f]\n",
					volumeCBData[20], volumeCBData[21], volumeCBData[22],
					volumeCBData[24], volumeCBData[25], volumeCBData[26]);
			}

			if (quadPSO) {
				commandList->SetPipelineState(quadPSO);
				commandList->SetGraphicsRootSignature(volumeRootSignature);
				commandList->SetGraphicsRootConstantBufferView(0, volumeCB->GetGPUVirtualAddress());
				commandList->SetGraphicsRootShaderResourceView(1, nanoVDBData.gpuBuffer->GetGPUVirtualAddress());
				commandList->IASetPrimitiveTopology(D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
				commandList->IASetVertexBuffers(0, 1, &D3D12_VERTEX_BUFFER_VIEW{ quadVBO->GetGPUVirtualAddress(), sizeof(quadVertices), sizeof(float) * 4 });
				commandList->IASetIndexBuffer(&D3D12_INDEX_BUFFER_VIEW{ quadIBO->GetGPUVirtualAddress(), sizeof(quadIndices), DXGI_FORMAT_R32_UINT });
				commandList->DrawIndexedInstanced(6, 1, 0, 0, 0);
			}

			EndRenderToSwapChain(commandList);
			EndCommandList();
			SwapD3D12Buffers();
			frameCount++;
			if (frameCount <= 3) {
				printf("Frame %d rendered\n", frameCount);
			}
		}
	}
	return 0;
}