#include "capture.hpp"

// clang-format off
#include "Utils.hpp"

#include "nvfbc_library.hpp"
#include <NvFBC/nvFBCToSys.h>
// clang-format on

namespace hazelnut {

	void DiffMapWidthHeight(NVFBC_TOSYS_DIFFMAP_BLOCKSIZE blockSize, u32& width, u32& height) {
		switch(blockSize) {
			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_16X16:
				width = 16;
				height = 16;
				break;

			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_32X32:
				width = 32;
				height = 32;
				break;

			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_64X64:
				width = 64;
				height = 64;
				break;

			default:
			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_128X128:
				width = 128;
				height = 128;
				break;
		}
	}

	// From Looking Glass, cleaned up to not be macro slop
	u32 DiffMapDimension(u32 srcDim, NVFBC_TOSYS_DIFFMAP_BLOCKSIZE blockSize) {
		u32 blockShift = 0;

		// Calculate shift
		switch(blockSize) {
			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_16X16: blockShift = 4; break;

			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_32X32: blockShift = 5; break;

			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_64X64: blockShift = 6; break;

			default:
			case NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_128X128: blockShift = 7; break;
		}

		return (((srcDim) + (1 << (blockShift)) - 1) >> (blockShift));
	}

	struct NVFBCDisplayCapture final : public DisplayCaptureBase {
		NVFBCRESULT nvfbcStatus;
		NvFBCLibrary nvfbcLib;

		NVFBC_TOSYS_GRAB_FRAME_PARAMS fbcSysGrabParams;
		NVFBC_TOSYS_DIFFMAP_BLOCKSIZE blockSize;
		NvFBCFrameGrabInfo grabInfo;
		NvFBCToSys* nvfbc;

		u32 width;
		u32 height;

		u32* pRawFramebuffer;
		u8* pDiffMap;

		u32 diffmapWidth;
		u32 diffmapHeight;

		// cached
		u32 tileWidth;
		u32 tileHeight;

	   public:
		virtual ~NVFBCDisplayCapture() { Shutdown(); }

		void Shutdown() {
			NvfbcDestroyInstance();
			nvfbcLib.Close();
		}

		bool Initialize() override {
			if(!nvfbcLib.Load()) {
				return false;
			}

			if(!NvfbcInitSession())
				return false;

			return true;
		}

		void NvfbcDestroyInstance() {
			if(nvfbc) {
				nvfbc->NvFBCToSysRelease();
				nvfbc = nullptr;

				// reset state while we're here I guess
				pRawFramebuffer = nullptr;
				pDiffMap = nullptr;

				width = 0;
				height = 0;
				diffmapWidth = 0;
				diffmapHeight = 0;
			}
		}

		bool NvfbcCreateInstance() {
			// Destroy an existing NvFBC instance to avoid resource leaks.
			if(nvfbc) {
				NvfbcDestroyInstance();
			}

			DWORD maxDisplayWidth = -1;
			DWORD maxDisplayHeight = -1;
			nvfbc = nvfbcLib.CreateToSys(&maxDisplayWidth, &maxDisplayHeight);

			if(!nvfbc) {
				return false;
			}

			return true;
		}

		bool NvfbcInitSession() {
			if(!NvfbcCreateInstance())
				return false;

			blockSize = NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_32X32;

			// set up a session
			NVFBC_TOSYS_SETUP_PARAMS fbcSysSetupParams {};
			fbcSysSetupParams.dwVersion = NVFBC_TOSYS_SETUP_PARAMS_VER;
			fbcSysSetupParams.eMode = NVFBC_TOSYS_ARGB;
			fbcSysSetupParams.bWithHWCursor = true;
			fbcSysSetupParams.ppBuffer = reinterpret_cast<void**>(&pRawFramebuffer);

			// enable the difference map with configured blocksize
			// (currently hardcoded to 32x32).
			fbcSysSetupParams.bDiffMap = TRUE;
			fbcSysSetupParams.ppDiffMap = reinterpret_cast<void**>(&pDiffMap);
			fbcSysSetupParams.eDiffMapBlockSize = blockSize;

			auto status = nvfbc->NvFBCToSysSetUp(&fbcSysSetupParams);
			if(status != NVFBC_SUCCESS)
				return false;

			fbcSysGrabParams = {};

			// set up grab params
			fbcSysGrabParams.dwVersion = NVFBC_TOSYS_GRAB_FRAME_PARAMS_VER;
			fbcSysGrabParams.dwFlags = NVFBC_TOSYS_WAIT_WITH_TIMEOUT;
			fbcSysGrabParams.dwTargetWidth = 0;
			fbcSysGrabParams.dwTargetHeight = 0;
			fbcSysGrabParams.dwStartX = 0;
			fbcSysGrabParams.dwStartY = 0;
			fbcSysGrabParams.eGMode = NVFBC_TOSYS_SOURCEMODE_FULL;
			fbcSysGrabParams.pNvFBCFrameGrabInfo = &grabInfo;
			fbcSysGrabParams.dwWaitTime = 1000;
			return true;
		}

		DisplayCaptureResult PaintWithTileOptimization(u32* pBufferData) {
			bool copiedTiles = false;

			for(u32 dy = 0; dy < diffmapHeight; ++dy) {
				for(u32 dx = 0; dx < diffmapWidth; ++dx) {
					auto& bl = pDiffMap[dy * diffmapWidth + dx];

					if(bl != 0) {
						copiedTiles = true;

#if 0
							x * (framebuffer.width / diff.diffmapWidth),   // x
							y * (framebuffer.height / diff.diffMapHeight), // y
							framebuffer.width / diff.diffmapWidth,		   // width
							framebuffer.height / diff.diffMapHeight		   // height
#endif

						u32 xOffset = dx * tileWidth;
						u32 yOffset = dy * tileHeight;

						for(u32 y = 0; y < tileHeight; ++y) {
							memcpy(&pBufferData[((y + yOffset) * width + xOffset)],
								   &pRawFramebuffer[((y + yOffset) * grabInfo.dwBufferWidth + xOffset)], tileWidth * 4);
						}

						// printf("copied %ux%u @ %ux%u tile\n", tileWidth, tileHeight, xOffset, yOffset);
					}
				}
			}

			if(!copiedTiles)
				return DisplayCaptureResult::OkUnchanged;

			return DisplayCaptureResult::Ok;
		}

		void PaintFull(u32* pBufferData) {
			for(u32 y = 0; y < height; ++y) {
				// Convert to BGRA
				// FIXME: Make this SIMD. I can't into this very well
#if 0
				usize srcStart = (y * grabInfo.dwBufferWidth) * 4;
				usize dstStart = (y * width) * 4;
				for(u32 x = 0; x < grabInfo.dwWidth * 4; x += 4) {
					pBufferData[(dstStart + x) + 0] = pSrcData[(srcStart + x) + 2]; // B
					pBufferData[(dstStart + x) + 1] = pSrcData[(srcStart + x) + 1]; // G
					pBufferData[(dstStart + x) + 2] = pSrcData[(srcStart + x) + 0]; // R
					pBufferData[(dstStart + x) + 3] = 0xff;							// A
				}
#endif

				memcpy(&pBufferData[(y * width)], &pRawFramebuffer[(y * grabInfo.dwBufferWidth)], width * 4);
			}
		}

		DisplayCaptureResult CaptureFrame(u32* pOut) override {
			auto nvStatus = nvfbc->NvFBCToSysGrabFrame(&this->fbcSysGrabParams);

			switch(nvStatus) {
				case NVFBC_SUCCESS: break;

				// Need to recreate the session. If it fails then we fail too.
				case NVFBC_ERROR_INVALIDATED_SESSION: {
					if(!NvfbcInitSession())
						return DisplayCaptureResult::Fail;

					// Recurse. This looks naughty, but will allow us to directly retry
					// the capture. (Plus if this causes issues whatever has happened is probably beyond saving)
					return CaptureFrame(pOut);
				} break;

				default: return DisplayCaptureResult::Fail;
			}

			auto result = DisplayCaptureResult::Ok;
			bool useTilePainting = true;

			if(width != grabInfo.dwWidth || height != grabInfo.dwHeight) {
				width = grabInfo.dwWidth;
				height = grabInfo.dwHeight;

				// update diffmap stuff
				diffmapWidth = DiffMapDimension(width, blockSize);
				diffmapHeight = DiffMapDimension(height, blockSize);

				tileWidth = width / diffmapWidth;
				tileHeight = height / diffmapHeight;

				result = DisplayCaptureResult::OkButResized;

				// Disable tile-painting optimization for the first frame
				useTilePainting = false;
			}

			if(useTilePainting) [[likely]] {
				result = PaintWithTileOptimization(&pOut[0]);
			} else {
				PaintFull(&pOut[0]);
			}

			return result;
		}

		FramebufferInformation GetFramebufferInformation() override { return FramebufferInformation { width, height }; }

		DiffInformation GetDiffInformation() override {
			// diffmapWidth = (u32)ceil((f32)width / 32);
			// diffmapHeight = (u32)ceil((f32)height / 32);

			return DiffInformation { pDiffMap, diffmapWidth, diffmapHeight };
		}
	};

	DisplayCaptureBase* CreateNVFBCDisplayCapture() {
		return new NVFBCDisplayCapture();
	}
} // namespace hazelnut