#include "stdafx.h"
#include "ScreenshotManager.h"

#include <d3d11.h>
#include <dxgi.h>
#include <wincodec.h>
#include <shlobj.h>
#include <chrono>
#include <iomanip>
#include <sstream>
#include <vector>

// link WIC
#pragma comment(lib, "windowscodecs.lib")

// defined in Windows64_Minecraft.cpp
extern IDXGISwapChain*			g_pSwapChain;
extern ID3D11Device*			g_pd3dDevice;
extern ID3D11DeviceContext*		g_pImmediateContext;
extern int						g_rScreenWidth;
extern int						g_rScreenHeight;

// make sure the output foldler exists
bool ScreenshotManager::ensureScreenshotsFolder(std::wstring& outPath) {
	wchar_t exePath[MAX_PATH] = {};
	if (!GetModuleFileNameW(nullptr, exePath, _countof(exePath))) return false;
	wchar_t* lastSlash = wcsrchr(exePath, L'\\');
	if (lastSlash) *(lastSlash + 1) = L'\0';

	std::wstring folder = exePath;
	folder += L"screenshots";
	// create if doesn't exist
	DWORD attrib = GetFileAttributesW(folder.c_str());
	if (attrib == INVALID_FILE_ATTRIBUTES || !(attrib & FILE_ATTRIBUTE_DIRECTORY)) {
		if (!CreateDirectoryW(folder.c_str(), nullptr)) {
			return false;
		}
	}

	outPath = folder + L"\\";
	return true;
}

// create a formatted name with the current timestamp
std::wstring ScreenshotManager::makeTimestampedFilename() {
	auto now = std::chrono::system_clock::now();
	auto t = std::chrono::system_clock::to_time_t(now);
	std::wstringstream ss;
	ss << L"screenshot_";
	std::tm tm{};
	localtime_s(&tm, &t);
	ss << std::put_time(&tm, L"%Y-%m-%d_%H-%M-%S") << ".png";
	return ss.str();
}

// takes a screenshot of the current frame
bool ScreenshotManager::takeScreenshot(const std::wstring& path) {
	if (!g_pSwapChain || !g_pd3dDevice || !g_pImmediateContext) return false;

	HRESULT hr = S_OK;
	ID3D11Texture2D* pBack = nullptr;
	hr = g_pSwapChain->GetBuffer(0, __uuidof(ID3D11Texture2D), (void**) &pBack);
	if (FAILED(hr) || !pBack) return false;

	D3D11_TEXTURE2D_DESC desc;
	pBack->GetDesc(&desc);

	D3D11_TEXTURE2D_DESC stagingDesc = desc;
	stagingDesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
	stagingDesc.Usage = D3D11_USAGE_STAGING;
	stagingDesc.BindFlags = 0;
	stagingDesc.MiscFlags = 0;
	stagingDesc.SampleDesc.Count = 1;
	stagingDesc.SampleDesc.Quality = 0;
	ID3D11Texture2D* pStaging = nullptr;
	hr = g_pd3dDevice->CreateTexture2D(&stagingDesc, nullptr, &pStaging);

	if (FAILED(hr) || !pStaging) { 
		pBack->Release(); 
		return false; 
	}

	g_pImmediateContext->CopyResource(pStaging, pBack);

	// map and read
	D3D11_MAPPED_SUBRESOURCE mapped;
	hr = g_pImmediateContext->Map(pStaging, 0, D3D11_MAP_READ, 0, &mapped);
	if (FAILED(hr)) { 
		pStaging->Release(); 
		pBack->Release(); 
		return false; 
	}

	// prepare WIC
	IWICImagingFactory* pFactory = nullptr;

	hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
	if (FAILED(hr) && hr != RPC_E_CHANGED_MODE) { 
		g_pImmediateContext->Unmap(pStaging, 0); 
		pStaging->Release(); 
		pBack->Release(); 
		return false; 
	}
	
	hr = CoCreateInstance(CLSID_WICImagingFactory, nullptr, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&pFactory));
	if (FAILED(hr)) { 
		g_pImmediateContext->Unmap(pStaging, 0); 
		pStaging->Release(); 
		pBack->Release(); 
		return false; 
	}

	std::wstring outFolder;
	if (!ensureScreenshotsFolder(outFolder)) {
		pFactory->Release();
		g_pImmediateContext->Unmap(pStaging, 0);
		pStaging->Release();
		pBack->Release();
		return false;
	}

	std::wstring fileName = path;
	if (fileName.empty()) {
		fileName = outFolder + makeTimestampedFilename();
	} else {
		if (fileName.find(L":\\") == std::wstring::npos) {
			fileName = outFolder + fileName;
		}
	}

	GUID wicFormat = GUID_WICPixelFormat32bppBGRA;

	const UINT width = desc.Width;
	const UINT height = desc.Height;
	const UINT srcRowPitch = mapped.RowPitch;
	const UINT dstStride = width * 4;

	// write data and shi
	std::vector<BYTE> imageData(dstStride * height);
	BYTE* dst = imageData.data();
	BYTE* srcBase = reinterpret_cast<BYTE*>(mapped.pData);
	for (UINT y = 0; y < height; ++y) {
		BYTE* src = srcBase + y * srcRowPitch;
		BYTE* row = dst + y * dstStride;
		for (UINT x = 0; x < width; ++x) {
			// read original rgba
			BYTE r = src[x * 4 + 0];
			BYTE g = src[x * 4 + 1];
			BYTE b = src[x * 4 + 2];
			BYTE a = src[x * 4 + 3];

			// black background to fix opacity issues (sorta)
			// (using integer math with rounding: (value * a + 127) / 255)
			UINT br = (r * a + 127) / 255;
			UINT bg = (g * a + 127) / 255;
			UINT bb = (b * a + 127) / 255;

			// we swizzle R / B because WIC expects BGRA for some fucking reason
			row[x * 4 + 0] = static_cast<BYTE>(bb);
			row[x * 4 + 1] = static_cast<BYTE>(bg);
			row[x * 4 + 2] = static_cast<BYTE>(br);
			row[x * 4 + 3] = 255;
		}
	}

	// create WIC bitmap from memory
	IWICBitmap* pBitmap = nullptr;
	hr = pFactory->CreateBitmapFromMemory(width, height, wicFormat, dstStride, static_cast<UINT>(imageData.size()), imageData.data(), &pBitmap);
	if (FAILED(hr)) { pFactory->Release(); g_pImmediateContext->Unmap(pStaging, 0); pStaging->Release(); pBack->Release(); return false; }

	// create encoder
	IWICStream* pStream = nullptr;
	IWICBitmapEncoder* pEncoder = nullptr;
	hr = pFactory->CreateStream(&pStream);
	if (FAILED(hr)) { pBitmap->Release(); pFactory->Release(); g_pImmediateContext->Unmap(pStaging, 0); pStaging->Release(); pBack->Release(); return false; }

	hr = pStream->InitializeFromFilename(fileName.c_str(), GENERIC_WRITE);
	if (FAILED(hr)) { pStream->Release(); pBitmap->Release(); pFactory->Release(); g_pImmediateContext->Unmap(pStaging, 0); pStaging->Release(); pBack->Release(); return false; }

	hr = pFactory->CreateEncoder(GUID_ContainerFormatPng, nullptr, &pEncoder);
	if (FAILED(hr)) { pStream->Release(); pBitmap->Release(); pFactory->Release(); g_pImmediateContext->Unmap(pStaging, 0); pStaging->Release(); pBack->Release(); return false; }

	hr = pEncoder->Initialize(pStream, WICBitmapEncoderNoCache);
	if (FAILED(hr)) { pEncoder->Release(); pStream->Release(); pBitmap->Release(); pFactory->Release(); g_pImmediateContext->Unmap(pStaging, 0); pStaging->Release(); pBack->Release(); return false; }

	IWICBitmapFrameEncode* pFrame = nullptr;
	IPropertyBag2* pProps = nullptr;
	hr = pEncoder->CreateNewFrame(&pFrame, &pProps);
	if (FAILED(hr)) { pEncoder->Release(); pStream->Release(); pBitmap->Release(); pFactory->Release(); g_pImmediateContext->Unmap(pStaging, 0); pStaging->Release(); pBack->Release(); return false; }

	hr = pFrame->Initialize(pProps);
	if (FAILED(hr)) { pFrame->Release(); if (pProps) pProps->Release(); pEncoder->Release(); pStream->Release(); pBitmap->Release(); pFactory->Release(); g_pImmediateContext->Unmap(pStaging, 0); pStaging->Release(); pBack->Release(); return false; }

	hr = pFrame->SetSize(width, height);
	if (FAILED(hr)) { pFrame->Release(); if (pProps) pProps->Release(); pEncoder->Release(); pStream->Release(); pBitmap->Release(); pFactory->Release(); g_pImmediateContext->Unmap(pStaging, 0); pStaging->Release(); pBack->Release(); return false; }

	hr = pFrame->SetPixelFormat(&wicFormat);  // prefer BGRA
	hr = pFrame->WriteSource(pBitmap, nullptr);
	hr = pFrame->Commit();
	hr = pEncoder->Commit();

	// cleanup
	if (pProps) pProps->Release();
	pFrame->Release();
	pEncoder->Release();
	pStream->Release();
	pBitmap->Release();
	pFactory->Release();

	g_pImmediateContext->Unmap(pStaging, 0);
	pStaging->Release();
	pBack->Release();

	CoUninitialize();

	return true;
}