#include "PngScreenshotWriter.h"

#include <windows.h>
#include <wincodec.h>
#include <atlbase.h>

#include <sstream>
#include <thread>

namespace
{
std::string HResultToString(HRESULT hr)
{
	std::ostringstream stream;
	stream << "HRESULT 0x" << std::hex << static_cast<unsigned long>(hr);
	return stream.str();
}

bool WritePngFile(
	const std::filesystem::path& outputPath,
	unsigned width,
	unsigned height,
	const std::vector<unsigned char>& bgraPixels,
	std::string& error)
{
	if (width == 0 || height == 0 || bgraPixels.size() < static_cast<std::size_t>(width) * height * 4)
	{
		error = "Invalid screenshot dimensions or pixel buffer.";
		return false;
	}

	HRESULT initializeResult = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
	const bool shouldUninitialize = SUCCEEDED(initializeResult);
	if (FAILED(initializeResult) && initializeResult != RPC_E_CHANGED_MODE)
	{
		error = "CoInitializeEx failed: " + HResultToString(initializeResult);
		return false;
	}

	CComPtr<IWICImagingFactory> factory;
	HRESULT result = CoCreateInstance(
		CLSID_WICImagingFactory,
		nullptr,
		CLSCTX_INPROC_SERVER,
		IID_PPV_ARGS(&factory));
	if (FAILED(result))
	{
		error = "Could not create WIC imaging factory: " + HResultToString(result);
		if (shouldUninitialize)
			CoUninitialize();
		return false;
	}

	CComPtr<IWICStream> stream;
	result = factory->CreateStream(&stream);
	if (SUCCEEDED(result))
		result = stream->InitializeFromFilename(outputPath.wstring().c_str(), GENERIC_WRITE);
	if (FAILED(result))
	{
		error = "Could not open screenshot output file: " + HResultToString(result);
		if (shouldUninitialize)
			CoUninitialize();
		return false;
	}

	CComPtr<IWICBitmapEncoder> encoder;
	result = factory->CreateEncoder(GUID_ContainerFormatPng, nullptr, &encoder);
	if (SUCCEEDED(result))
		result = encoder->Initialize(stream, WICBitmapEncoderNoCache);
	if (FAILED(result))
	{
		error = "Could not initialize PNG encoder: " + HResultToString(result);
		if (shouldUninitialize)
			CoUninitialize();
		return false;
	}

	CComPtr<IWICBitmapFrameEncode> frame;
	CComPtr<IPropertyBag2> propertyBag;
	result = encoder->CreateNewFrame(&frame, &propertyBag);
	if (SUCCEEDED(result))
		result = frame->Initialize(propertyBag);
	if (SUCCEEDED(result))
		result = frame->SetSize(width, height);

	WICPixelFormatGUID pixelFormat = GUID_WICPixelFormat32bppBGRA;
	if (SUCCEEDED(result))
		result = frame->SetPixelFormat(&pixelFormat);
	if (SUCCEEDED(result) && pixelFormat != GUID_WICPixelFormat32bppBGRA)
	{
		error = "PNG encoder did not accept BGRA pixel format.";
		result = E_FAIL;
	}

	const UINT stride = width * 4;
	const UINT imageSize = stride * height;
	if (SUCCEEDED(result))
		result = frame->WritePixels(height, stride, imageSize, const_cast<BYTE*>(bgraPixels.data()));
	if (SUCCEEDED(result))
		result = frame->Commit();
	if (SUCCEEDED(result))
		result = encoder->Commit();

	if (shouldUninitialize)
		CoUninitialize();

	if (FAILED(result))
	{
		error = "Could not write screenshot PNG: " + HResultToString(result);
		std::error_code ignored;
		std::filesystem::remove(outputPath, ignored);
		return false;
	}

	return true;
}
}

void WritePngFileAsync(
	const std::filesystem::path& outputPath,
	unsigned width,
	unsigned height,
	std::vector<unsigned char> rgbaPixels)
{
	std::thread(
		[outputPath, width, height, pixels = std::move(rgbaPixels)]() mutable
		{
			for (std::size_t index = 0; index + 3 < pixels.size(); index += 4)
				std::swap(pixels[index], pixels[index + 2]);

			std::string error;
			if (!WritePngFile(outputPath, width, height, pixels, error))
				OutputDebugStringA(("Screenshot write failed: " + error + "\n").c_str());
			else
				OutputDebugStringA(("Screenshot written: " + outputPath.string() + "\n").c_str());
		}).detach();
}