SimpleRemoter/server/2015Remote/McpServer.h

#pragma once

#include <Windows.h>
#include <string>
#include <thread>
#include <atomic>
#include <chrono>
#include <mutex>
#include <condition_variable>
#include <map>
#include <vector>
#include <cstdint>

// httplib 与 Windows 头部的 min/max 宏冲突，按 file_server.h 的既有约定处理。
#undef min
#undef max
#include "httplib.h"
#ifndef max
#define max(a,b) (((a) > (b)) ? (a) : (b))
#endif
#ifndef min
#define min(a,b) (((a) < (b)) ? (a) : (b))
#endif

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

class CMy2015RemoteDlg;
class context;

// MCP (Model Context Protocol) 服务端：httplib 封装 + JSON-RPC 2.0 分发 + 静态 token 校验。
// 与 CWebService 平级、互不依赖；默认禁用，经「扩展 → MCP设置」开启后监听
// （默认 127.0.0.1:6544，仅本机回环）。见 docs/Mcp_Design.md。
class CMcpServer {
public:
    static CMcpServer& Instance();

    void SetParentDlg(CMy2015RemoteDlg* pDlg) { m_parent = pDlg; }
    void SetToken(const std::string& token) { m_token = token; }

    bool Start(const std::string& bind, int port);
    void Stop();
    bool IsRunning() const { return m_running.load(); }

    // ===== P2b：无请求 id 的一次性响应 → 每 host 单飞行 =====
    // 进程/窗口列表（TOKEN_PSLIST/TOKEN_WSLIST，一次性子链接）与历史活动
    // （TOKEN_REPORT_ACTIVITY，主连接 RPC）都不带请求关联，同一 host 同一时刻
    // 只允许一个挂起请求，避免同 host 并发请求响应归属歧义。
    // 注意：单飞行只约束 MCP 侧同 host 并发，不约束与 MFC 对话框并存（各用独立子链接）。

    // 该 host 是否已有挂起请求（MessageHandle 在 TOKEN_PSLIST/TOKEN_WSLIST 分支调用）。
    bool IsPending(uint64_t device_id);

    // 拷贝响应缓冲到挂起结果并唤醒等待者（在 MessageHandle 内同步调用，数据此时仍在
    // context 缓冲中、IO 线程被阻塞，拷贝是安全的）。
    void TakeMainResponse(uint64_t device_id, const BYTE* data, ULONG len);

    // 工具线程：登记挂起（false = 该 host 已有挂起请求，设备忙）。
    bool BeginPending(uint64_t device_id, const std::string& tool);

    // 工具线程：等待响应；成功返回 true 并把缓冲写入 out，超时/失败返回 false（并清理）。
    bool WaitPending(uint64_t device_id, std::vector<BYTE>& out, int timeoutMs);

    // 工具线程：清理挂起状态（发送失败等提前退出路径）。
    void ClearPending(uint64_t device_id);

    // ===== P2c：list_files / get_screenshot 扩展 =====

    // 登记挂起（带目录路径；list_files 专用，path 为空表示只列盘）。
    bool BeginPending(uint64_t device_id, const std::string& tool, const std::string& path);

    // get_screenshot：工具线程生成 16 位 reqId（跳过 0，0 表示未设置），用于丢弃过期响应。
    uint16_t NextPreviewReqId();

    // get_screenshot：登记期望的 reqId（必须在发送 COMMAND_SCREEN_PREVIEW_REQ 前调用）。
    void SetPendingReqId(uint64_t device_id, uint16_t reqId);

    // get_screenshot：MessageHandle(TOKEN_SCREEN_PREVIEW_RSP) 在 IO 线程同步调用；
    // 挂起命中且 reqId 一致则拷贝响应并返回 true，否则 false（回落到 MFC 预览路径）。
    bool TakePreviewResponse(uint64_t device_id, uint16_t reqId, const BYTE* data, ULONG len);

    // list_files：MessageHandle(TOKEN_DRIVE_LIST) 调用。
    // 返回 true = 已取走盘列表（调用方 CancelIO）；false = 已转向下发 COMMAND_LIST_FILES、
    // 等 TOKEN_FILE_LIST（调用方不关子链接）。
    bool OnDriveList(uint64_t device_id, context* subCtx, const BYTE* buf, ULONG len);

    // ===== P3：list_registry（注册表查询，三阶段：COMMAND_REGEDIT → TOKEN_REGEDIT →
    // COMMAND_REG_FIND → TOKEN_REG_PATH + TOKEN_REG_KEY）=====

    // MessageHandle(TOKEN_REGEDIT) 调用：挂起命中且为 list_registry 则下发 COMMAND_REG_FIND
    // 并返回 true（接管子链接，不打开 MFC 对话框）；否则 false（回落 MFC）。
    bool OnRegeditReady(uint64_t device_id, context* subCtx);

    // MessageHandle(TOKEN_REG_PATH / TOKEN_REG_KEY) 调用：两包都到齐后置 done 并唤醒等待者。
    void TakeRegPath(uint64_t device_id, const BYTE* data, ULONG len);
    void TakeRegKey(uint64_t device_id, const BYTE* data, ULONG len);

    // 工具线程：等待注册表两包（子键 + 值）；成功返回 true 并把两包分别写入 out，超时返回 false。
    bool WaitPendingRegistry(uint64_t device_id, std::vector<BYTE>& subkeys,
                             std::vector<BYTE>& values, int timeoutMs);

private:
    CMcpServer();
    ~CMcpServer();
    CMcpServer(const CMcpServer&) = delete;
    CMcpServer& operator=(const CMcpServer&) = delete;

    // POST /mcp 处理器：token 校验 + JSON-RPC 分发。
    void HandleMcp(const httplib::Request& req, httplib::Response& res);

    httplib::Server    m_server;
    std::thread        m_thread;
    std::atomic<bool>  m_running{false};
    std::string        m_token;
    CMy2015RemoteDlg*  m_parent = nullptr;

    // 挂起请求注册表（受 m_PendingMutex 保护，键 = device_id）。
    struct PendingRequest {
        std::string       tool;
        std::string       path;              // list_files：目录路径；list_registry：rootToken(1B)+相对子键路径
        uint16_t          expectedReqId = 0; // get_screenshot：期望的预览 reqId（0 = 未设置）
        std::vector<BYTE> data;
        std::vector<BYTE> regPath;           // list_registry：TOKEN_REG_PATH 子键包（含 token 字节）
        std::vector<BYTE> regKey;            // list_registry：TOKEN_REG_KEY 值包（含 token 字节）
        bool              pathDone = false;  // list_registry：已收到 TOKEN_REG_PATH
        bool              keyDone  = false;  // list_registry：已收到 TOKEN_REG_KEY
        bool              done = false;
    };
    std::mutex                   m_PendingMutex;
    std::condition_variable      m_PendingCv;
    std::map<uint64_t, PendingRequest> m_Pending;

    // get_screenshot 的 reqId 发生器（16 位，跳过 0）。与 MFC 预览的 m_PreviewReqId 各自独立计数。
    std::atomic<uint16_t>        m_PreviewReqId{1};
};

// 全局访问器（仿 WebService()，见 WebService.h 末尾）
inline CMcpServer& McpServer() { return CMcpServer::Instance(); }

// 生成 32 hex（128-bit）随机 token；供运行时兜底与「MCP设置」对话框预填复用。
std::string GenerateRandomToken();