#include "StdAfx.h"
#include "cadsdk.h"
#include <winsock2.h>
#include <ws2tcpip.h>
#include <thread>
#include <atomic>
#include <nlohmann/json.hpp>

#pragma comment(lib, "ws2_32.lib")

#define WM_MCP_EXECUTE_TASK (WM_USER + 2048)

static std::thread g_serverThread;
static std::atomic<bool> g_bRunning(false);
static SOCKET g_listenSocket = INVALID_SOCKET;

static HWND g_hMsgWnd = NULL;
static const wchar_t* g_szMsgClassName = L"AutoAgentMcpMsgWndClass";

// 用于跨线程传递数据的结构体
struct TaskData
{
	std::string rawJson;
};

// 此函数在 AutoCAD 主 UI 线程运行，绝对安全.
void ExecuteTaskInMainThread(TaskData *pTask)
{
	if (pTask == nullptr)
	{
		return;
	}

	try
	{
		auto j = nlohmann::json::parse(pTask->rawJson);
		if (j.contains("cmd") && j["cmd"] == "draw_circle")
		{
			double x = j.value("x", 0.0);
			double y = j.value("y", 0.0);
			double r = j.value("r", 100.0);

			AcApDocument *pDoc = acDocManager->curDocument();
			if (pDoc)
			{
				// 此时在主线程，安全锁文档
				Acad::ErrorStatus es = acDocManager->lockDocument(pDoc, AcAp::kWrite, nullptr, nullptr, false);
				if (es == Acad::eOk)
				{
					AcDbDatabase *pDb = pDoc->database();
					AcDbBlockTable *pBlockTable = nullptr;
					pDb->getBlockTable(pBlockTable, AcDb::kForRead);

					AcDbBlockTableRecord *pBTR = nullptr;
					pBlockTable->getAt(ACDB_MODEL_SPACE, pBTR, AcDb::kForWrite);
					pBlockTable->close();

					AcGePoint3d center(x, y, 0.0);
					AcGeVector3d normal(0.0, 0.0, 1.0);
					AcDbCircle *pCircle = new AcDbCircle(center, normal, r);

					AcDbObjectId circleId;
					pBTR->appendAcDbEntity(circleId, pCircle);

					pBTR->close();
					pCircle->close();

					acDocManager->unlockDocument(pDoc);
					acutPrintf(_T("\n[MCP Server] 成功在坐标(%.2f, %.2f)绘制了半径为%.2f的圆\n"), x, y, r);
				}
			}
		}
	}
	catch (...)
	{
		acutPrintf(_T("\n[MCP Server] JSON 解析或绘图异常!\n"));
	}
}

// 隐式窗口的回调函数 (属于主线程)
LRESULT CALLBACK McpWndProc(HWND hWnd, UINT message, WPARAM wParam, LPARAM lParam)
{
	if (message == WM_MCP_EXECUTE_TASK)
	{
		TaskData *pTask = reinterpret_cast<TaskData *>(lParam);
		if (pTask)
		{
			ExecuteTaskInMainThread(pTask);
			delete pTask; // 执行完毕，清理内存
		}
		return 0;
	}
	return DefWindowProc(hWnd, message, wParam, lParam);
}

// TCP 监听后台线程
static void TcpServerLoop()
{
	WSADATA wsaData;
	WSAStartup(MAKEWORD(2, 2), &wsaData);

	g_listenSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	sockaddr_in serverAddr = { 0 };
	serverAddr.sin_family = AF_INET;
	inet_pton(AF_INET, "127.0.0.1", &serverAddr.sin_addr);
	serverAddr.sin_port = htons(8080);

	bind(g_listenSocket, (sockaddr *)&serverAddr, sizeof(serverAddr));
	listen(g_listenSocket, SOMAXCONN);

	while (g_bRunning)
	{
		// 使用 select 模型实现带超时的轮询，确保插件卸载时线程能正常退出
		fd_set readSet;
		FD_ZERO(&readSet);
		FD_SET(g_listenSocket, &readSet);

		timeval timeout = { 1, 0 }; // 1秒超时机制
		int ret = select(0, &readSet, nullptr, nullptr, &timeout);

		if (ret > 0 && FD_ISSET(g_listenSocket, &readSet))
		{
			SOCKET clientSocket = accept(g_listenSocket, nullptr, nullptr);
			if (clientSocket != INVALID_SOCKET)
			{
				char recvBuf[2048] = { 0 };
				int bytesReceived = recv(clientSocket, recvBuf, sizeof(recvBuf) - 1, 0);
				if (bytesReceived > 0)
				{
					// 将数据拷贝到堆内存，准备派发给主线程
					TaskData *pTask = new TaskData;
					pTask->rawJson = recvBuf;

					// 关键转变：通过 PostMessage 发送给主线程拥有的隐式窗口
					if (g_hMsgWnd != NULL)
					{
						PostMessage(g_hMsgWnd, WM_MCP_EXECUTE_TASK, 0, reinterpret_cast<LPARAM>(pTask));
					}
					else
					{
						delete pTask;
					}
				}
				closesocket(clientSocket);
			}
		}
	}

	closesocket(g_listenSocket);
	WSACleanup();
}

namespace mcp
{
	bool init_mcp()
	{
		// 1. 在主线程注册并创建 Message-Only 窗口
		WNDCLASSEX wc = { 0 };
		wc.cbSize = sizeof(WNDCLASSEX);
		wc.lpfnWndProc = McpWndProc;
		wc.hInstance = GetModuleHandle(NULL);
		wc.lpszClassName = g_szMsgClassName;
		RegisterClassEx(&wc);

		g_hMsgWnd = CreateWindowEx(
			0,
			g_szMsgClassName,
			L"AutoAgentMcpMsgWindow",
			0, 0, 0, 0, 0,
			HWND_MESSAGE, // 纯消息窗口，无界面、无感知
			NULL,
			GetModuleHandle(NULL),
			NULL
		);
		g_bRunning = true;
		try
		{
			g_serverThread = std::thread(TcpServerLoop);
			acutPrintf(_T("\n[MCP Server] 后台 TCP 监听线程已启动 (端口: 8080)...\n"));
		}
		catch (const std::system_error &e)
		{
			// 线程创建失败，重置状态标志并输出错误
			g_bRunning = false;
			// e.what() 返回的是 const char*，所以在 acutPrintf 中要用 %hs 而不是 %s 或 %ls
			acutPrintf(_T("\n[MCP Server] 致命错误：无法创建后台监听线程！原因: %hs\n"), e.what());
		}
		return g_bRunning;
	}

	bool close_mcp()
	{
		g_bRunning = false;
		if (g_serverThread.joinable())
		{
			g_serverThread.join();
		}

		if (g_hMsgWnd != NULL)
		{
			DestroyWindow(g_hMsgWnd);
			UnregisterClass(g_szMsgClassName, GetModuleHandle(NULL));
			g_hMsgWnd = NULL;
		}

		acutPrintf(_T("\n[MCP Server] 插件已卸载，线程已关闭。\n"));
		return true;
	}
}