RAG/src/main/java/com/wok/supportbot/config/McpClientManager.java

package com.wok.supportbot.config;

import com.wok.supportbot.entity.McpServerConfig;
import com.wok.supportbot.service.McpServerConfigService;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
import io.modelcontextprotocol.client.transport.ServerParameters;
import io.modelcontextprotocol.client.transport.StdioClientTransport;
import io.modelcontextprotocol.spec.McpSchema;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

import java.time.Duration;
import java.time.Instant;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * MCP 客户端生命周期管理器
 * 负责 MCP Client 的创建、缓存、刷新和销毁。
 * <p>
 * 支持两种传输模式：
 * - SSE（Server-Sent Events）：通过 HttpClientSseClientTransport 连接远程 MCP Server
 * - stdio（标准输入输出）：通过 StdioClientTransport 启动本地 MCP Server 进程
 */
@Component
@Slf4j
public class McpClientManager {

    @Autowired
    private McpServerConfigService mcpServerConfigService;

    /**
     * 客户端缓存：key = 配置ID 的字符串形式，value = MCP Client 实例
     * volatile 保证 refreshAll() 双缓冲切换时的可见性
     */
    private volatile ConcurrentHashMap<String, McpSyncClient> clientCache = new ConcurrentHashMap<>();

    /**
     * 不可用配置集合：记录已知不存在或未启用的配置ID，避免重复查询 DB
     * 配置被启用或新建时需同步移除此集合中的对应条目
     */
    private final Set<String> unavailableConfigs = ConcurrentHashMap.newKeySet();

    /**
     * 健康状态缓存：key = 配置ID 的字符串形式，value = 健康检查结果
     */
    private final ConcurrentHashMap<String, HealthStatus> healthCache = new ConcurrentHashMap<>();

    /**
     * MCP Server 健康状态记录
     */
    public record HealthStatus(
            /** 状态：ONLINE / OFFLINE / UNKNOWN */
            String status,
            /** 响应延迟（毫秒） */
            long latencyMs,
            /** 最后检查时间 */
            Instant lastCheckTime,
            /** 错误信息（健康时为 null） */
            String errorMessage
    ) {
        public static HealthStatus online(long latencyMs) {
            return new HealthStatus("ONLINE", latencyMs, Instant.now(), null);
        }

        public static HealthStatus offline(long latencyMs, String errorMessage) {
            return new HealthStatus("OFFLINE", latencyMs, Instant.now(), errorMessage);
        }

        public static HealthStatus unknown(String errorMessage) {
            return new HealthStatus("UNKNOWN", 0, Instant.now(), errorMessage);
        }
    }

    /**
     * MCP 客户端请求超时时间，支持 Duration 格式（如 30s、PT30S、60s）
     */
    @Value("${mcp.client.request-timeout:30s}")
    private String requestTimeoutStr;

    /**
     * MCP 客户端初始化超时时间，支持 Duration 格式（如 15s、PT15S、30s）
     */
    @Value("${mcp.client.init-timeout:15s}")
    private String initTimeoutStr;

    /**
     * 应用启动时自动加载所有已启用的 MCP Server 配置并建立连接
     * 确保第一次对话请求时 clientCache 已就绪，MCP 工具可以被注册到 ChatClient
     */
    @PostConstruct
    public void init() {
        log.info("McpClientManager 初始化，开始加载已启用的 MCP Server 配置...");
        refreshAll();
    }

    /**
     * 刷新所有 MCP 客户端连接（双缓冲策略）
     * 先在新 Map 中构建所有客户端，再原子切换引用，最后关闭旧客户端。
     * 避免清空与重建之间的请求全部失败。
     */
    public void refreshAll() {
        log.info("开始刷新所有 MCP 客户端连接...");

        // 获取所有启用的配置（一次查询代替两次按类型过滤）
        List<McpServerConfig> allActiveConfigs = mcpServerConfigService.listAllActiveConfigs();

        // 第一步：在全新 Map 中构建所有客户端（不影响现有缓存）
        ConcurrentHashMap<String, McpSyncClient> newCache = new ConcurrentHashMap<>();
        for (McpServerConfig config : allActiveConfigs) {
            try {
                McpSyncClient client = createClientDirectly(config);
                if (client != null) {
                    newCache.put(config.getId().toString(), client);
                }
            } catch (Exception e) {
                log.error("创建客户端失败: id={}, name={}, transportType={}, error={}",
                        config.getId(), config.getName(), config.getTransportType(), e.getMessage());
            }
        }

        // 第二步：原子切换引用（volatile 保证其他线程立即可见）
        ConcurrentHashMap<String, McpSyncClient> oldCache = clientCache;
        clientCache = newCache;
        unavailableConfigs.clear();

        // 第三步：关闭旧缓存中的客户端（不影响新请求）
        for (Map.Entry<String, McpSyncClient> entry : oldCache.entrySet()) {
            try {
                McpSyncClient client = entry.getValue();
                if (client != null) {
                    client.close();
                    log.debug("已关闭旧 MCP 客户端: configId={}", entry.getKey());
                }
            } catch (Exception e) {
                log.error("关闭旧 MCP 客户端失败: configId={}, error={}", entry.getKey(), e.getMessage());
            }
        }

        log.info("MCP 客户端刷新完成，当前缓存数量: {}", clientCache.size());
    }

    /**
     * 获取指定配置的客户端（懒加载）
     * 使用 computeIfAbsent 保证同一配置只创建一次客户端，避免并发竞态。
     * 如果缓存中不存在且不在不可用集合中，从 DB 读取配置后原子创建并缓存。
     *
     * @param configId 配置ID
     * @return MCP Client 实例，不存在或未启用返回 null
     */
    public McpSyncClient getClient(Long configId) {
        String key = configId.toString();

        // 快速路径：缓存命中
        McpSyncClient cached = clientCache.get(key);
        if (cached != null) {
            return cached;
        }

        // 快速路径：已知不可用的配置，直接跳过
        if (unavailableConfigs.contains(key)) {
            return null;
        }

        // 缓存未命中，使用 computeIfAbsent 保证同一 key 只创建一次客户端
        // ConcurrentHashMap.computeIfAbsent 对同一 key 加锁，避免并发线程重复创建
        // 注意：mapping function 不能返回 null（会抛 NPE），因此不可用的情况记录到 unavailableConfigs
        McpServerConfig config = mcpServerConfigService.getConfigById(configId);
        if (config == null) {
            log.warn("MCP 配置不存在: id={}", configId);
            unavailableConfigs.add(key);
            return null;
        }
        if (!Boolean.TRUE.equals(config.getIsActive())) {
            log.warn("MCP 配置未启用，跳过创建客户端: id={}, name={}", configId, config.getName());
            unavailableConfigs.add(key);
            return null;
        }

        // 配置存在且启用，通过 computeIfAbsent 原子创建（防止并发重复创建）
        return clientCache.computeIfAbsent(key, k -> createClientDirectly(config));
    }

    /**
     * 获取所有已启用的 MCP 工具描述
     * 遍历所有已缓存的客户端，调用 listTools() 获取真实的工具列表
     *
     * @return 工具描述列表，每个元素包含 config_id、name、transport_type、description、tools
     */
    public List<Map<String, Object>> listAvailableTools() {
        List<Map<String, Object>> tools = new ArrayList<>();

        log.info("🔍 listAvailableTools: clientCache 大小 = {}", clientCache.size());

        if (clientCache.isEmpty()) {
            log.warn("⚠️ clientCache 为空！MCP Server 可能未配置或连接失败。");
            log.warn("   请检查：1) 是否添加了 MCP Server 配置  2) 配置是否已启用  3) 启动日志是否有连接错误");
            return tools;
        }

        for (Map.Entry<String, McpSyncClient> entry : clientCache.entrySet()) {
            Long configId;
            try {
                configId = Long.parseLong(entry.getKey());
            } catch (NumberFormatException e) {
                log.warn("缓存 key 格式异常: {}", entry.getKey());
                continue;
            }

            McpServerConfig config = mcpServerConfigService.getConfigById(configId);
            if (config == null || !Boolean.TRUE.equals(config.getIsActive())) {
                continue;
            }

            Map<String, Object> toolInfo = new LinkedHashMap<>();
            toolInfo.put("config_id", config.getId().toString());
            toolInfo.put("name", config.getName());
            toolInfo.put("transport_type", config.getTransportType());
            toolInfo.put("description", config.getDescription());

            // 从 MCP Client 获取真实的工具列表
            List<Map<String, Object>> toolList = new ArrayList<>();
            try {
                McpSyncClient client = entry.getValue();
                McpSchema.ListToolsResult listToolsResult = client.listTools();
                if (listToolsResult != null && listToolsResult.tools() != null) {
                    log.info("   📦 MCP Server [{}] 暴露了 {} 个工具", config.getName(), listToolsResult.tools().size());
                    for (McpSchema.Tool tool : listToolsResult.tools()) {
                        Map<String, Object> toolMeta = new LinkedHashMap<>();
                        toolMeta.put("name", tool.name());
                        toolMeta.put("description", tool.description());
                        // JsonSchema 对象序列化为 JSON 字符串，供 McpToolCallback 构建 ToolDefinition 使用
                        if (tool.inputSchema() != null) {
                            toolMeta.put("inputSchema", ModelOptionsUtils.toJsonString(tool.inputSchema()));
                        } else {
                            toolMeta.put("inputSchema", "{}");
                        }
                        toolList.add(toolMeta);
                    }
                }
                log.info("   📦 MCP Server [{}] 获取到 {} 个工具", config.getName(), toolList.size());
                if (toolList.isEmpty()) {
                    log.warn("   ⚠️ MCP Server [{}] 未暴露任何工具，AI 模型将无法调用", config.getName());
                }
            } catch (Exception e) {
                log.error("获取 MCP 工具列表失败: serverId={}, name={}, error={}",
                        configId, config.getName(), e.getMessage());
            }

            toolInfo.put("tools", toolList);
            tools.add(toolInfo);
        }

        return tools;
    }

    /**
     * 定时健康检查：每 5 分钟执行一次
     * 遍历所有缓存的客户端，尝试 listTools() 探测连通性，记录状态到 healthCache
     */
    @Scheduled(fixedDelay = 300000)
    public void healthCheck() {
        log.info("开始 MCP Server 定时健康检查，当前缓存数量: {}", clientCache.size());
        for (Map.Entry<String, McpSyncClient> entry : clientCache.entrySet()) {
            String key = entry.getKey();
            checkSingleClient(key, entry.getValue());
        }
        log.info("MCP Server 定时健康检查完成");
    }

    /**
     * 手动触发单个配置的健康检查
     *
     * @param configId 配置ID
     * @return 该配置的最新健康状态
     */
    public HealthStatus checkHealthForConfig(Long configId) {
        String key = configId.toString();
        McpSyncClient client = clientCache.get(key);
        if (client == null) {
            HealthStatus status = HealthStatus.unknown("客户端未缓存或配置不存在");
            healthCache.put(key, status);
            return status;
        }
        return checkSingleClient(key, client);
    }

    /**
     * 对单个客户端执行健康探测
     * 通过调用 listTools() 判断服务端是否存活
     *
     * @param key    缓存 key（配置ID 字符串）
     * @param client MCP 客户端实例
     * @return 健康状态
     */
    private HealthStatus checkSingleClient(String key, McpSyncClient client) {
        long startTime = System.currentTimeMillis();
        try {
            McpSchema.ListToolsResult result = client.listTools();
            long latencyMs = System.currentTimeMillis() - startTime;
            int toolsCount = (result != null && result.tools() != null) ? result.tools().size() : 0;
            HealthStatus status = HealthStatus.online(latencyMs);
            healthCache.put(key, status);
            log.debug("MCP Server 健康检查通过: configId={}, latencyMs={}, toolsCount={}", key, latencyMs, toolsCount);
            return status;
        } catch (Exception e) {
            long latencyMs = System.currentTimeMillis() - startTime;
            HealthStatus status = HealthStatus.offline(latencyMs, e.getMessage());
            healthCache.put(key, status);
            log.warn("MCP Server 健康检查失败: configId={}, latencyMs={}, error={}", key, latencyMs, e.getMessage());
            return status;
        }
    }

    /**
     * 获取指定配置的健康状态
     *
     * @param configId 配置ID
     * @return 健康状态，未检查过返回 UNKNOWN
     */
    public HealthStatus getHealthStatus(Long configId) {
        HealthStatus status = healthCache.get(configId.toString());
        if (status == null) {
            return HealthStatus.unknown("尚未执行健康检查");
        }
        return status;
    }

    /**
     * 获取所有已缓存配置的健康状态
     *
     * @return Map：key = 配置ID 字符串，value = 健康状态
     */
    public Map<String, HealthStatus> getAllHealthStatus() {
        Map<String, HealthStatus> result = new LinkedHashMap<>();
        // 遍历所有缓存的客户端，未在 healthCache 中的标记为 UNKNOWN
        for (String key : clientCache.keySet()) {
            HealthStatus status = healthCache.get(key);
            if (status == null) {
                status = HealthStatus.unknown("尚未执行健康检查");
            }
            result.put(key, status);
        }
        return result;
    }

    /**
     * 关闭所有 MCP 客户端连接
     * 在应用关闭或手动刷新时调用
     */
    @PreDestroy
    public void shutdown() {
        log.info("关闭所有 MCP 客户端连接，当前数量: {}", clientCache.size());
        for (Map.Entry<String, McpSyncClient> entry : clientCache.entrySet()) {
            try {
                McpSyncClient client = entry.getValue();
                if (client != null) {
                    client.close();
                    log.debug("已关闭 MCP 客户端: configId={}", entry.getKey());
                }
            } catch (Exception e) {
                log.error("关闭 MCP 客户端失败: configId={}, error={}", entry.getKey(), e.getMessage());
            }
        }
        clientCache.clear();
        unavailableConfigs.clear();
        healthCache.clear();
        log.info("所有 MCP 客户端连接已关闭");
    }

    /**
     * 根据配置创建 MCP 客户端（不负责缓存，由调用方决定是否缓存）
     * <p>
     * SSE 模式：使用 HttpClientSseClientTransport 连接远程 MCP Server
     * stdio 模式：使用 StdioClientTransport 启动本地 MCP Server 进程
     *
     * @param config MCP Server 配置
     * @return 创建并初始化后的客户端实例，失败返回 null
     */
    private McpSyncClient createClientDirectly(McpServerConfig config) {
        try {
            McpSyncClient client;

            if ("sse".equals(config.getTransportType())) {
                // 创建 SSE 传输层：连接远程 MCP Server
                // HttpClientSseClientTransport 会自动将 /sse 作为 SSE 端点附加到 baseUri
                String serverUrl = config.getServerUrl().trim();
                // 如果 URL 以 /sse 结尾，去掉后缀作为 baseUri（transport 层会自动附加）
                String baseUri = serverUrl.replaceAll("/sse/?$", "");
                HttpClientSseClientTransport transport = HttpClientSseClientTransport.builder(baseUri).build();

                client = buildSyncClient(transport);
                log.info("正在初始化 SSE MCP 客户端: id={}, url={}", config.getId(), baseUri);

            } else if ("stdio".equals(config.getTransportType())) {
                // 创建 stdio 传输层：启动本地 MCP Server 进程
                String command = config.getCommand().trim();

                // 解析参数列表（逗号分隔）
                List<String> argsList = new ArrayList<>();
                if (config.getArgs() != null && !config.getArgs().trim().isEmpty()) {
                    argsList = Arrays.asList(config.getArgs().split("\\s*,\\s*"));
                }

                // 构建 ServerParameters
                ServerParameters.Builder paramsBuilder = ServerParameters.builder(command);
                if (!argsList.isEmpty()) {
                    paramsBuilder.args(argsList);
                }

                // 环境变量（Map<String, Object> -> Map<String, String>）
                if (config.getEnvVars() != null && !config.getEnvVars().isEmpty()) {
                    Map<String, String> envStr = new HashMap<>();
                    config.getEnvVars().forEach((k, v) -> envStr.put(k, v != null ? v.toString() : ""));
                    paramsBuilder.env(envStr);
                }

                StdioClientTransport transport = new StdioClientTransport(paramsBuilder.build());
                client = buildSyncClient(transport);
                log.info("正在初始化 stdio MCP 客户端: id={}, command={}", config.getId(), command);

            } else {
                log.error("不支持的传输类型: {}", config.getTransportType());
                return null;
            }

            // 初始化连接（握手协议）
            client.initialize();
            log.info("MCP 客户端初始化成功: id={}, name={}, transportType={}",
                    config.getId(), config.getName(), config.getTransportType());

            return client;

        } catch (Exception e) {
            log.error("创建 MCP 客户端失败: id={}, name={}, transportType={}, error={}",
                    config.getId(), config.getName(), config.getTransportType(), e.getMessage(), e);
            return null;
        }
    }

    /**
     * 构建 McpSyncClient 实例
     * 统一设置客户端信息和超时参数，超时时间从配置文件读取
     *
     * @param transport 传输层实例（SSE 或 stdio）
     * @return 构建的 McpSyncClient
     */
    private McpSyncClient buildSyncClient(io.modelcontextprotocol.spec.McpClientTransport transport) {
        Duration requestTimeout = parseDuration(requestTimeoutStr, Duration.ofSeconds(30));
        Duration initTimeout = parseDuration(initTimeoutStr, Duration.ofSeconds(15));

        return io.modelcontextprotocol.client.McpClient.sync(transport)
                .clientInfo(new McpSchema.Implementation("supportbot-mcp-client", "1.0.0"))
                .requestTimeout(requestTimeout)
                .initializationTimeout(initTimeout)
                .build();
    }

    /**
     * 解析 Duration 字符串，兼容 ISO-8601 格式（PT30S）和简化格式（30s、60s）
     * 解析失败时返回默认值
     *
     * @param durationStr  Duration 字符串
     * @param defaultValue 解析失败时的默认值
     * @return 解析后的 Duration
     */
    private Duration parseDuration(String durationStr, Duration defaultValue) {
        if (durationStr == null || durationStr.isBlank()) {
            return defaultValue;
        }
        try {
            // 先尝试 ISO-8601 格式（如 PT30S、PT1M30S）
            return Duration.parse(durationStr.trim());
        } catch (Exception e) {
            // 尝试简化格式：纯数字秒（如 "30"）或带 "s" 后缀（如 "30s"）
            String trimmed = durationStr.trim().toLowerCase();
            String numStr = trimmed.replaceAll("[^0-9.]+$", "");
            if (!numStr.isEmpty()) {
                try {
                    long seconds = Long.parseLong(numStr);
                    return Duration.ofSeconds(seconds);
                } catch (NumberFormatException ignored) {
                    // 继续尝试小数
                }
            }
            log.warn("无法解析 Duration 配置 '{}'，使用默认值: {}", durationStr, defaultValue);
            return defaultValue;
        }
    }

    // ==================== 增量缓存管理（供 Controller 层 CRUD 操作使用） ====================

    /**
     * 新增配置时：创建客户端并加入缓存
     * 仅对新建的单个配置创建连接，不影响其他已有连接。
     *
     * @param configId 新建配置的ID
     */
    public void addClient(Long configId) {
        String key = configId.toString();
        // 从不可用集合中移除（新建的配置可能之前被标记为不可用）
        unavailableConfigs.remove(key);

        McpServerConfig config = mcpServerConfigService.getConfigById(configId);
        if (config == null || !Boolean.TRUE.equals(config.getIsActive())) {
            return;
        }
        McpSyncClient client = createClientDirectly(config);
        if (client != null) {
            clientCache.put(key, client);
            log.info("增量添加 MCP 客户端: id={}, name={}", configId, config.getName());
        }
    }

    /**
     * 更新配置时：关闭旧客户端，重建新客户端
     * 仅对更新的单个配置重建连接，不影响其他已有连接。
     *
     * @param configId 更新的配置ID
     */
    public void rebuildClient(Long configId) {
        String key = configId.toString();
        // 先移除旧客户端
        removeClient(configId);

        McpServerConfig config = mcpServerConfigService.getConfigById(configId);
        if (config == null) {
            unavailableConfigs.add(key);
            return;
        }
        if (!Boolean.TRUE.equals(config.getIsActive())) {
            unavailableConfigs.add(key);
            return;
        }

        unavailableConfigs.remove(key);
        McpSyncClient client = createClientDirectly(config);
        if (client != null) {
            clientCache.put(key, client);
            log.info("增量重建 MCP 客户端: id={}, name={}", configId, config.getName());
        }
    }

    /**
     * 禁用配置时：关闭并移除客户端，标记为不可用
     *
     * @param configId 禁用的配置ID
     */
    public void disableClient(Long configId) {
        removeClient(configId);
        unavailableConfigs.add(configId.toString());
    }

    /**
     * 移除指定配置的客户端缓存
     * 配置删除或禁用时调用
     *
     * @param configId 配置ID
     */
    public void removeClient(Long configId) {
        String key = configId.toString();
        McpSyncClient removed = clientCache.remove(key);
        if (removed != null) {
            try {
                removed.close();
            } catch (Exception e) {
                log.error("关闭被移除的 MCP 客户端失败: configId={}, error={}", configId, e.getMessage());
            }
            log.info("已移除并关闭 MCP 客户端缓存: configId={}", configId);
        }
    }

    /**
     * 测试指定配置的 MCP Server 连接
     * 创建临时客户端 -> initialize -> listTools -> close，不影响正常缓存
     *
     * @param configId 配置ID
     * @return 测试结果 Map：success / latencyMs / toolsCount / errorMessage
     */
    public Map<String, Object> testConnection(Long configId) {
        McpServerConfig config = mcpServerConfigService.getConfigById(configId);
        if (config == null) {
            Map<String, Object> result = new LinkedHashMap<>();
            result.put("success", false);
            result.put("latencyMs", 0);
            result.put("toolsCount", 0);
            result.put("errorMessage", "配置不存在: id=" + configId);
            return result;
        }

        McpSyncClient tempClient = null;
        long startTime = System.currentTimeMillis();
        try {
            // 创建临时客户端（不加入缓存），复用 createClientDirectly 的传输层构建逻辑
            tempClient = createClientDirectly(config);
            if (tempClient == null) {
                Map<String, Object> result = new LinkedHashMap<>();
                result.put("success", false);
                result.put("latencyMs", System.currentTimeMillis() - startTime);
                result.put("toolsCount", 0);
                result.put("errorMessage", "客户端创建失败，请检查配置");
                return result;
            }

            // 获取工具列表
            McpSchema.ListToolsResult listToolsResult = tempClient.listTools();
            int toolsCount = (listToolsResult != null && listToolsResult.tools() != null)
                    ? listToolsResult.tools().size() : 0;

            long latencyMs = System.currentTimeMillis() - startTime;
            log.info("MCP 连接测试成功: id={}, name={}, latencyMs={}, toolsCount={}",
                    configId, config.getName(), latencyMs, toolsCount);

            Map<String, Object> result = new LinkedHashMap<>();
            result.put("success", true);
            result.put("latencyMs", latencyMs);
            result.put("toolsCount", toolsCount);
            result.put("errorMessage", null);
            return result;

        } catch (Exception e) {
            long latencyMs = System.currentTimeMillis() - startTime;
            log.warn("MCP 连接测试失败: id={}, name={}, latencyMs={}, error={}",
                    configId, config.getName(), latencyMs, e.getMessage());

            Map<String, Object> result = new LinkedHashMap<>();
            result.put("success", false);
            result.put("latencyMs", latencyMs);
            result.put("toolsCount", 0);
            result.put("errorMessage", e.getMessage());
            return result;

        } finally {
            // 确保关闭临时客户端
            if (tempClient != null) {
                try {
                    tempClient.close();
                } catch (Exception e) {
                    log.debug("关闭临时测试客户端失败: configId={}, error={}", configId, e.getMessage());
                }
            }
        }
    }

    /**
     * 获取当前缓存的客户端数量
     *
     * @return 缓存数量
     */
    public int getCacheSize() {
        return clientCache.size();
    }
}