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

#pragma comment(lib, "ws2_32.lib")
#define WM_MCP_EXECUTE_TASK (WM_USER + 2048)
extern HINSTANCE _hdllInstance;

#define MCP_SERVER_PORT 8080
#define MCP_SERVER_CONFIG _T("mcp_config.json")
#define MCP_SERVER_CONFIG_ALT _T("cad-mcp-frame.json")
#define MCP_MAX_REQUEST_BYTES (4 * 1024 * 1024)

namespace
{
	std::thread g_serverThread;
	std::atomic<bool> g_bRunning(false);
	SOCKET g_listenSocket = INVALID_SOCKET;
	int    g_nPort = MCP_SERVER_PORT;
	std::string g_sToken;          // 非空时要求客户端带 token
	std::wstring g_sConfigPath;    // 实际使用的配置文件全路径

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

	// 声明一个全局工具集对象
	mcp_Tools g_mcpTools;

	// 用于跨线程传递数据的结构体
	struct TaskData
	{
		~TaskData()
		{
			if (clientSocket != INVALID_SOCKET)
			{
				::closesocket(clientSocket);
				clientSocket = INVALID_SOCKET;
			}
		}
		SOCKET clientSocket = INVALID_SOCKET;
		std::string rawJson;
	};

	// ------------------------------------------------------------------
	// 报文边界：客户端发完不等应答就等结果（不做 shutdown），所以服务端只能靠
	// "JSON 结构是否闭合" 判断一个请求是否收完。字符串内的括号不算。
	// ------------------------------------------------------------------
	bool IsJsonComplete(const std::string &s)
	{
		int nBrace = 0;
		int nBracket = 0;
		bool bInStr = false;
		bool bEsc = false;
		bool bAny = false;

		for (size_t i = 0; i < s.size(); ++i)
		{
			const char c = s[i];
			if (bInStr)
			{
				if (bEsc) { bEsc = false; }
				else if (c == '\\') { bEsc = true; }
				else if (c == '"') { bInStr = false; }
				continue;
			}
			switch (c)
			{
			case '"': bInStr = true; bAny = true; break;
			case '{': ++nBrace; bAny = true; break;
			case '}': --nBrace; bAny = true; break;
			case '[': ++nBracket; bAny = true; break;
			case ']': --nBracket; bAny = true; break;
			default: break;
			}
			if (bAny && nBrace == 0 && nBracket == 0 && !bInStr)
				return true;
		}
		return false;
	}

	// 循环接收直到 JSON 闭合（或超时/超长）。
	// 原来只 recv 一次、缓冲固定 2048，大请求会被截断，且 TCP 分包必失败。
	bool RecvOneRequest(SOCKET s, std::string &out)
	{
		out.clear();
		char buf[4096];
		const ULONGLONG nDeadline = ::GetTickCount64() + 30000;   // 整个请求 30s

		for (;;)
		{
			fd_set readSet;
			FD_ZERO(&readSet);
			FD_SET(s, &readSet);
			timeval tv = { 5, 0 };                                  // 单次等待 5s
			const int ret = ::select(0, &readSet, nullptr, nullptr, &tv);
			if (ret < 0)
				return false;
			if (ret == 0)
			{
				if (::GetTickCount64() > nDeadline) return false;
				continue;
			}

			const int nRecv = ::recv(s, buf, (int)sizeof(buf), 0);
			if (nRecv <= 0)
				return IsJsonComplete(out);                         // 对端关闭：已收完整才算成功

			out.append(buf, (size_t)nRecv);
			if (out.size() > MCP_MAX_REQUEST_BYTES)
				return false;
			if (IsJsonComplete(out))
				return true;
			if (::GetTickCount64() > nDeadline)
				return IsJsonComplete(out);
		}
	}

	// send 可能只发一部分，必须循环（tools/list 的应答可能上万字节）
	bool SendAll(SOCKET s, const std::string &data)
	{
		size_t nSent = 0;
		while (nSent < data.size())
		{
			const int n = ::send(s, data.c_str() + nSent, (int)(data.size() - nSent), 0);
			if (n <= 0)
				return false;
			nSent += (size_t)n;
		}
		return true;
	}

	std::string MakeErrorResponse(const char *pszMessage, int nCode = -32603, int nId = 0)
	{
		nlohmann::json err;
		err["jsonrpc"] = "2.0";
		err["id"] = nId;
		err["error"] = { { "code", nCode }, { "message", pszMessage } };
		return err.dump();
	}

	// 处理工具查询请求
	std::string HandleToolsList(int msgId)
	{
		nlohmann::json resp;
		resp["jsonrpc"] = "2.0";
		resp["id"] = msgId;

		// 直接向全局工具集要数据
		resp["result"]["tools"] = g_mcpTools.ToLLM();

		return resp.dump();
	}

	// 处理工具调用请求
	std::string HandleToolsCall(int msgId, const nlohmann::json &params)
	{
		nlohmann::json resp;
		resp["jsonrpc"] = "2.0";
		resp["id"] = msgId;

		std::string toolName = params.value("name", "");
		nlohmann::json args = params.value("arguments", nlohmann::json::object());

		// 让工具集去执行，返回 payload 可能是 {"result": ...} 也可能是 {"error": ...}
		nlohmann::json payload = g_mcpTools.Call(toolName, args);

		// 将 payload 注入到外层信封中
		resp.update(payload);

		return resp.dump();
	}

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

		std::string responseStr;

		// 关闭 nlohmann 的异常抛出机制，安全解析
		nlohmann::json req = nlohmann::json::parse(pTask->rawJson, nullptr, false);

		if (req.is_discarded())
		{
			// 非法 JSON 报文
			responseStr = MakeErrorResponse("Parse error", -32700, 0);
		}
		else
		{
			std::string method = req.value("method", "");
			int msgId = req.value("id", 0);

			// 鉴权：配置里写了 token 就必须带对（请求根级的 token 字段）
			if (!g_sToken.empty())
			{
				const std::string sGot = req.value("token", "");
				if (sGot != g_sToken)
				{
					acutPrintf(_T("\n[MCP Server] 请求被拒：token 缺失或不匹配（客户端 IP 与本机相同路径也要带 token）。"));
					responseStr = MakeErrorResponse("Unauthorized", -32001, msgId);
					if (pTask->clientSocket != INVALID_SOCKET)
						SendAll(pTask->clientSocket, responseStr);
					return;
				}
			}

			if (method == "tools/list")
			{
				responseStr = HandleToolsList(msgId);
			}
			else if (method == "tools/call")
			{
				responseStr = HandleToolsCall(msgId, req.value("params", nlohmann::json::object()));
			}
			else
			{
				// 找不到方法
				responseStr = MakeErrorResponse("Method not found", -32601, msgId);
			}
		}

		// 通过传递过来的 Socket 将结果发送回大模型端
		if (pTask->clientSocket != INVALID_SOCKET)
		{
			SendAll(pTask->clientSocket, responseStr);
		}
	}

	// 隐式窗口的回调函数 (属于主线程)
	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 监听后台线程（监听 socket 由主线程创建好并传入，失败在启动时就暴露）
	void TcpServerLoop(SOCKET listenSocket)
	{
		while (g_bRunning)
		{
			fd_set readSet;
			FD_ZERO(&readSet);
			FD_SET(listenSocket, &readSet);

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

			if (ret > 0 && FD_ISSET(listenSocket, &readSet))
			{
				SOCKET clientSocket = accept(listenSocket, nullptr, nullptr);
				if (clientSocket == INVALID_SOCKET)
					continue;

				// 循环收直到 JSON 闭合；原来只 recv 一次 2048，大请求会被截断
				std::string rawJson;
				const bool bOk = RecvOneRequest(clientSocket, rawJson);
				if (!bOk || rawJson.empty())
				{
					// 收包失败/超时：回一个错误并关闭，避免客户端干等
					SendAll(clientSocket, MakeErrorResponse("Bad request", -32700));
					::closesocket(clientSocket);   // 原来这里没关，会泄漏 socket
					continue;
				}

				// 将数据拷贝到堆内存，准备派发给主线程
				TaskData *pTask = new TaskData;
				pTask->clientSocket = clientSocket; // 记录 Socket，由 TaskData 析构关闭
				pTask->rawJson = rawJson;

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

		closesocket(listenSocket);
		WSACleanup();
	}

}

namespace mcp
{
	// 兼容两种配置文件名：先找 mcp_config.json，再找 cad-mcp-frame.json
	static std::wstring ResolveConfigPath()
	{
		const fs::path dir = mcp_utils::getModelDirPath();
		fs::path p1 = dir / MCP_SERVER_CONFIG;
		if (fs::exists(p1)) return p1.wstring();
		fs::path p2 = dir / MCP_SERVER_CONFIG_ALT;
		if (fs::exists(p2)) return p2.wstring();
		return p1.wstring();   // 都不存在：返回默认名，后续加载会失败并报错
	}

	// 读取 server.port / server.token（可选）
	static void LoadServerSettings(const std::wstring &wsConfig)
	{
		try
		{
			std::ifstream ifs(wsConfig);
			if (!ifs.is_open()) return;
			nlohmann::json cfg;
			ifs >> cfg;
			if (!cfg.contains("server")) return;
			const nlohmann::json &srv = cfg["server"];
			const int nPort = srv.value("port", MCP_SERVER_PORT);
			if (nPort > 0 && nPort < 65536) g_nPort = nPort;
			g_sToken = srv.value("token", std::string());
		}
		catch (...)
		{
			// 配置读坏了不影响默认值
		}
	}

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

		if (g_hMsgWnd != NULL)
		{
			// 排空仍挂在本窗口消息队列里的任务，否则 TaskData（含 socket）会泄漏
			MSG msg;
			int nGuard = 0;
			while (nGuard++ < 1000 && PeekMessage(&msg, g_hMsgWnd, WM_MCP_EXECUTE_TASK, WM_MCP_EXECUTE_TASK, PM_REMOVE))
			{
				TaskData *pTask = reinterpret_cast<TaskData *>(msg.lParam);
				if (pTask != NULL) delete pTask;
			}

			DestroyWindow(g_hMsgWnd);
			UnregisterClass(g_szMsgClassName, _hdllInstance);
			g_hMsgWnd = NULL;
		}

		g_mcpTools.clear();

		acutPrintf(_T("\n[MCP Server] The plugin has been uninstalled, and the thread has been closed.\n"));
		return true;
	}

	bool init_mcp_server()
	{
		// 0. 配置文件与服务器设置（端口/可选 token）
		g_sConfigPath = ResolveConfigPath();
		LoadServerSettings(g_sConfigPath);

		// 1. 先把监听 socket 建好：失败就在启动时暴露，不再静默
		WSADATA wsaData;
		if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
		{
			acutPrintf(_T("\n[MCP Server] Fatal: WSAStartup failed.\n"));
			return false;
		}

		SOCKET listenSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
		if (listenSocket == INVALID_SOCKET)
		{
			acutPrintf(_T("\n[MCP Server] Fatal: socket() failed, err=%d.\n"), ::WSAGetLastError());
			WSACleanup();
			return false;
		}

		sockaddr_in serverAddr = { 0 };
		serverAddr.sin_family = AF_INET;
		serverAddr.sin_port = htons((u_short)g_nPort);
		inet_pton(AF_INET, "127.0.0.1", &serverAddr.sin_addr);

		if (bind(listenSocket, (sockaddr *)&serverAddr, sizeof(serverAddr)) != 0)
		{
			const int nErr = ::WSAGetLastError();
			acutPrintf(_T("\n[MCP Server] Fatal: bind 127.0.0.1:%d failed, err=%d（端口被占用？）\n"), g_nPort, nErr);
			closesocket(listenSocket);
			WSACleanup();
			return false;
		}

		if (listen(listenSocket, SOMAXCONN) != 0)
		{
			acutPrintf(_T("\n[MCP Server] Fatal: listen failed, err=%d.\n"), ::WSAGetLastError());
			closesocket(listenSocket);
			WSACleanup();
			return false;
		}

		g_listenSocket = listenSocket;

		// 2. 在主线程注册并创建 Message-Only 窗口
		WNDCLASSEX wc = { 0 };
		wc.cbSize = sizeof(WNDCLASSEX);
		wc.lpfnWndProc = McpWndProc;
		wc.hInstance = _hdllInstance;
		wc.lpszClassName = g_szMsgClassName;
		RegisterClassEx(&wc);

		g_hMsgWnd = CreateWindowEx(
			0,
			g_szMsgClassName,
			L"AutoAgentMcpMsgWindow",
			0, 0, 0, 0, 0,
			HWND_MESSAGE, // 纯消息窗口，无界面、无感知
			NULL,
			_hdllInstance,
			NULL
		);
		if (g_hMsgWnd == NULL)
		{
			acutPrintf(_T("\n[MCP Server] Fatal: CreateWindowEx failed, err=%d.\n"), ::GetLastError());
			closesocket(listenSocket);
			g_listenSocket = INVALID_SOCKET;
			UnregisterClass(g_szMsgClassName, _hdllInstance);
			WSACleanup();
			return false;
		}

		g_bRunning = true;
		try
		{
			g_serverThread = std::thread(TcpServerLoop, listenSocket);
			acutPrintf(_T("\n[MCP Server] Background TCP listening thread started (Port: %d, token: %s)...\n"),
				g_nPort, g_sToken.empty() ? _T("none") : _T("required"));
		}
		catch (const std::system_error &e)
		{
			g_bRunning = false;
			acutPrintf(_T("\n[MCP Server] Fatal error: Unable to create background listening thread! Reason: %hs\n"), e.what());
			DestroyWindow(g_hMsgWnd);
			g_hMsgWnd = NULL;
			UnregisterClass(g_szMsgClassName, _hdllInstance);
			closesocket(listenSocket);
			g_listenSocket = INVALID_SOCKET;
			WSACleanup();
			return false;
		}

		// 3. 监听跑起来了才注册工具
		if (!mcp_ConfigParser::LoadAndRegister(g_sConfigPath, g_mcpTools))
		{
			acutPrintf(_T("\n[MCP Server] Warning: Failed to load tool configuration!")
				_T("\n[MCP Server] 只支持 mcp_config.json 或 cad-mcp-frame.json，且必须是合法 JSON。Automatic Shutdown.\n"));
			close_mcp_server();
			return false;
		}

		g_mcpTools.print();
		return g_bRunning;
	}
}