#include "d3d11_window.hpp"

#include <wincodec.h> // WIC PNG encoder for screenshots

#include <imgui_impl_win32.h>

// Forward declared in the ImGui Win32 backend; lets ImGui consume input first.
extern IMGUI_IMPL_API LRESULT ImGui_ImplWin32_WndProcHandler(HWND hwnd, UINT msg, WPARAM wparam, LPARAM lparam);

using Microsoft::WRL::ComPtr;

namespace coop {

namespace {
constexpr wchar_t kWindowClass[] = L"CoopAllTheThingsWindow";

// Encode a tightly-packed/row-pitched RGBA8 image to a PNG file via WIC. `src` is the
// CPU-mapped staging copy of the back buffer (stride = row_pitch). Pure COM, so no
// extra link dependency. Returns false on any failure (caller swallows it -- a missed
// screenshot is never fatal).
bool write_rgba8_png(const std::wstring& path, UINT width, UINT height, const BYTE* src, UINT row_pitch)
{
	ComPtr<IWICImagingFactory> factory;
	if (FAILED(CoCreateInstance(CLSID_WICImagingFactory, nullptr, CLSCTX_INPROC_SERVER,
	                            IID_PPV_ARGS(factory.GetAddressOf())))) {
		return false;
	}
	ComPtr<IWICBitmap> bitmap; // wrap the back-buffer bytes (RGBA, matches the swap chain)
	if (FAILED(factory->CreateBitmapFromMemory(width, height, GUID_WICPixelFormat32bppRGBA, row_pitch,
	                                           row_pitch * height, const_cast<BYTE*>(src), bitmap.GetAddressOf()))) {
		return false;
	}
	ComPtr<IWICStream> stream;
	if (FAILED(factory->CreateStream(stream.GetAddressOf()))
	    || FAILED(stream->InitializeFromFilename(path.c_str(), GENERIC_WRITE))) {
		return false;
	}
	ComPtr<IWICBitmapEncoder> encoder;
	if (FAILED(factory->CreateEncoder(GUID_ContainerFormatPng, nullptr, encoder.GetAddressOf()))
	    || FAILED(encoder->Initialize(stream.Get(), WICBitmapEncoderNoCache))) {
		return false;
	}
	ComPtr<IWICBitmapFrameEncode> frame;
	ComPtr<IPropertyBag2> props;
	if (FAILED(encoder->CreateNewFrame(frame.GetAddressOf(), props.GetAddressOf()))
	    || FAILED(frame->Initialize(props.Get())) || FAILED(frame->SetSize(width, height))) {
		return false;
	}
	// Let the encoder pick its native pixel format; WriteSource converts our RGBA to it.
	WICPixelFormatGUID fmt = GUID_WICPixelFormat32bppBGRA;
	frame->SetPixelFormat(&fmt);
	if (FAILED(frame->WriteSource(bitmap.Get(), nullptr)) || FAILED(frame->Commit()) || FAILED(encoder->Commit())) {
		return false;
	}
	return true;
}
} // namespace

D3D11Window::~D3D11Window()
{
	release_render_target();
	if (hwnd_ != nullptr) {
		DestroyWindow(hwnd_);
		hwnd_ = nullptr;
	}
	UnregisterClassW(kWindowClass, GetModuleHandleW(nullptr));
}

bool D3D11Window::create(const wchar_t* title)
{
	const HINSTANCE instance = GetModuleHandleW(nullptr);

	WNDCLASSEXW wc = {};
	wc.cbSize = sizeof(wc);
	wc.style = CS_HREDRAW | CS_VREDRAW;
	wc.lpfnWndProc = &D3D11Window::wnd_proc;
	wc.hInstance = instance;
	wc.hCursor = LoadCursorW(nullptr, IDC_ARROW);
	wc.lpszClassName = kWindowClass;
	if (RegisterClassExW(&wc) == 0) {
		return false;
	}

	// Borderless popup covering the primary monitor. Steam RPT only streams the
	// focused window, and exclusive-fullscreen swap chains can't be captured, so
	// a plain WS_POPUP is exactly what we want.
	const int width = GetSystemMetrics(SM_CXSCREEN);
	const int height = GetSystemMetrics(SM_CYSCREEN);

	hwnd_ = CreateWindowExW(0, kWindowClass, title, WS_POPUP, 0, 0, width, height, nullptr, nullptr, instance, this);
	if (hwnd_ == nullptr) {
		return false;
	}

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

	ShowWindow(hwnd_, SW_SHOW);
	UpdateWindow(hwnd_);
	SetForegroundWindow(hwnd_);
	return true;
}

bool D3D11Window::create_device()
{
	DXGI_SWAP_CHAIN_DESC1 desc = {};
	desc.Width = 0; // derive from the window client area
	desc.Height = 0;
	desc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
	desc.SampleDesc.Count = 1;
	desc.BufferUsage = DXGI_USAGE_RENDER_TARGET_OUTPUT;
	desc.BufferCount = 2;
	desc.Scaling = DXGI_SCALING_STRETCH;
	desc.SwapEffect = DXGI_SWAP_EFFECT_FLIP_DISCARD;
	desc.AlphaMode = DXGI_ALPHA_MODE_IGNORE;

	UINT flags = 0;
#ifdef _DEBUG
	flags |= D3D11_CREATE_DEVICE_DEBUG;
#endif
	const D3D_FEATURE_LEVEL levels[] = {D3D_FEATURE_LEVEL_11_1, D3D_FEATURE_LEVEL_11_0};

	if (FAILED(D3D11CreateDevice(nullptr, D3D_DRIVER_TYPE_HARDWARE, nullptr, flags, levels, _countof(levels),
	                             D3D11_SDK_VERSION, device_.GetAddressOf(), nullptr, context_.GetAddressOf()))) {
		return false;
	}

	ComPtr<IDXGIDevice> dxgi_device;
	if (FAILED(device_.As(&dxgi_device))) {
		return false;
	}
	ComPtr<IDXGIAdapter> adapter;
	if (FAILED(dxgi_device->GetAdapter(adapter.GetAddressOf()))) {
		return false;
	}
	ComPtr<IDXGIFactory2> factory;
	if (FAILED(adapter->GetParent(IID_PPV_ARGS(factory.GetAddressOf())))) {
		return false;
	}

	if (FAILED(factory->CreateSwapChainForHwnd(device_.Get(), hwnd_, &desc, nullptr, nullptr,
	                                           swap_chain_.GetAddressOf()))) {
		return false;
	}
	// Don't let DXGI swallow Alt+Enter into an exclusive-fullscreen transition.
	factory->MakeWindowAssociation(hwnd_, DXGI_MWA_NO_ALT_ENTER);

	create_render_target();
	return true;
}

bool D3D11Window::note_device_loss(HRESULT hr)
{
	if (hr != DXGI_ERROR_DEVICE_REMOVED && hr != DXGI_ERROR_DEVICE_RESET) {
		return false;
	}
	// GetDeviceRemovedReason gives the specific cause (HUNG / driver internal / removed); a plain
	// RESET may report S_OK there, so fall back to the originating error in that case.
	const HRESULT reason = device_ ? device_->GetDeviceRemovedReason() : hr;
	device_lost_reason_ = (reason != S_OK) ? reason : hr;
	device_lost_ = true;
	return true;
}

void D3D11Window::create_render_target()
{
	ComPtr<ID3D11Texture2D> back_buffer;
	if (SUCCEEDED(swap_chain_->GetBuffer(0, IID_PPV_ARGS(back_buffer.GetAddressOf())))) {
		const HRESULT hr = device_->CreateRenderTargetView(back_buffer.Get(), nullptr, rtv_.ReleaseAndGetAddressOf());
		note_device_loss(hr); // a removed device surfaces here too; the render loop checks device_lost()
	}
}

void D3D11Window::release_render_target()
{
	rtv_.Reset();
}

void D3D11Window::handle_resize(UINT width, UINT height)
{
	if (swap_chain_ == nullptr || width == 0 || height == 0) {
		return;
	}
	release_render_target();
	const HRESULT hr = swap_chain_->ResizeBuffers(0, width, height, DXGI_FORMAT_UNKNOWN, 0);
	if (note_device_loss(hr)) {
		return; // device gone; the render loop will see device_lost() and stop
	}
	create_render_target();
}

bool D3D11Window::pump_messages()
{
	MSG msg;
	while (PeekMessageW(&msg, nullptr, 0, 0, PM_REMOVE)) {
		if (msg.message == WM_QUIT) {
			return false;
		}
		TranslateMessage(&msg);
		DispatchMessageW(&msg);
	}
	if (resize_pending_) {
		handle_resize(resize_width_, resize_height_);
		resize_pending_ = false;
	}
	return true;
}

void D3D11Window::render_frame(const RenderCallback& render, UINT sync_interval)
{
	const float clear[4] = {0.06f, 0.06f, 0.08f, 1.0f};
	context_->OMSetRenderTargets(1, rtv_.GetAddressOf(), nullptr);
	context_->ClearRenderTargetView(rtv_.Get(), clear);

	if (render) {
		render();
	}

	// Screenshot (F10): capture after the overlay is drawn but before Present -- the
	// flip-model back buffer is undefined once presented.
	if (!pending_screenshot_.empty()) {
		if (save_backbuffer_png(pending_screenshot_)) {
			saved_screenshot_ = pending_screenshot_;
		}
		pending_screenshot_.clear();
	}

	// sync_interval 1 (default) vsyncs to the monitor; 0 presents immediately so the
	// caller can pace the flip itself (frame-sync to the game's published frames).
	const HRESULT hr = swap_chain_->Present(sync_interval, 0);
	note_device_loss(hr); // a TDR/driver reset on the host surfaces here; the render loop halts on it
}

void D3D11Window::request_screenshot(std::wstring path)
{
	pending_screenshot_ = std::move(path);
}

std::wstring D3D11Window::take_screenshot_result()
{
	std::wstring result;
	result.swap(saved_screenshot_);
	return result;
}

bool D3D11Window::save_backbuffer_png(const std::wstring& path)
{
	ComPtr<ID3D11Texture2D> back;
	if (FAILED(swap_chain_->GetBuffer(0, IID_PPV_ARGS(back.GetAddressOf())))) {
		return false;
	}
	D3D11_TEXTURE2D_DESC desc{};
	back->GetDesc(&desc);

	// A CPU-readable staging copy we can Map (the back buffer itself isn't readable).
	D3D11_TEXTURE2D_DESC staging = desc;
	staging.Usage = D3D11_USAGE_STAGING;
	staging.BindFlags = 0;
	staging.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
	staging.MiscFlags = 0;
	ComPtr<ID3D11Texture2D> cpu;
	if (FAILED(device_->CreateTexture2D(&staging, nullptr, cpu.GetAddressOf()))) {
		return false;
	}
	context_->CopyResource(cpu.Get(), back.Get());

	D3D11_MAPPED_SUBRESOURCE map{};
	if (FAILED(context_->Map(cpu.Get(), 0, D3D11_MAP_READ, 0, &map))) {
		return false;
	}
	// The swap chain is DXGI_FORMAT_R8G8B8A8_UNORM (see create_device), i.e. RGBA bytes.
	const bool ok = write_rgba8_png(path, desc.Width, desc.Height, static_cast<const BYTE*>(map.pData), map.RowPitch);
	context_->Unmap(cpu.Get(), 0);
	return ok;
}

LRESULT CALLBACK D3D11Window::wnd_proc(HWND hwnd, UINT msg, WPARAM wparam, LPARAM lparam)
{
	if (msg == WM_NCCREATE) {
		auto* create = reinterpret_cast<CREATESTRUCTW*>(lparam);
		SetWindowLongPtrW(hwnd, GWLP_USERDATA, reinterpret_cast<LONG_PTR>(create->lpCreateParams));
	}

	if (ImGui_ImplWin32_WndProcHandler(hwnd, msg, wparam, lparam)) {
		return true;
	}

	auto* self = reinterpret_cast<D3D11Window*>(GetWindowLongPtrW(hwnd, GWLP_USERDATA));

	switch (msg) {
	case WM_SIZE:
		if (self != nullptr && wparam != SIZE_MINIMIZED) {
			self->resize_pending_ = true;
			self->resize_width_ = LOWORD(lparam);
			self->resize_height_ = HIWORD(lparam);
		}
		return 0;
	case WM_DPICHANGED:
		// Per-monitor-v2: the DPI of the display we're on changed. Resize to the rect Windows suggests
		// in lparam (its recommended handling; the resulting WM_SIZE repaints the swap chain via the
		// deferred-resize path above), then latch the new DPI for the overlay to rescale its font/style.
		if (self != nullptr) {
			if (const auto* suggested = reinterpret_cast<const RECT*>(lparam); suggested != nullptr) {
				SetWindowPos(hwnd, nullptr, suggested->left, suggested->top, suggested->right - suggested->left,
				             suggested->bottom - suggested->top, SWP_NOZORDER | SWP_NOACTIVATE);
			}
			self->dpi_pending_ = true;
			self->pending_dpi_ = HIWORD(wparam); // X and Y DPI are equal; HIWORD is the Y value
		}
		return 0;
	case WM_SYSKEYDOWN:
		// F10 is our screenshot key; ImGui already saw this message (handler runs above),
		// so swallow it here to stop DefWindowProc from flicking into Win32 menu mode.
		// Alt+F4 (VK_F4) falls through to DefWindowProc so it still closes the window.
		if (wparam == VK_F10) {
			return 0;
		}
		break;
	case WM_DESTROY:
		// Reached by the close button, Alt+F4 (DefWindowProc turns it into WM_CLOSE ->
		// DestroyWindow), and our own teardown. Esc deliberately does NOT quit -- it's a
		// common in-game key, so quitting is via the File -> Exit menu item or Alt+F4.
		PostQuitMessage(0);
		return 0;
	default:
		break;
	}
	return DefWindowProcW(hwnd, msg, wparam, lparam);
}

} // namespace coop