#include "capture.hpp"

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

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

namespace hazelnut {

	class FramebufferCapture_NVFBC final : public IFramebufferCapture {
		NVFBCRESULT nvfbcStatus;
		NvFBCLibrary nvfbcLib;

		NVFBC_TOSYS_GRAB_FRAME_PARAMS fbcSysGrabParams;
		NvFBCFrameGrabInfo grabInfo;
		NvFBCToSys* nvfbc;

		u32 width;
		u32 height;

		u32* pRawFramebuffer;
		unique_buffer<u32> convertedFramebuffer;

		u8* pDiffMap;
		u32 diffmapWidth;
		u32 diffmapHeight;

	   public:
		virtual ~FramebufferCapture_NVFBC() { 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 session to avoid memory leak
			if(nvfbc) {
				NvfbcDestroyInstance();
			}

			DWORD maxDisplayWidth = -1;
			DWORD maxDisplayHeight = -1;
			nvfbc = (NvFBCToSys*)nvfbcLib.create(NVFBC_TO_SYS, &maxDisplayWidth, &maxDisplayHeight);

			if(!nvfbc) {
				return false;
			}

			return true;
		}

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

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

			// enable a 32x32 difference map
			fbcSysSetupParams.bDiffMap = TRUE;
			fbcSysSetupParams.ppDiffMap = (void**)&pDiffMap;
			fbcSysSetupParams.eDiffMapBlockSize = NVFBC_TOSYS_DIFFMAP_BLOCKSIZE_32X32;

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

			memset(&fbcSysGrabParams, 0, sizeof(NVFBC_TOSYS_GRAB_FRAME_PARAMS));

			// 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 = 16;
			return true;
		}

		DisplayCaptureResult CaptureFrame() 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();
				} break;

				default: return DisplayCaptureResult::Fail;
			}

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

				diffmapWidth = (u32)ceil((f32)width / 32);
				diffmapHeight = (u32)ceil((f32)height / 32);

				convertedFramebuffer.resize(width * height);

				return DisplayCaptureResult::OkButResized;
			}

			// Splat to the converted framebuffer. It doesn't have padding on it.
			auto* pBufferData = convertedFramebuffer.data();
			for(u32 y = 0; y < grabInfo.dwHeight; ++y) {
				memcpy(&pBufferData[y * width], &pRawFramebuffer[(y * grabInfo.dwBufferWidth)], grabInfo.dwWidth * 4);
			}

			return DisplayCaptureResult::Ok;
		}

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

		DiffInformation GetDiffInformation() override { return DiffInformation { pDiffMap, diffmapWidth, diffmapHeight }; }
	};

	IFramebufferCapture* CreateFramebufferCapture_NVFBC() {
		return new FramebufferCapture_NVFBC();
	}
} // namespace hazelnut