/*
 *  Copyright (c) 2023-2026, Easy-Agents (fuhai999@gmail.com).
 *  <p>
 *  Licensed under the Apache License, Version 2.0 (the "License");
 */
package com.easyagents.mcp.client;

import com.alibaba.fastjson2.JSON;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.spec.McpSchema;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;

/**
 * MCP 配置与运行环境检测器。
 */
public class McpEnvironmentChecker {

    private static final Logger log = LoggerFactory.getLogger(McpEnvironmentChecker.class);
    private static final Duration COMMAND_TIMEOUT = Duration.ofSeconds(3);
    private static final Duration MCP_REQUEST_TIMEOUT = Duration.ofSeconds(10);
    private static final Set<String> SUPPORTED_TRANSPORTS = Set.of("stdio", "http-sse", "http-stream");
    private static final Set<String> KNOWN_VERSION_COMMANDS = Set.of(
        "node", "npm", "npx", "pnpm", "python", "python3", "pip", "pip3");
    private final boolean probeEnabled;
    private final Function<String, McpTransportFactory> transportFactoryProvider;

    /**
     * 创建启用连接探测的检测器。
     */
    public McpEnvironmentChecker() {
        this(true);
    }

    /**
     * 创建可控制连接探测行为的检测器。
     *
     * @param probeEnabled 是否启用 MCP 连接探测
     */
    McpEnvironmentChecker(boolean probeEnabled) {
        this(probeEnabled, McpEnvironmentChecker::defaultTransportFactory);
    }

    /**
     * 创建可注入 transport 工厂的检测器。
     *
     * @param probeEnabled 是否启用 MCP 连接探测
     * @param transportFactoryProvider transport 工厂提供器
     */
    McpEnvironmentChecker(boolean probeEnabled,
                          Function<String, McpTransportFactory> transportFactoryProvider) {
        this.probeEnabled = probeEnabled;
        this.transportFactoryProvider = transportFactoryProvider;
    }

    /**
     * 检测 MCP 配置。
     *
     * @param configJson MCP 配置 JSON
     * @return 检测结果
     */
    public McpEnvironmentCheckResult check(String configJson) {
        McpEnvironmentCheckResult result = new McpEnvironmentCheckResult();
        if (configJson == null || configJson.isBlank()) {
            result.setOverallStatus(McpCheckStatus.FAILED);
            result.addServer(failedServer("config", null, "configJson", "MCP 配置 JSON 不能为空", null));
            return result;
        }

        McpConfig config;
        try {
            config = JSON.parseObject(configJson, McpConfig.class);
        } catch (Exception error) {
            result.setOverallStatus(McpCheckStatus.FAILED);
            result.addServer(failedServer("config", null, "json", "MCP 配置 JSON 格式错误", sanitize(error)));
            return result;
        }

        if (config == null || config.getMcpServers() == null || config.getMcpServers().isEmpty()) {
            result.setOverallStatus(McpCheckStatus.FAILED);
            result.addServer(failedServer("config", null, "mcpServers", "mcpServers 不能为空", null));
            return result;
        }

        for (Map.Entry<String, McpConfig.ServerSpec> entry : config.getMcpServers().entrySet()) {
            result.addServer(checkServer(entry.getKey(), entry.getValue()));
        }
        return result;
    }

    private McpServerCheckResult checkServer(String serverName, McpConfig.ServerSpec spec) {
        McpServerCheckResult result = new McpServerCheckResult();
        result.setServerName(serverName);
        result.setTransport(transport(spec));

        if (serverName == null || serverName.isBlank()) {
            result.addCheck(McpCheckItem.of("serverName", McpCheckStatus.FAILED,
                "MCP 服务名称不能为空", null));
        }
        if (spec == null) {
            result.addCheck(McpCheckItem.of("server", McpCheckStatus.FAILED,
                "MCP 服务配置不能为空", null));
            return result;
        }

        String transport = transport(spec);
        result.setTransport(transport);
        if (!SUPPORTED_TRANSPORTS.contains(transport)) {
            result.addCheck(McpCheckItem.of("transport", McpCheckStatus.FAILED,
                "不支持的 MCP 传输类型", transport));
            return result;
        }
        result.addCheck(McpCheckItem.of("transport", McpCheckStatus.SUCCESS,
            "MCP 传输类型可用", transport));

        Map<String, String> resolvedEnv = resolveEnv(spec.getEnv(), result);
        if ("stdio".equals(transport)) {
            validateStdio(spec, result);
        } else {
            validateHttp(spec, result);
        }

        if (probeEnabled && result.getStatus() != McpCheckStatus.FAILED) {
            probe(serverName, spec, resolvedEnv, result);
        }
        return result;
    }

    private void validateStdio(McpConfig.ServerSpec spec, McpServerCheckResult result) {
        if (spec.getCommand() == null || spec.getCommand().isBlank()) {
            result.addCheck(McpCheckItem.of("command", McpCheckStatus.FAILED,
                "stdio MCP 必须配置 command", null));
            return;
        }
        result.addCheck(checkCommand(spec.getCommand()));
    }

    private void validateHttp(McpConfig.ServerSpec spec, McpServerCheckResult result) {
        if (spec.getUrl() == null || spec.getUrl().isBlank()) {
            result.addCheck(McpCheckItem.of("url", McpCheckStatus.FAILED,
                "HTTP MCP 必须配置 url", null));
            return;
        }
        result.addCheck(McpCheckItem.of("url", McpCheckStatus.SUCCESS,
            "MCP 连接地址已配置", spec.getUrl()));
    }

    private McpCheckItem checkCommand(String command) {
        String executable = executableName(command);
        boolean known = KNOWN_VERSION_COMMANDS.contains(executable);
        ProcessBuilder builder = known
            ? new ProcessBuilder(command, "--version")
            : new ProcessBuilder(command);
        try {
            Process process = builder.redirectErrorStream(true).start();
            boolean finished = process.waitFor(COMMAND_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
            if (!finished) {
                process.destroyForcibly();
                return McpCheckItem.of("command", McpCheckStatus.SUCCESS,
                    command + " 可启动", "版本检测超时，已终止检测进程");
            }
            String output = new String(process.getInputStream().readAllBytes(), StandardCharsets.UTF_8).trim();
            if (process.exitValue() == 0 || !known) {
                return McpCheckItem.of("command", McpCheckStatus.SUCCESS,
                    command + " 可用", firstLine(output));
            }
            return McpCheckItem.of("command", McpCheckStatus.WARNING,
                command + " 可启动但返回非零状态", firstLine(output));
        } catch (Exception error) {
            return McpCheckItem.of("command", McpCheckStatus.FAILED,
                "容器内未找到命令：" + command, sanitize(error));
        }
    }

    private void probe(String serverName,
                       McpConfig.ServerSpec spec,
                       Map<String, String> resolvedEnv,
                       McpServerCheckResult result) {
        CloseableTransport transport = null;
        McpSyncClient client = null;
        try {
            transport = transportFactoryProvider.apply(spec.getTransport()).create(spec, resolvedEnv);
            client = McpClient.sync(transport.getTransport())
                .requestTimeout(MCP_REQUEST_TIMEOUT)
                .build();
            client.initialize();
            McpSchema.ListToolsResult toolsResult = client.listTools();
            int toolCount = toolsResult == null || toolsResult.tools() == null ? 0 : toolsResult.tools().size();
            result.setToolCount(toolCount);
            if (toolCount == 0) {
                result.addCheck(McpCheckItem.of("tools", McpCheckStatus.WARNING,
                    "MCP 已连接，但没有发现工具", null));
            } else {
                result.addCheck(McpCheckItem.of("tools", McpCheckStatus.SUCCESS,
                    "MCP 工具列表获取成功", String.valueOf(toolCount)));
            }
        } catch (Exception error) {
            log.debug("MCP check failed for server: {}", serverName, error);
            result.addCheck(McpCheckItem.of("connection", McpCheckStatus.FAILED,
                "MCP 初始化或工具发现失败", sanitize(error)));
        } finally {
            closeQuietly(client);
            closeQuietly(transport);
        }
    }

    private Map<String, String> resolveEnv(Map<String, String> env, McpServerCheckResult result) {
        Map<String, String> resolved = new HashMap<>();
        if (env == null || env.isEmpty()) {
            result.addCheck(McpCheckItem.of("env", McpCheckStatus.SUCCESS,
                "未配置额外环境变量", null));
            return resolved;
        }
        for (Map.Entry<String, String> entry : env.entrySet()) {
            String key = entry.getKey();
            String value = entry.getValue();
            if (value != null && value.startsWith("${input:") && value.endsWith("}")) {
                String inputId = value.substring("${input:".length(), value.length() - 1);
                String resolvedValue = System.getProperty("mcp.input." + inputId);
                if (resolvedValue == null || resolvedValue.isBlank()) {
                    result.addCheck(McpCheckItem.of("env", McpCheckStatus.FAILED,
                        "环境变量未解析：" + key, "input:" + inputId));
                    continue;
                }
                resolved.put(key, resolvedValue);
                result.addCheck(McpCheckItem.of("env", McpCheckStatus.SUCCESS,
                    "环境变量已解析：" + key, "input:" + inputId));
                continue;
            }
            resolved.put(key, value);
            result.addCheck(McpCheckItem.of("env", McpCheckStatus.SUCCESS,
                "环境变量已配置：" + key, null));
        }
        return resolved;
    }

    private McpServerCheckResult failedServer(String serverName,
                                              String transport,
                                              String name,
                                              String message,
                                              String detail) {
        McpServerCheckResult server = new McpServerCheckResult();
        server.setServerName(serverName);
        server.setTransport(transport);
        server.addCheck(McpCheckItem.of(name, McpCheckStatus.FAILED, message, detail));
        return server;
    }

    private static McpTransportFactory defaultTransportFactory(String transportType) {
        return switch (transport(transportType)) {
            case "stdio" -> new StdioTransportFactory();
            case "http-sse" -> new HttpSseTransportFactory();
            case "http-stream" -> new HttpStreamTransportFactory();
            default -> throw new IllegalArgumentException("Unsupported transport: " + transportType);
        };
    }

    private String transport(McpConfig.ServerSpec spec) {
        return spec == null ? "stdio" : transport(spec.getTransport());
    }

    private static String transport(String value) {
        return value == null || value.isBlank() ? "stdio" : value.toLowerCase(Locale.ROOT);
    }

    private String executableName(String command) {
        int slash = Math.max(command.lastIndexOf('/'), command.lastIndexOf('\\'));
        String name = slash >= 0 ? command.substring(slash + 1) : command;
        return name.endsWith(".cmd") ? name.substring(0, name.length() - 4) : name;
    }

    private String firstLine(String value) {
        if (value == null || value.isBlank()) {
            return null;
        }
        int lineEnd = value.indexOf('\n');
        String line = lineEnd >= 0 ? value.substring(0, lineEnd) : value;
        return sanitize(line.trim());
    }

    private String sanitize(Throwable error) {
        if (error == null) {
            return null;
        }
        String message = error.getMessage();
        return message == null || message.isBlank() ? error.getClass().getSimpleName() : sanitize(message);
    }

    private String sanitize(String value) {
        if (value == null) {
            return null;
        }
        String sanitized = value.replaceAll("(?i)(api[_-]?key|token|secret|password)=([^\\s,;]+)", "$1=******");
        return sanitized.length() > 500 ? sanitized.substring(0, 500) : sanitized;
    }

    private void closeQuietly(AutoCloseable closeable) {
        if (closeable == null) {
            return;
        }
        try {
            closeable.close();
        } catch (Exception error) {
            log.debug("Failed to close MCP check resource.", error);
        }
    }
}