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 Tool 清单冻结为只读白名单视图。
 */
final class FrozenMcpClientWrapper extends McpClientWrapper {

    private final McpClientWrapper delegate;
    private final List<McpSchema.Tool> frozenTools;
    private final Map<String, String> runtimeToRaw = new LinkedHashMap<>();

    /**
     * 创建冻结 MCP client 视图。
     *
     * @param delegate 原始 client
     * @param actualTools 已一次性读取并验证的远端 Tool
     * @param manifest 冻结清单
     * @param aliases 显式运行别名
     * @param prefix 运行名前缀
     */
    FrozenMcpClientWrapper(McpClientWrapper delegate,
                           List<McpSchema.Tool> actualTools,
                           List<McpToolManifestEntry> manifest,
                           Map<String, String> aliases,
                           String prefix) {
        super(delegate == null ? "mcp" : delegate.getName());
        this.delegate = delegate;
        this.frozenTools = freeze(actualTools, manifest, aliases, prefix);
        this.frozenTools.forEach(tool -> cachedTools.put(tool.name(), tool));
    }

    /** {@inheritDoc} */
    @Override
    public Mono<Void> initialize() {
        return delegate.initialize().doOnSuccess(ignored -> initialized = delegate.isInitialized());
    }

    /** {@inheritDoc} */
    @Override
    public Mono<List<McpSchema.Tool>> listTools() {
        return Mono.just(frozenTools);
    }

    /** {@inheritDoc} */
    @Override
    public Mono<McpSchema.CallToolResult> callTool(String toolName, Map<String, Object> arguments) {
        return delegate.callTool(runtimeToRaw.getOrDefault(toolName, toolName), arguments);
    }

    /** {@inheritDoc} */
    @Override
    public void close() {
        delegate.close();
        initialized = false;
    }

    /**
     * 按冻结 manifest 顺序裁剪并应用稳定运行别名。
     *
     * @param actualTools 远端当前 Tool
     * @param manifest 冻结清单
     * @param aliases 显式别名
     * @param prefix 动态前缀
     * @return 不可变 Tool 白名单
     */
    private List<McpSchema.Tool> freeze(List<McpSchema.Tool> actualTools,
                                        List<McpToolManifestEntry> manifest,
                                        Map<String, String> aliases,
                                        String prefix) {
        Map<String, McpSchema.Tool> actualByName = new LinkedHashMap<>();
        if (actualTools != null) {
            actualTools.stream().filter(tool -> tool != null && tool.name() != null)
                    .forEach(tool -> actualByName.put(tool.name(), tool));
        }
        Map<String, String> usedRuntimeNames = new LinkedHashMap<>();
        List<McpSchema.Tool> result = new ArrayList<>();
        for (McpToolManifestEntry entry : manifest) {
            McpSchema.Tool actual = actualByName.get(entry.getName());
            if (actual == null) {
                throw new IllegalStateException("Frozen MCP tool is missing after validation: " + entry.getName());
            }
            String runtimeName = uniqueRuntimeName(
                    runtimeName(actual.name(), aliases, prefix), actual.name(), usedRuntimeNames);
            runtimeToRaw.put(runtimeName, actual.name());
            Map<String, Object> meta = new LinkedHashMap<>();
            if (actual.meta() != null) {
                meta.putAll(actual.meta());
            }
            if (!runtimeName.equals(actual.name())) {
                meta.put(AliasedMcpClientWrapper.RAW_TOOL_NAME_META_KEY, actual.name());
            }
            // 模型可见描述也必须来自发布时冻结清单，避免远端描述在运行中漂移。
            result.add(new McpSchema.Tool(runtimeName, actual.title(), entry.getDescription(),
                    actual.inputSchema(), actual.outputSchema(), actual.annotations(), meta));
        }
        return List.copyOf(result);
    }

    /**
     * 计算单个 Tool 的运行名。
     *
     * @param rawName 原始名称
     * @param aliases 显式别名
     * @param prefix 动态前缀
     * @return 运行名
     */
    private String runtimeName(String rawName, Map<String, String> aliases, String prefix) {
        String alias = aliases == null ? null : aliases.get(rawName);
        if (alias != null && !alias.isBlank()) {
            return alias;
        }
        if (prefix == null || prefix.isBlank()) {
            return rawName;
        }
        String segment = String.valueOf(rawName == null ? "" : rawName).trim()
                .replaceAll("[^A-Za-z0-9_-]", "_")
                .replaceAll("_+", "_");
        return prefix.trim() + (segment.isBlank() ? "tool" : segment);
    }

    /**
     * 避免别名碰撞。
     *
     * @param candidate 候选运行名
     * @param rawName 原始名称
     * @param used 已使用运行名
     * @return 唯一运行名
     */
    private String uniqueRuntimeName(String candidate,
                                     String rawName,
                                     Map<String, String> used) {
        String existing = used.get(candidate);
        if (existing == null || existing.equals(rawName)) {
            used.put(candidate, rawName);
            return candidate;
        }
        int suffix = 2;
        String value = candidate + "_" + suffix;
        while (used.containsKey(value)) {
            suffix++;
            value = candidate + "_" + suffix;
        }
        used.put(value, rawName);
        return value;
    }
}
