#include "stdafx.h"
#include "McpServer.h"
#include "jsoncpp/json.h"
#include "HostJson.h"        // BuildHostJson（单台主机序列化公共函数）
#include "context.h"         // context 接口
#include "2015RemoteDlg.h"   // CMy2015RemoteDlg 成员（m_HostList/m_cs/m_ClientMap）+ VERSION_STR

#include <sstream>

#ifndef _WIN64
#ifdef _DEBUG
#pragma comment(lib, "jsoncpp/jsoncppd.lib")
#else
#pragma comment(lib, "jsoncpp/jsoncpp.lib")
#endif
#else
#ifdef _DEBUG
#pragma comment(lib, "jsoncpp/jsoncpp_x64d.lib")
#else
#pragma comment(lib, "jsoncpp/jsoncpp_x64.lib")
#endif
#endif

namespace {

// P2b 工具等待响应的超时（ms）。MCP 一次性请求：等待子连接回传进程/窗口列表。
static const int kMcpToolTimeoutMs = 20000;

// Json::Value → 紧凑 JSON 字符串
std::string JsonToString(const Json::Value& v) {
    Json::StreamWriterBuilder b;
    b["indentation"] = "";
    return Json::writeString(b, v);
}

// JSON-RPC 2.0 成功响应
std::string BuildResult(const Json::Value& id, const Json::Value& result) {
    Json::Value resp(Json::objectValue);
    resp["jsonrpc"] = "2.0";
    resp["id"] = id;
    resp["result"] = result;
    return JsonToString(resp);
}

// JSON-RPC 2.0 错误响应
std::string BuildError(const Json::Value& id, int code, const std::string& msg) {
    Json::Value resp(Json::objectValue);
    resp["jsonrpc"] = "2.0";
    resp["id"] = id;
    Json::Value err(Json::objectValue);
    err["code"] = code;
    err["message"] = msg;
    resp["error"] = err;
    return JsonToString(resp);
}

// initialize 握手（MCP 规范：protocolVersion + capabilities + serverInfo）
std::string BuildInitializeResult(const Json::Value& id) {
    Json::Value result(Json::objectValue);
    result["protocolVersion"] = "2025-06-18";
    Json::Value caps(Json::objectValue);
    caps["tools"] = Json::Value(Json::objectValue);
    result["capabilities"] = caps;
    Json::Value serverInfo(Json::objectValue);
    serverInfo["name"] = "yama";
    serverInfo["version"] = VERSION_STR;
    result["serverInfo"] = serverInfo;
    return BuildResult(id, result);
}

// ping 健康检查：返回空 result
std::string BuildPingResult(const Json::Value& id) {
    return BuildResult(id, Json::Value(Json::objectValue));
}

// ========== P2a 通用辅助 ==========

// 小写化（仅 ASCII，UTF-8 多字节原样保留）：用于不区分大小写的子串匹配
std::string ToLowerAscii(const std::string& s) {
    std::string r = s;
    for (char& c : r) if (c >= 'A' && c <= 'Z') c = (char)(c - 'A' + 'a');
    return r;
}

// 不区分大小写的子串匹配
bool ContainsCI(const std::string& haystack, const std::string& needle) {
    if (needle.empty()) return true;
    return ToLowerAscii(haystack).find(ToLowerAscii(needle)) != std::string::npos;
}

// 取对象的字符串字段，缺失/非字符串返回 ""
std::string JsonStrField(const Json::Value& v, const char* key) {
    if (v.isObject() && v.isMember(key) && v[key].isString())
        return v[key].asString();
    return "";
}

// 是否纯数字（host id 为 uint64 十进制字符串）
bool IsDigits(const std::string& s) {
    if (s.empty()) return false;
    for (char c : s) if (c < '0' || c > '9') return false;
    return true;
}

// 读取 tools/call 的入参（MCP 规范：params.arguments 为工具入参对象）
Json::Value GetCallArguments(const Json::Value& params) {
    if (params.isObject() && params.isMember("arguments") && params["arguments"].isObject())
        return params["arguments"];
    return Json::Value(Json::objectValue);
}

// 读取可选字符串入参，缺失返回 ""
std::string GetStringArg(const Json::Value& args, const char* key) {
    return JsonStrField(args, key);
}

// ========== P2b 辅助 ==========

// 解析 id 入参（必填、纯数字）为 uint64；非法返回 false。
bool ParseHostIdArg(const Json::Value& args, uint64_t& out, std::string& err) {
    std::string sid = GetStringArg(args, "id");
    if (sid.empty()) {
        err = "Missing required parameter: id";
        return false;
    }
    if (!IsDigits(sid)) {
        err = "Invalid id: expected a decimal host id string";
        return false;
    }
    out = strtoull(sid.c_str(), nullptr, 10);
    return true;
}

// 在 m_HostList 中按 clientID 找在线主 context（复刻 CollectOnlineHosts 的锁内遍历）。
context* FindMainContext(CMy2015RemoteDlg* parent, uint64_t id) {
    if (!parent) return nullptr;
    context* found = nullptr;
    EnterCriticalSection(&parent->m_cs);
    for (context* ctx : parent->m_HostList) {
        if (ctx && ctx->GetClientID() == id && ctx->IsLogin()) {
            found = ctx;
            break;
        }
    }
    LeaveCriticalSection(&parent->m_cs);
    return found;
}

// 有界 strlen：p 最多可读 avail 字节，返回 n < avail 表示遇到 '\0'。
size_t BoundedStrlen(const char* p, size_t avail) {
    size_t n = 0;
    while (n < avail && p[n]) ++n;
    return n;
}

// 从源编码 cp 转为 UTF-8（输出到 JSON）。空/失败返回 ""。
std::string ToUtf8(const char* s, UINT cp) {
    if (!s || !*s) return "";
    int wlen = MultiByteToWideChar(cp, 0, s, -1, NULL, 0);
    if (wlen <= 0) return "";
    std::wstring w(wlen - 1, L'\0');
    MultiByteToWideChar(cp, 0, s, -1, &w[0], wlen);
    int u8len = WideCharToMultiByte(CP_UTF8, 0, w.c_str(), -1, NULL, 0, NULL, NULL);
    if (u8len <= 0) return "";
    std::string out(u8len - 1, '\0');
    WideCharToMultiByte(CP_UTF8, 0, w.c_str(), -1, &out[0], u8len, NULL, NULL);
    return out;
}

// 从 "title|status|pid|r1|r2" 解析（自末尾 4 个 '|' 反推），兼容只有标题的老客户端。
// 逻辑复刻 SystemDlg.cpp::ParseWindowAttrs。
void ParseWindowAttrsStr(const char* s, std::string& title, std::string& status, uint64_t& pid) {
    title.clear();
    status = "normal";
    pid = 0;
    if (!s || !*s) return;
    std::string t(s);
    int len = (int)t.size();
    int pipePos[4] = { -1, -1, -1, -1 };
    int cnt = 0;
    for (int i = len - 1; i >= 0 && cnt < 4; --i) {
        if (t[i] == '|') pipePos[cnt++] = i;
    }
    if (cnt < 4) {  // 老格式：只有标题
        title = t;
        return;
    }
    pid = strtoull(t.c_str() + pipePos[2] + 1, nullptr, 10);
    status = t.substr(pipePos[3] + 1, pipePos[2] - pipePos[3] - 1);
    title = t.substr(0, pipePos[3]);
}

// 解析 TOKEN_PSLIST 缓冲（data[0]=token，其后为 [pid:4][name:arch\0][path\0] 记录）。
// 进程名/路径来自客户端 ANSI（CP_ACP），转 UTF-8 输出。
Json::Value ParseProcessList(const std::vector<BYTE>& data) {
    Json::Value arr(Json::arrayValue);
    if (data.size() < 2) return arr;
    const char* p = (const char*)data.data();
    size_t len = data.size();
    size_t off = 1;  // 跳过 TOKEN 字节
    while (off + sizeof(DWORD) <= len) {
        DWORD pid = *(const DWORD*)(p + off);
        off += sizeof(DWORD);

        const char* exeFile = p + off;
        size_t exeLen = BoundedStrlen(exeFile, len - off);
        if (exeLen >= len - off) break;  // 未以 '\0' 结尾，异常数据
        if (exeLen == 0) break;          // 空进程名 = 尾部零填充（LocalSize 对齐），停止解析
        off += exeLen + 1;

        const char* fullPath = p + off;
        size_t pathLen = BoundedStrlen(fullPath, len - off);
        if (pathLen >= len - off) break;
        off += pathLen + 1;

        std::string name(exeFile, exeLen);
        std::string arch;
        size_t colon = name.find(':');
        if (colon != std::string::npos) {
            arch = name.substr(colon + 1);
            name = name.substr(0, colon);
        }

        Json::Value item(Json::objectValue);
        item["pid"]  = (Json::UInt64)pid;
        item["name"] = ToUtf8(name.c_str(), CP_ACP);
        item["arch"] = arch.empty() ? "N/A" : ToUtf8(arch.c_str(), CP_ACP);
        item["path"] = ToUtf8(fullPath, CP_ACP);
        arr.append(item);
    }
    return arr;
}

// 解析 TOKEN_WSLIST 缓冲（data[0]=token，其后为 [hwnd:4][title|status|pid|r1|r2\0] 记录）。
// 窗口标题为客户端 UTF-8（老客户端为 CP_ACP），用 cp 解码后转 UTF-8 输出。
Json::Value ParseWindowList(const std::vector<BYTE>& data, UINT cp) {
    Json::Value arr(Json::arrayValue);
    if (data.size() < 2) return arr;
    const char* p = (const char*)data.data();
    size_t len = data.size();
    size_t off = 1;
    while (off + sizeof(DWORD) <= len) {
        DWORD hwnd = *(const DWORD*)(p + off);
        off += sizeof(DWORD);

        const char* titleWithAttrs = p + off;
        size_t tlen = BoundedStrlen(titleWithAttrs, len - off);
        if (tlen >= len - off) break;
        if (tlen == 0) break;           // 空记录 = 尾部零填充（LocalSize 对齐），停止解析
        off += tlen + 1;

        std::string title, status;
        uint64_t pid = 0;
        ParseWindowAttrsStr(titleWithAttrs, title, status, pid);

        Json::Value item(Json::objectValue);
        item["hwnd"]   = (Json::UInt64)hwnd;
        item["title"]  = ToUtf8(title.c_str(), cp);
        item["status"] = status;
        item["pid"]    = (Json::UInt64)pid;
        arr.append(item);
    }
    return arr;
}

// 收集所有在线主机 JSON 数组（m_cs 锁内遍历，复用 BuildHostJson 序列化，方案 C）
void CollectOnlineHosts(CMy2015RemoteDlg* parent, Json::Value& hosts) {
    if (!parent) return;
    EnterCriticalSection(&parent->m_cs);
    for (context* ctx : parent->m_HostList) {
        if (!ctx || !ctx->IsLogin()) continue;
        hosts.append(BuildHostJson(ctx, parent->m_ClientMap));
    }
    LeaveCriticalSection(&parent->m_cs);
}

// ========== 工具 schema ==========

// 单台主机字段 schema（hosts 数组元素 / 单机详情共用的形状）
Json::Value BuildHostItemSchema() {
    Json::Value itemProps(Json::objectValue);
    const char* strFields[] = {
        "id", "name", "remark", "ip", "os", "location", "rtt",
        "version", "activeWindow", "group", "screen", "clientType"
    };
    for (const char* f : strFields) {
        Json::Value p(Json::objectValue);
        p["type"] = "string";
        itemProps[f] = p;
    }
    Json::Value onlineProp(Json::objectValue);
    onlineProp["type"] = "boolean";
    itemProps["online"] = onlineProp;
    return itemProps;
}

// list_online_hosts / search_hosts 的 outputSchema（hosts 数组）
Json::Value BuildHostOutputSchema() {
    Json::Value props(Json::objectValue);

    Json::Value hostsProp(Json::objectValue);
    hostsProp["type"] = "array";
    Json::Value items(Json::objectValue);
    items["type"] = "object";
    items["properties"] = BuildHostItemSchema();
    hostsProp["items"] = items;
    props["hosts"] = hostsProp;

    Json::Value schema(Json::objectValue);
    schema["type"] = "object";
    schema["properties"] = props;
    Json::Value required(Json::arrayValue);
    required.append("hosts");
    schema["required"] = required;
    return schema;
}

// get_host_detail 的 outputSchema（单台主机）
Json::Value BuildHostDetailOutputSchema() {
    Json::Value props(Json::objectValue);

    Json::Value hostProp(Json::objectValue);
    hostProp["type"] = "object";
    hostProp["properties"] = BuildHostItemSchema();
    props["host"] = hostProp;

    Json::Value schema(Json::objectValue);
    schema["type"] = "object";
    schema["properties"] = props;
    Json::Value required(Json::arrayValue);
    required.append("host");
    schema["required"] = required;
    return schema;
}

// search_hosts 的 inputSchema（全部可选）
Json::Value BuildSearchHostsInputSchema() {
    Json::Value props(Json::objectValue);
    const char* strParams[] = { "name", "ip", "group", "os" };
    for (const char* p : strParams) {
        Json::Value s(Json::objectValue);
        s["type"] = "string";
        props[p] = s;
    }
    Json::Value onlineProp(Json::objectValue);
    onlineProp["type"] = "boolean";
    props["online"] = onlineProp;

    Json::Value schema(Json::objectValue);
    schema["type"] = "object";
    schema["properties"] = props;
    return schema;
}

// get_host_detail 的 inputSchema（id 必填）
Json::Value BuildGetHostDetailInputSchema() {
    Json::Value props(Json::objectValue);

    Json::Value idProp(Json::objectValue);
    idProp["type"] = "string";
    idProp["description"] = u8"主机 id，取 list_online_hosts / search_hosts 返回的 id 字段";
    props["id"] = idProp;

    Json::Value schema(Json::objectValue);
    schema["type"] = "object";
    schema["properties"] = props;
    Json::Value required(Json::arrayValue);
    required.append("id");
    schema["required"] = required;
    return schema;
}

// list_processes 的 outputSchema（processes 数组）
Json::Value BuildProcessListOutputSchema() {
    Json::Value props(Json::objectValue);

    Json::Value procsProp(Json::objectValue);
    procsProp["type"] = "array";
    Json::Value items(Json::objectValue);
    items["type"] = "object";
    Json::Value itemProps(Json::objectValue);
    Json::Value pidProp(Json::objectValue);
    pidProp["type"] = "integer";
    itemProps["pid"] = pidProp;
    const char* strFields[] = { "name", "arch", "path" };
    for (const char* f : strFields) {
        Json::Value s(Json::objectValue);
        s["type"] = "string";
        itemProps[f] = s;
    }
    items["properties"] = itemProps;
    procsProp["items"] = items;
    props["processes"] = procsProp;

    Json::Value schema(Json::objectValue);
    schema["type"] = "object";
    schema["properties"] = props;
    Json::Value required(Json::arrayValue);
    required.append("processes");
    schema["required"] = required;
    return schema;
}

// list_windows 的 outputSchema（windows 数组）
Json::Value BuildWindowListOutputSchema() {
    Json::Value props(Json::objectValue);

    Json::Value winsProp(Json::objectValue);
    winsProp["type"] = "array";
    Json::Value items(Json::objectValue);
    items["type"] = "object";
    Json::Value itemProps(Json::objectValue);
    Json::Value hwndProp(Json::objectValue);
    hwndProp["type"] = "integer";
    itemProps["hwnd"] = hwndProp;
    Json::Value pidProp(Json::objectValue);
    pidProp["type"] = "integer";
    itemProps["pid"] = pidProp;
    const char* strFields[] = { "title", "status" };
    for (const char* f : strFields) {
        Json::Value s(Json::objectValue);
        s["type"] = "string";
        itemProps[f] = s;
    }
    items["properties"] = itemProps;
    winsProp["items"] = items;
    props["windows"] = winsProp;

    Json::Value schema(Json::objectValue);
    schema["type"] = "object";
    schema["properties"] = props;
    Json::Value required(Json::arrayValue);
    required.append("windows");
    schema["required"] = required;
    return schema;
}

// get_activity_history 的 outputSchema（records 数组 + 原始文本）
Json::Value BuildActivityHistoryOutputSchema() {
    Json::Value props(Json::objectValue);

    Json::Value recordsProp(Json::objectValue);
    recordsProp["type"] = "array";
    Json::Value items(Json::objectValue);
    items["type"] = "string";
    recordsProp["items"] = items;
    props["records"] = recordsProp;

    Json::Value rawProp(Json::objectValue);
    rawProp["type"] = "string";
    props["activityHistory"] = rawProp;

    Json::Value schema(Json::objectValue);
    schema["type"] = "object";
    schema["properties"] = props;
    return schema;
}

// tools/list
std::string BuildToolsListResult(const Json::Value& id) {
    Json::Value result(Json::objectValue);
    Json::Value tools(Json::arrayValue);

    // 1) list_online_hosts
    {
        Json::Value tool(Json::objectValue);
        tool["name"] = "list_online_hosts";
        // 说明文字为 UTF-8：项目 /execution-charset:.936 会把普通窄字面量编译成 GBK，
        // 故用 u8 前缀确保输出到 JSON 的字节是 UTF-8。
        tool["description"] = u8"获取当前所有在线主机的列表，包含计算机名、IP、操作系统、版本、备注、分组、活动窗口、延迟等实时信息。";

        Json::Value inputSchema(Json::objectValue);
        inputSchema["type"] = "object";
        inputSchema["properties"] = Json::Value(Json::objectValue);
        inputSchema["required"] = Json::Value(Json::arrayValue);
        tool["inputSchema"] = inputSchema;

        tool["outputSchema"] = BuildHostOutputSchema();

        tools.append(tool);
    }

    // 2) search_hosts（P2a：纯内存过滤，无子链接）
    {
        Json::Value tool(Json::objectValue);
        tool["name"] = "search_hosts";
        tool["description"] = u8"按计算机名/备注、IP、分组、操作系统过滤在线主机。所有条件均可选、按 AND 组合；子串匹配（ASCII 不区分大小写）。只返回在线主机。";

        tool["inputSchema"] = BuildSearchHostsInputSchema();
        tool["outputSchema"] = BuildHostOutputSchema();

        tools.append(tool);
    }

    // 3) get_host_detail（P2a：单机详情，纯内存）
    {
        Json::Value tool(Json::objectValue);
        tool["name"] = "get_host_detail";
        tool["description"] = u8"获取单台在线主机的详细信息（id、计算机名、IP、操作系统、备注、分组、活动窗口、屏幕分辨率、客户端类型等）。";

        tool["inputSchema"] = BuildGetHostDetailInputSchema();
        tool["outputSchema"] = BuildHostDetailOutputSchema();

        tools.append(tool);
    }

    // 4) list_processes（P2b：主连接 RPC，进程列表）
    {
        Json::Value tool(Json::objectValue);
        tool["name"] = "list_processes";
        tool["description"] = u8"获取指定在线主机的进程列表（PID、映像名称、架构、程序完整路径）。通过主连接下发命令、子连接回传，一次性返回。";

        tool["inputSchema"] = BuildGetHostDetailInputSchema();  // 复用 { id } 必填 schema
        tool["outputSchema"] = BuildProcessListOutputSchema();

        tools.append(tool);
    }

    // 5) list_windows（P2b：主连接 RPC，窗口列表）
    {
        Json::Value tool(Json::objectValue);
        tool["name"] = "list_windows";
        tool["description"] = u8"获取指定在线主机的顶层窗口列表（句柄、窗口标题、窗口状态、所属进程 PID）。一次性返回。";

        tool["inputSchema"] = BuildGetHostDetailInputSchema();  // 复用 { id } 必填 schema
        tool["outputSchema"] = BuildWindowListOutputSchema();

        tools.append(tool);
    }

    // 6) get_activity_history（P2b：主连接 RPC，历史活动记录）
    {
        Json::Value tool(Json::objectValue);
        tool["name"] = "get_activity_history";
        tool["description"] = u8"获取指定在线主机的历史活动记录（前台窗口驻留时长，每行一条「[时间] [标题] 时长」）。主连接 RPC、一次性返回。";

        tool["inputSchema"] = BuildGetHostDetailInputSchema();  // 复用 { id } 必填 schema
        tool["outputSchema"] = BuildActivityHistoryOutputSchema();

        tools.append(tool);
    }

    result["tools"] = tools;
    return BuildResult(id, result);
}

// tools/call：list_online_hosts
std::string BuildListOnlineHosts(const Json::Value& id, CMy2015RemoteDlg* parent) {
    Json::Value hosts(Json::arrayValue);
    CollectOnlineHosts(parent, hosts);
    int count = (int)hosts.size();

    Json::Value result(Json::objectValue);
    Json::Value structuredContent(Json::objectValue);
    structuredContent["hosts"] = hosts;
    result["structuredContent"] = structuredContent;

    Json::Value content(Json::arrayValue);
    Json::Value item(Json::objectValue);
    item["type"] = "text";
    item["text"] = std::string(u8"共 ") + std::to_string(count) + std::string(u8" 台主机在线。");
    content.append(item);
    result["content"] = content;
    result["isError"] = false;

    return BuildResult(id, result);
}

// tools/call：search_hosts
std::string BuildSearchHosts(const Json::Value& id, const Json::Value& args, CMy2015RemoteDlg* parent) {
    Json::Value all(Json::arrayValue);
    CollectOnlineHosts(parent, all);

    std::string fName  = GetStringArg(args, "name");
    std::string fIp    = GetStringArg(args, "ip");
    std::string fGroup = GetStringArg(args, "group");
    std::string fOs    = GetStringArg(args, "os");
    bool hasOnline = args.isObject() && args.isMember("online") && args["online"].isBool();
    bool wantOnline = hasOnline ? args["online"].asBool() : true;

    Json::Value hosts(Json::arrayValue);
    // 列表只含在线主机：显式 online=false 时直接空结果
    if (!hasOnline || wantOnline) {
        for (unsigned int i = 0; i < all.size(); ++i) {
            const Json::Value& h = all[i];
            if (!fName.empty()) {
                std::string name   = JsonStrField(h, "name");
                std::string remark = JsonStrField(h, "remark");
                if (!ContainsCI(name, fName) && !ContainsCI(remark, fName)) continue;
            }
            if (!fIp.empty()    && !ContainsCI(JsonStrField(h, "ip"),    fIp))    continue;
            if (!fGroup.empty() && !ContainsCI(JsonStrField(h, "group"), fGroup)) continue;
            if (!fOs.empty()    && !ContainsCI(JsonStrField(h, "os"),    fOs))    continue;
            hosts.append(h);
        }
    }
    int count = (int)hosts.size();

    Json::Value result(Json::objectValue);
    Json::Value structuredContent(Json::objectValue);
    structuredContent["hosts"] = hosts;
    result["structuredContent"] = structuredContent;

    Json::Value content(Json::arrayValue);
    Json::Value item(Json::objectValue);
    item["type"] = "text";
    item["text"] = std::string(u8"共 ") + std::to_string(count) + std::string(u8" 台主机匹配。");
    content.append(item);
    result["content"] = content;
    result["isError"] = false;

    return BuildResult(id, result);
}

// tools/call：get_host_detail
std::string BuildGetHostDetail(const Json::Value& id, const Json::Value& args, CMy2015RemoteDlg* parent) {
    std::string sid = GetStringArg(args, "id");
    if (sid.empty()) {
        return BuildError(id, -32602, "Missing required parameter: id");
    }
    if (!IsDigits(sid)) {
        return BuildError(id, -32602, "Invalid id: expected a decimal host id string");
    }

    Json::Value all(Json::arrayValue);
    CollectOnlineHosts(parent, all);
    for (unsigned int i = 0; i < all.size(); ++i) {
        const Json::Value& h = all[i];
        if (JsonStrField(h, "id") == sid) {
            Json::Value result(Json::objectValue);
            Json::Value structuredContent(Json::objectValue);
            structuredContent["host"] = h;
            result["structuredContent"] = structuredContent;

            Json::Value content(Json::arrayValue);
            Json::Value item(Json::objectValue);
            item["type"] = "text";
            std::string name = JsonStrField(h, "name");
            std::string ip   = JsonStrField(h, "ip");
            item["text"] = std::string(u8"主机 ") + name + " (" + ip + ")" + u8" 的详情。";
            content.append(item);
            result["content"] = content;
            result["isError"] = false;

            return BuildResult(id, result);
        }
    }

    return BuildError(id, -32002, "Host not found or offline: " + sid);
}

// tools/call：list_processes（主连接 RPC，子连接一次性回传）
std::string BuildListProcesses(const Json::Value& id, const Json::Value& args, CMy2015RemoteDlg* parent) {
    uint64_t devId = 0;
    std::string err;
    if (!ParseHostIdArg(args, devId, err))
        return BuildError(id, -32602, err);

    context* ctx = FindMainContext(parent, devId);
    if (!ctx)
        return BuildError(id, -32002, "Host not found or offline: " + std::to_string(devId));

    CMcpServer& mcp = CMcpServer::Instance();
    if (!mcp.BeginPending(devId, "list_processes"))
        return BuildError(id, -32003, "Device busy: another request is pending for this host");

    BYTE cmd = COMMAND_SYSTEM;
    if (!ctx->Send2Client(&cmd, 1)) {
        mcp.ClearPending(devId);
        return BuildError(id, -32004, "Failed to send command to host");
    }

    std::vector<BYTE> data;
    if (!mcp.WaitPending(devId, data, kMcpToolTimeoutMs))
        return BuildError(id, -32001, "Timeout waiting for process list");

    Json::Value procs = ParseProcessList(data);
    int count = (int)procs.size();

    Json::Value result(Json::objectValue);
    Json::Value structuredContent(Json::objectValue);
    structuredContent["processes"] = procs;
    result["structuredContent"] = structuredContent;

    Json::Value content(Json::arrayValue);
    Json::Value item(Json::objectValue);
    item["type"] = "text";
    item["text"] = std::string(u8"共 ") + std::to_string(count) + std::string(u8" 个进程。");
    content.append(item);
    result["content"] = content;
    result["isError"] = false;

    return BuildResult(id, result);
}

// tools/call：list_windows（主连接 RPC，子连接一次性回传）
std::string BuildListWindows(const Json::Value& id, const Json::Value& args, CMy2015RemoteDlg* parent) {
    uint64_t devId = 0;
    std::string err;
    if (!ParseHostIdArg(args, devId, err))
        return BuildError(id, -32602, err);

    context* ctx = FindMainContext(parent, devId);
    if (!ctx)
        return BuildError(id, -32002, "Host not found or offline: " + std::to_string(devId));

    CMcpServer& mcp = CMcpServer::Instance();
    if (!mcp.BeginPending(devId, "list_windows"))
        return BuildError(id, -32003, "Device busy: another request is pending for this host");

    BYTE cmd = COMMAND_WSLIST;
    if (!ctx->Send2Client(&cmd, 1)) {
        mcp.ClearPending(devId);
        return BuildError(id, -32004, "Failed to send command to host");
    }

    std::vector<BYTE> data;
    if (!mcp.WaitPending(devId, data, kMcpToolTimeoutMs))
        return BuildError(id, -32001, "Timeout waiting for window list");

    // 窗口标题编码由客户端能力位决定（新客户端 UTF-8，老客户端 CP_ACP）。
    UINT cp = GetClientEncoding(ctx);
    Json::Value wins = ParseWindowList(data, cp);
    int count = (int)wins.size();

    Json::Value result(Json::objectValue);
    Json::Value structuredContent(Json::objectValue);
    structuredContent["windows"] = wins;
    result["structuredContent"] = structuredContent;

    Json::Value content(Json::arrayValue);
    Json::Value item(Json::objectValue);
    item["type"] = "text";
    item["text"] = std::string(u8"共 ") + std::to_string(count) + std::string(u8" 个窗口。");
    content.append(item);
    result["content"] = content;
    result["isError"] = false;

    return BuildResult(id, result);
}

// tools/call：get_activity_history（主连接 RPC，无子链接、不弹框、不 CancelIO）
std::string BuildGetActivityHistory(const Json::Value& id, const Json::Value& args, CMy2015RemoteDlg* parent) {
    uint64_t devId = 0;
    std::string err;
    if (!ParseHostIdArg(args, devId, err))
        return BuildError(id, -32602, err);

    context* ctx = FindMainContext(parent, devId);
    if (!ctx)
        return BuildError(id, -32002, "Host not found or offline: " + std::to_string(devId));

    CMcpServer& mcp = CMcpServer::Instance();
    if (!mcp.BeginPending(devId, "get_activity_history"))
        return BuildError(id, -32003, "Device busy: another request is pending for this host");

    BYTE cmd = COMMAND_QUERY_ACTIVITY;
    if (!ctx->Send2Client(&cmd, 1)) {
        mcp.ClearPending(devId);
        return BuildError(id, -32004, "Failed to send command to host");
    }

    std::vector<BYTE> data;
    if (!mcp.WaitPending(devId, data, kMcpToolTimeoutMs))
        return BuildError(id, -32001, "Timeout waiting for activity history");

    // data[0]=token，其后为客户端 ActivityHistory::Dump() 的 UTF-8 纯文本
    // （标题在客户端已由 GetActiveWindowTitle 转 UTF-8，服务端无需再转码）。
    std::string text;
    if (data.size() > 1)
        text.assign((const char*)data.data() + 1, data.size() - 1);

    // 拆行：每行一条记录，过滤空行（含 Dump 末尾的换行）。
    Json::Value records(Json::arrayValue);
    size_t start = 0;
    while (start <= text.size()) {
        size_t nl = text.find('\n', start);
        std::string line = text.substr(start, (nl == std::string::npos ? text.size() : nl) - start);
        if (!line.empty()) records.append(line);
        if (nl == std::string::npos) break;
        start = nl + 1;
    }
    int count = (int)records.size();

    Json::Value result(Json::objectValue);
    Json::Value structuredContent(Json::objectValue);
    structuredContent["records"] = records;
    structuredContent["activityHistory"] = text;
    result["structuredContent"] = structuredContent;

    Json::Value content(Json::arrayValue);
    Json::Value item(Json::objectValue);
    item["type"] = "text";
    item["text"] = std::string(u8"共 ") + std::to_string(count) + std::string(u8" 条历史活动记录。");
    content.append(item);
    result["content"] = content;
    result["isError"] = false;

    return BuildResult(id, result);
}

// tools/call 分派
std::string BuildToolsCall(const Json::Value& root, CMy2015RemoteDlg* parent) {
    const Json::Value& id = root["id"];
    Json::Value params = root.isMember("params") ? root["params"] : Json::Value(Json::objectValue);

    std::string toolName;
    if (params.isObject() && params.isMember("name") && params["name"].isString()) {
        toolName = params["name"].asString();
    }

    const Json::Value args = GetCallArguments(params);

    if (toolName == "list_online_hosts") return BuildListOnlineHosts(id, parent);
    if (toolName == "search_hosts")      return BuildSearchHosts(id, args, parent);
    if (toolName == "get_host_detail")   return BuildGetHostDetail(id, args, parent);
    if (toolName == "list_processes")    return BuildListProcesses(id, args, parent);
    if (toolName == "list_windows")      return BuildListWindows(id, args, parent);
    if (toolName == "get_activity_history") return BuildGetActivityHistory(id, args, parent);

    return BuildError(id, -32602,
                      "Unknown tool: " + (toolName.empty() ? std::string("(empty)") : toolName));
}

}  // namespace

//////////////////////////////////////////////////////////////////////////
// CMcpServer Implementation
//////////////////////////////////////////////////////////////////////////

CMcpServer& CMcpServer::Instance() {
    static CMcpServer instance;
    return instance;
}

CMcpServer::CMcpServer() {
    m_server.Post("/mcp", [this](const httplib::Request& req, httplib::Response& res) {
        HandleMcp(req, res);
    });
}

CMcpServer::~CMcpServer() {
    Stop();  // 兜底：确保监听线程 join，避免 std::thread 析构触发 terminate
}

bool CMcpServer::Start(const std::string& bind, int port) {
    if (m_running.load()) return true;  // 已在运行

    m_thread = std::thread([this, bind, port]() {
        m_server.listen(bind, port);
    });

    // 给 listen 一点时间绑定端口；httplib::Server::is_running() 在 listen 内部置位。
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    m_running.store(m_server.is_running());
    return m_running.load();
}

void CMcpServer::Stop() {
    m_server.stop();
    if (m_thread.joinable()) {
        m_thread.join();
    }
    m_running.store(false);
}

void CMcpServer::HandleMcp(const httplib::Request& req, httplib::Response& res) {
    res.set_header("Content-Type", "application/json");

    // 静态 token 校验：Authorization: Bearer <token>（Start 前经 SetToken 保证非空）
    if (req.get_header_value("Authorization") != ("Bearer " + m_token)) {
        res.status = 401;
        res.set_content(BuildError(Json::nullValue, -32000, "Unauthorized"), "application/json");
        return;
    }

    // 解析 JSON-RPC 请求体
    Json::Value root;
    Json::CharReaderBuilder rbuilder;
    std::string errs;
    std::istringstream iss(req.body);
    if (!Json::parseFromStream(rbuilder, iss, &root, &errs) || !root.isObject()) {
        res.set_content(BuildError(Json::nullValue, -32700, "Parse error"), "application/json");
        return;
    }

    // 通知（无 id）→ 不返回 JSON-RPC 响应（如 notifications/initialized）
    if (!root.isMember("id")) {
        res.status = 202;
        res.set_content("", "application/json");
        return;
    }

    // 结构校验：缺 method
    if (!root.isMember("method") || !root["method"].isString()) {
        res.set_content(BuildError(root["id"], -32600, "Invalid Request"), "application/json");
        return;
    }

    std::string method = root["method"].asString();

    if (method == "initialize") {
        res.set_content(BuildInitializeResult(root["id"]), "application/json");
        return;
    }
    if (method == "ping") {
        res.set_content(BuildPingResult(root["id"]), "application/json");
        return;
    }
    if (method == "tools/list") {
        res.set_content(BuildToolsListResult(root["id"]), "application/json");
        return;
    }
    if (method == "tools/call") {
        res.set_content(BuildToolsCall(root, m_parent), "application/json");
        return;
    }

    // 未实现的方法
    res.set_content(BuildError(root["id"], -32601, "Method not found"), "application/json");
}

// ===== P2b 挂起请求注册表实现 =====

bool CMcpServer::IsPending(uint64_t device_id) {
    std::lock_guard<std::mutex> lk(m_PendingMutex);
    return m_Pending.find(device_id) != m_Pending.end();
}

void CMcpServer::TakeMainResponse(uint64_t device_id, const BYTE* data, ULONG len) {
    std::lock_guard<std::mutex> lk(m_PendingMutex);
    auto it = m_Pending.find(device_id);
    if (it == m_Pending.end()) return;  // 已超时清理 → 迟到数据，丢弃
    it->second.data.assign(data, data + len);
    it->second.done = true;
    m_PendingCv.notify_one();
}

bool CMcpServer::BeginPending(uint64_t device_id, const std::string& tool) {
    std::lock_guard<std::mutex> lk(m_PendingMutex);
    if (m_Pending.find(device_id) != m_Pending.end()) return false;  // 设备忙
    PendingRequest r;
    r.tool = tool;
    m_Pending[device_id] = std::move(r);
    return true;
}

bool CMcpServer::WaitPending(uint64_t device_id, std::vector<BYTE>& out, int timeoutMs) {
    std::unique_lock<std::mutex> lk(m_PendingMutex);
    auto it = m_Pending.find(device_id);
    if (it == m_Pending.end()) return false;

    bool signaled = m_PendingCv.wait_for(lk, std::chrono::milliseconds(timeoutMs),
                                         [&] { return it->second.done; });
    if (!signaled || it->second.data.empty()) {
        m_Pending.erase(it);  // 超时/空数据 → 清理
        return false;
    }
    out = std::move(it->second.data);
    m_Pending.erase(it);
    return true;
}

void CMcpServer::ClearPending(uint64_t device_id) {
    std::lock_guard<std::mutex> lk(m_PendingMutex);
    m_Pending.erase(device_id);
}

// rand_s：Windows CRT 加密安全随机源（基于系统 CSPRNG）。其声明需在 <stdlib.h> 前
// 定义 _CRT_RAND_S；为避免依赖 PCH 的包含顺序，这里手动声明其导出原型（errno_t == int）。
extern "C" int __cdecl rand_s(unsigned int* randomValue);

std::string GenerateRandomToken() {
    static const char hex[] = "0123456789abcdef";
    std::string out;
    out.reserve(32);
    for (int i = 0; i < 16; ++i) {
        unsigned int v = 0;
        if (rand_s(&v) != 0) {
            // rand_s 失败（罕见）：退化为时间 + 地址熵，保证仍返回非空 token。
            v = (unsigned int)(GetTickCount() ^ (ULONG_PTR)&out);
        }
        out.push_back(hex[(v >> 4) & 0xF]);
        out.push_back(hex[v & 0xF]);
    }
    return out;
}