#include "KTCad.h"
#include "mcp_plugin_api.h"
#include <Windows.h>
#include "tchar.h"
#include "rxregsvc.h"
#include <gdiplus.h>
#pragma comment(lib, "gdiplus.lib")

#ifdef ZWCAD
	// 中望的库是 ZwRx.lib，符号也换成了 zcrx 家族（acrxEntryPoint 会被 SDK 宏改名，
	// 但显式写出来的 GetApiVersion 不会，所以这里按平台分开）。
	#pragma comment(linker, "/export:zcrxGetApiVersion,PRIVATE")
#else
	#pragma comment(lib, "rxapi.lib")
	#pragma comment(linker, "/export:acrxGetApiVersion,PRIVATE")
#endif

EXPORT_MCP_PLUGIN_VERSION()

namespace
{
	// 辅助函数：Base64 编码
	std::string Base64Encode(const unsigned char *data, size_t length)
	{
		static const char encoding_table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
		size_t out_len = 4 * ((length + 2) / 3);
		std::string ret(out_len, '\0');
		size_t i;
		char *p = const_cast<char *>(ret.c_str());

		for (i = 0; i < length - 2; i += 3)
		{
			*p++ = encoding_table[(data[i] >> 2) & 0x3F];
			*p++ = encoding_table[((data[i] & 0x3) << 4) | ((int)(data[i + 1] & 0xF0) >> 4)];
			*p++ = encoding_table[((data[i + 1] & 0xF) << 2) | ((int)(data[i + 2] & 0xC0) >> 6)];
			*p++ = encoding_table[data[i + 2] & 0x3F];
		}
		if (i < length)
		{
			*p++ = encoding_table[(data[i] >> 2) & 0x3F];
			if (i == (length - 1))
			{
				*p++ = encoding_table[((data[i] & 0x3) << 4)];
				*p++ = '=';
			}
			else
			{
				*p++ = encoding_table[((data[i] & 0x3) << 4) | ((int)(data[i + 1] & 0xF0) >> 4)];
				*p++ = encoding_table[((data[i + 1] & 0xF) << 2)];
			}
			*p++ = '=';
		}
		return ret;
	}

	// 辅助函数：获取 GDI+ 图像编码器的 Clsid (用于转 PNG)
	int GetEncoderClsid(const WCHAR *format, CLSID *pClsid)
	{
		UINT num = 0, size = 0;
		Gdiplus::GetImageEncodersSize(&num, &size);
		if (size == 0) return -1;
		Gdiplus::ImageCodecInfo *pImageCodecInfo = (Gdiplus::ImageCodecInfo *)(malloc(size));
		if (pImageCodecInfo == NULL) return -1;
		Gdiplus::GetImageEncoders(num, size, pImageCodecInfo);
		for (UINT j = 0; j < num; ++j)
		{
			if (_tcscmp(pImageCodecInfo[j].MimeType, format) == 0)
			{
				*pClsid = pImageCodecInfo[j].Clsid;
				free(pImageCodecInfo);
				return j;
			}
		}
		free(pImageCodecInfo);
		return -1;
	}
}

class mcp_Tool_ViewportScreenshot: public mcp_Tool
{
public:
	virtual nlohmann::json Execute(const nlohmann::json &args) override
	{
		// 1. 获取 AutoCAD 当前文档的绘图区句柄
		HWND hWnd = adsw_acadDocWnd();
		if (!hWnd)
			return mcp_Tool_Utility::make_error(-32000, L"无法获取CAD绘图区窗口句柄");

		// 2. 初始化 GDI+
		Gdiplus::GdiplusStartupInput gdiplusStartupInput;
		ULONG_PTR gdiplusToken;
		Gdiplus::GdiplusStartup(&gdiplusToken, &gdiplusStartupInput, NULL);

		std::string base64Result;
		bool success = false;

		// 3. 开始 Windows GDI 抓屏逻辑 (使用花括号限定生命周期)
		{
			RECT rc;
			GetClientRect(hWnd, &rc);
			int width = rc.right - rc.left;
			int height = rc.bottom - rc.top;

			HDC hdcScreen = GetDC(hWnd);
			HDC hdcMem = CreateCompatibleDC(hdcScreen);
			HBITMAP hBitmap = CreateCompatibleBitmap(hdcScreen, width, height);
			HBITMAP hOldBitmap = (HBITMAP)SelectObject(hdcMem, hBitmap);

			// 将屏幕内容块拷贝到内存位图中
			BitBlt(hdcMem, 0, 0, width, height, hdcScreen, 0, 0, SRCCOPY);

			// 将 GDI 位图转换为 GDI+ 位图
			Gdiplus::Bitmap bmp(hBitmap, NULL);

			// 创建内存流，将位图以 PNG 格式保存到内存，而非硬盘
			IStream *pStream = nullptr;
			if (CreateStreamOnHGlobal(NULL, TRUE, &pStream) == S_OK)
			{
				CLSID pngClsid;
				GetEncoderClsid(L"image/png", &pngClsid);

				if (bmp.Save(pStream, &pngClsid, NULL) == Gdiplus::Ok)
				{
					// 从内存流中读取二进制数据
					STATSTG stat;
					pStream->Stat(&stat, STATFLAG_NONAME);
					ULONG size = stat.cbSize.LowPart;

					std::vector<BYTE> buffer(size);
					LARGE_INTEGER liZero = {};
					pStream->Seek(liZero, STREAM_SEEK_SET, NULL);

					ULONG bytesRead;
					pStream->Read(buffer.data(), size, &bytesRead);

					// 转换为 Base64
					base64Result = Base64Encode(buffer.data(), size);
					success = true;
				}
				pStream->Release();
			}

			// 清理 GDI 资源
			SelectObject(hdcMem, hOldBitmap);
			DeleteObject(hBitmap);
			DeleteDC(hdcMem);
			ReleaseDC(hWnd, hdcScreen);
		}

		// 4. 关闭 GDI+
		Gdiplus::GdiplusShutdown(gdiplusToken);

		// 5. 按照 MCP 多模态协议返回数据
		if (success)
		{
			return mcp_Tool_Utility::make_mixed_result(L"已成功截取 AutoCAD 当前视口画面。", base64Result, "image/png");
		}
		else
		{
			return mcp_Tool_Utility::make_error(-32001, L"截图或图像编码失败");
		}
	}
};

EXPORT_MCP_TOOL(mcp_Tool_ViewportScreenshot, CreateViewportScreenshotTool)