package com.easyagents.agent.runtime.mcp;

import io.agentscope.core.tool.mcp.McpClientWrapper;
import io.modelcontextprotocol.spec.McpSchema;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * 为 MCP client 增加运行时工具别名。
 */
class AliasedMcpClientWrapper extends McpClientWrapper {

    static final String RAW_TOOL_NAME_META_KEY = "easyagentsRawMcpToolName";

    private final McpClientWrapper delegate;
    private final Map<String, String> rawToAlias;
    private final Map<String, String> aliasToRaw;
    private final String toolNamePrefix;

    /**
     * 创建 MCP client 别名包装器。
     *
     * @param delegate 原始 MCP client
     * @param rawToAlias 原始工具名到运行时工具名的映射
     */
    AliasedMcpClientWrapper(McpClientWrapper delegate, Map<String, String> rawToAlias) {
        this(delegate, rawToAlias, null);
    }

    /**
     * 创建 MCP client 别名包装器。
     *
     * @param delegate 原始 MCP client
     * @param rawToAlias 原始工具名到运行时工具名的映射
     * @param toolNamePrefix 动态工具名前缀
     */
    AliasedMcpClientWrapper(McpClientWrapper delegate, Map<String, String> rawToAlias, String toolNamePrefix) {
        super(delegate == null ? "mcp" : delegate.getName());
        this.delegate = delegate;
        this.rawToAlias = rawToAlias == null ? Map.of() : new LinkedHashMap<>(rawToAlias);
        this.aliasToRaw = new LinkedHashMap<>();
        this.toolNamePrefix = toolNamePrefix == null || toolNamePrefix.isBlank() ? null : toolNamePrefix.trim();
        this.rawToAlias.forEach((rawName, aliasName) -> {
            if (rawName != null && aliasName != null && !aliasName.isBlank()) {
                this.aliasToRaw.put(aliasName, rawName);
            }
        });
    }

    /**
     * 初始化底层 MCP client。
     *
     * @return 初始化完成信号
     */
    @Override
    public Mono<Void> initialize() {
        return delegate.initialize().doOnSuccess(ignored -> initialized = delegate.isInitialized());
    }

    /**
     * 返回已替换为运行时别名的工具列表。
     *
     * @return 工具列表
     */
    @Override
    public Mono<List<McpSchema.Tool>> listTools() {
        return delegate.listTools().map(this::aliasTools);
    }

    /**
     * 调用 MCP 工具，运行时别名会映射回原始工具名。
     *
     * @param toolName 运行时工具名
     * @param arguments 工具参数
     * @return 工具调用结果
     */
    @Override
    public Mono<McpSchema.CallToolResult> callTool(String toolName, Map<String, Object> arguments) {
        return delegate.callTool(rawToolName(toolName), arguments);
    }

    /**
     * 关闭底层 MCP client。
     */
    @Override
    public void close() {
        delegate.close();
        initialized = false;
    }

    private List<McpSchema.Tool> aliasTools(List<McpSchema.Tool> tools) {
        if (tools == null || tools.isEmpty()) {
            cachedTools.clear();
            return List.of();
        }
        List<McpSchema.Tool> aliased = new ArrayList<>();
        cachedTools.clear();
        Map<String, String> usedAliases = new LinkedHashMap<>();
        for (McpSchema.Tool tool : tools) {
            if (tool == null) {
                continue;
            }
            McpSchema.Tool aliasTool = aliasTool(tool, usedAliases);
            cachedTools.put(aliasTool.name(), aliasTool);
            aliased.add(aliasTool);
        }
        return aliased;
    }

    private McpSchema.Tool aliasTool(McpSchema.Tool tool, Map<String, String> usedAliases) {
        String aliasName = uniqueAliasName(aliasName(tool.name()), tool.name(), usedAliases);
        if (aliasName == null || aliasName.isBlank() || aliasName.equals(tool.name())) {
            return tool;
        }
        aliasToRaw.put(aliasName, tool.name());
        Map<String, Object> meta = new LinkedHashMap<>();
        if (tool.meta() != null) {
            meta.putAll(tool.meta());
        }
        meta.put(RAW_TOOL_NAME_META_KEY, tool.name());
        return new McpSchema.Tool(aliasName, tool.title(), tool.description(), tool.inputSchema(),
                tool.outputSchema(), tool.annotations(), meta);
    }

    private String uniqueAliasName(String aliasName, String rawName, Map<String, String> usedAliases) {
        if (aliasName == null || aliasName.isBlank()) {
            return aliasName;
        }
        String existingRawName = usedAliases.get(aliasName);
        if (existingRawName == null || existingRawName.equals(rawName)) {
            usedAliases.put(aliasName, rawName);
            return aliasName;
        }
        int suffix = 2;
        String candidate = aliasName + "_" + suffix;
        while (usedAliases.containsKey(candidate)) {
            suffix++;
            candidate = aliasName + "_" + suffix;
        }
        usedAliases.put(candidate, rawName);
        return candidate;
    }

    private String aliasName(String rawName) {
        String explicitAlias = rawToAlias.get(rawName);
        if (explicitAlias != null && !explicitAlias.isBlank()) {
            return explicitAlias;
        }
        if (toolNamePrefix == null) {
            return rawName;
        }
        return toolNamePrefix + safeToolNameSegment(rawName);
    }

    private String safeToolNameSegment(String value) {
        String normalized = String.valueOf(value == null ? "" : value).trim()
                .replaceAll("[^A-Za-z0-9_-]", "_")
                .replaceAll("_+", "_");
        if (normalized.isBlank()) {
            return "tool";
        }
        return normalized;
    }

    private String rawToolName(String toolName) {
        return aliasToRaw.getOrDefault(toolName, toolName);
    }
}
