package com.easyagents.agent.runtime.mcp;

import com.easyagents.agent.runtime.AgentRuntimeException;
import com.easyagents.agent.runtime.hitl.AgentToolApprovalRequest;
import com.easyagents.agent.runtime.tool.AgentToolCategory;
import com.easyagents.agent.runtime.tool.AgentToolSpec;
import com.easyagents.agent.runtime.tool.AgentToolVisibility;
import io.agentscope.core.tool.Toolkit;
import io.agentscope.core.tool.mcp.McpClientWrapper;
import io.agentscope.core.tool.mcp.McpTool;
import io.modelcontextprotocol.spec.McpSchema;

import java.util.*;
import java.util.Locale;

/**
 * 将 MCP 运行时声明注册到 AgentScope Toolkit。
 */
public class McpToolkitAdapter {

    private final McpClientFactory clientFactory;

    /**
     * 使用默认 MCP client factory 创建适配器。
     */
    public McpToolkitAdapter() {
        this(new McpClientFactory());
    }

    /**
     * 使用指定 MCP client factory 创建适配器。
     *
     * @param clientFactory MCP client factory
     */
    public McpToolkitAdapter(McpClientFactory clientFactory) {
        this.clientFactory = clientFactory == null ? new McpClientFactory() : clientFactory;
    }

    /**
     * 将 MCP 声明注册到 Toolkit。
     *
     * @param specs MCP 声明
     * @param toolkit AgentScope Toolkit
     * @return MCP 注册结果
     */
    public McpRegistration register(List<McpSpec> specs, Toolkit toolkit) {
        if (specs == null || specs.isEmpty()) {
            return McpRegistration.empty();
        }
        if (toolkit == null) {
            throw new AgentRuntimeException("AgentScope toolkit is required for MCP registration.");
        }
        List<McpClientWrapper> clients = new ArrayList<>();
        List<AgentToolSpec> toolSpecs = new ArrayList<>();
        List<McpSkillRegistration> skillRegistrations = new ArrayList<>();
        try {
            for (McpSpec spec : specs) {
                if (spec == null) {
                    continue;
                }
                McpSpecValidator.validateConnection(spec);
                McpClientWrapper client = clientFactory.create(spec);
                clients.add(client);
                if (isSkillBound(spec)) {
                    List<McpSchema.Tool> actualTools = initializeAndListTools(client);
                    McpToolManifest.assertFrozenManifest(spec, actualTools);
                    client = new FrozenMcpClientWrapper(client, actualTools,
                            spec.getFrozenToolManifest(), spec.getToolAliases(), spec.getToolNamePrefix());
                } else {
                    client = applyAliases(spec, client);
                }
                clients.set(clients.size() - 1, client);
                if (isSkillBound(spec)) {
                    // Skill MCP 必须以冻结 manifest 派生白名单，调用方不能通过空列表放宽到远端全部 Tool。
                    spec.setEnableTools(frozenRuntimeToolNames(spec, client));
                    skillRegistrations.add(new McpSkillRegistration(
                            spec.getSkillId(), client, spec.getEnableTools(), spec.getDisableTools(),
                            spec.getPresetParameters()));
                } else {
                    registerClient(spec, client, toolkit);
                }
                toolSpecs.addAll(toToolSpecs(spec, registeredTools(spec, client)));
            }
        } catch (RuntimeException error) {
            closeQuietly(clients);
            throw error;
        }
        return new McpRegistration(clients, toolSpecs, skillRegistrations);
    }

    /**
     * 根据冻结原始 Tool 名称和别名后的远端清单生成强制运行白名单。
     *
     * @param spec Skill MCP 声明
     * @param client 已应用运行别名的 client
     * @return 冻结 Tool 对应的运行名
     */
    private List<String> frozenRuntimeToolNames(McpSpec spec, McpClientWrapper client) {
        Set<String> frozenRawNames = new LinkedHashSet<>();
        for (McpToolManifestEntry entry : spec.getFrozenToolManifest()) {
            if (entry != null && entry.getName() != null && !entry.getName().isBlank()) {
                frozenRawNames.add(entry.getName());
            }
        }
        List<String> names = new ArrayList<>();
        for (McpSchema.Tool tool : listTools(client)) {
            if (tool != null && frozenRawNames.contains(rawToolName(spec, tool))) {
                names.add(tool.name());
            }
        }
        if (names.size() != frozenRawNames.size()) {
            throw new AgentRuntimeException("Skill-bound MCP frozen tool aliases are incomplete: "
                    + spec.getName());
        }
        return List.copyOf(names);
    }

    private McpClientWrapper applyAliases(McpSpec spec, McpClientWrapper client) {
        boolean hasExplicitAliases = spec.getToolAliases() != null && !spec.getToolAliases().isEmpty();
        boolean hasDynamicPrefix = spec.getToolNamePrefix() != null && !spec.getToolNamePrefix().isBlank();
        if (!hasExplicitAliases && !hasDynamicPrefix) {
            return client;
        }
        return new AliasedMcpClientWrapper(client, spec.getToolAliases(), spec.getToolNamePrefix());
    }

    private void registerClient(McpSpec spec, McpClientWrapper client, Toolkit toolkit) {
        String groupName = blankToNull(spec.getGroupName());
        if (groupName != null && toolkit.getToolGroup(groupName) == null) {
            toolkit.createToolGroup(groupName, spec.getDescription(), true);
        }
        toolkit.registration()
                .mcpClient(client)
                .enableTools(emptyToNull(spec.getEnableTools()))
                .disableTools(emptyToNull(spec.getDisableTools()))
                .group(groupName)
                .presetParameters(emptyToNull(spec.getPresetParameters()))
                .apply();
    }

    private List<McpSchema.Tool> registeredTools(McpSpec spec, McpClientWrapper client) {
        List<McpSchema.Tool> tools = listTools(client);
        if (tools == null || tools.isEmpty()) {
            return List.of();
        }
        List<McpSchema.Tool> filtered = new ArrayList<>();
        for (McpSchema.Tool tool : tools) {
            if (tool != null && shouldRegister(tool.name(), spec.getEnableTools(), spec.getDisableTools())) {
                filtered.add(tool);
            }
        }
        return filtered;
    }

    private List<McpSchema.Tool> listTools(McpClientWrapper client) {
        List<McpSchema.Tool> tools = client.listTools().block();
        return tools == null ? List.of() : tools;
    }

    /**
     * 初始化 client 后读取一次远端 Tool 清单。
     *
     * @param client MCP client
     * @return Tool 清单
     */
    private List<McpSchema.Tool> initializeAndListTools(McpClientWrapper client) {
        // AgentScope validates the initialized flag when listTools() is invoked. Build the
        // second publisher only after initialization has completed, otherwise eager publisher
        // assembly can fail even though the server initializes successfully moments later.
        client.initialize().block();
        List<McpSchema.Tool> tools = client.listTools().block();
        return tools == null ? List.of() : tools;
    }

    private boolean isSkillBound(McpSpec spec) {
        return spec.getSkillId() != null && !spec.getSkillId().isBlank();
    }

    private boolean shouldRegister(String toolName, List<String> enableTools, List<String> disableTools) {
        if (enableTools != null && !enableTools.isEmpty()) {
            return enableTools.contains(toolName);
        }
        return disableTools == null || disableTools.isEmpty() || !disableTools.contains(toolName);
    }

    private List<AgentToolSpec> toToolSpecs(McpSpec spec, List<McpSchema.Tool> tools) {
        if (tools == null || tools.isEmpty()) {
            return List.of();
        }
        List<AgentToolSpec> toolSpecs = new ArrayList<>();
        for (McpSchema.Tool tool : tools) {
            AgentToolSpec toolSpec = new AgentToolSpec();
            Set<String> excludedPresetNames = Set.of();
            Map<String, Object> toolPresetParameters = spec.getPresetParameters() == null
                    ? null
                    : spec.getPresetParameters().get(tool.name());
            if (toolPresetParameters != null) {
                excludedPresetNames = toolPresetParameters.keySet();
            }
            toolSpec.setName(tool.name());
            toolSpec.setDescription(tool.description());
            toolSpec.setCategory(AgentToolCategory.MCP);
            toolSpec.setVisibility(AgentToolVisibility.VISIBLE);
            toolSpec.setParametersSchema(McpTool.convertMcpSchemaToParameters(tool.inputSchema(), excludedPresetNames));
            toolSpec.setOutputSchema(tool.outputSchema());
            AgentToolApprovalRequest toolApprovalRequest = toolApprovalRequest(spec, tool.name());
            toolSpec.setApprovalRequired(spec.isApprovalRequired() || toolApprovalRequest != null);
            toolSpec.setApprovalRequest(toolApprovalRequest == null ? spec.getApprovalRequest() : toolApprovalRequest);
            toolSpec.setMetadata(metadata(spec, tool));
            toolSpecs.add(toolSpec);
        }
        return toolSpecs;
    }

    private AgentToolApprovalRequest toolApprovalRequest(McpSpec spec, String toolName) {
        if (spec.getToolApprovalRequests() == null || spec.getToolApprovalRequests().isEmpty()) {
            return null;
        }
        return spec.getToolApprovalRequests().get(toolName);
    }

    private Map<String, Object> metadata(McpSpec spec, McpSchema.Tool tool) {
        Map<String, Object> metadata = new LinkedHashMap<>();
        if (spec.getMetadata() != null) {
            spec.getMetadata().forEach((key, value) -> {
                if (!isSensitiveMetadataKey(key)) {
                    metadata.put(key, value);
                }
            });
        }
        metadata.put("source", "MCP");
        metadata.put("mcpName", spec.getName());
        metadata.put("mcpToolName", tool.name());
        metadata.put("rawMcpToolName", rawToolName(spec, tool));
        metadata.put("toolDisplayName", toolDisplayName(spec, tool));
        metadata.put("transportType", spec.getTransportType().configValue());
        if (isSkillBound(spec)) {
            metadata.put("skillId", spec.getSkillId());
        }
        return metadata;
    }

    private String rawToolName(McpSpec spec, McpSchema.Tool tool) {
        if (tool != null && tool.meta() != null) {
            Object rawName = tool.meta().get(AliasedMcpClientWrapper.RAW_TOOL_NAME_META_KEY);
            if (rawName != null && !String.valueOf(rawName).isBlank()) {
                return String.valueOf(rawName);
            }
        }
        String toolName = tool == null ? null : tool.name();
        if (spec.getToolAliases() != null && !spec.getToolAliases().isEmpty()) {
            for (Map.Entry<String, String> entry : spec.getToolAliases().entrySet()) {
                if (Objects.equals(entry.getValue(), toolName)) {
                    return entry.getKey();
                }
            }
        }
        return toolName;
    }

    private String toolDisplayName(McpSpec spec, McpSchema.Tool tool) {
        String rawToolName = rawToolName(spec, tool);
        String mcpName = spec == null ? null : spec.getDescription();
        if (mcpName == null || mcpName.isBlank()) {
            mcpName = spec == null ? null : spec.getName();
        }
        if (mcpName == null || mcpName.isBlank()) {
            return rawToolName;
        }
        if (rawToolName == null || rawToolName.isBlank()) {
            return mcpName;
        }
        return mcpName + " - " + rawToolName;
    }

    private boolean isSensitiveMetadataKey(String key) {
        if (key == null || key.isBlank()) {
            return false;
        }
        String normalized = key.toLowerCase(Locale.ROOT).replace("-", "").replace("_", "");
        return normalized.contains("key")
                || normalized.contains("token")
                || normalized.contains("secret")
                || normalized.contains("password")
                || normalized.contains("authorization")
                || normalized.contains("credential");
    }

    private void closeQuietly(List<McpClientWrapper> clients) {
        for (McpClientWrapper client : clients) {
            if (client == null) {
                continue;
            }
            try {
                client.close();
            } catch (Exception ignored) {
            }
        }
    }

    private <T> List<T> emptyToNull(List<T> values) {
        return values == null || values.isEmpty() ? null : values;
    }

    private <K, V> Map<K, V> emptyToNull(Map<K, V> values) {
        return values == null || values.isEmpty() ? null : values;
    }

    private String blankToNull(String value) {
        return value == null || value.isBlank() ? null : value;
    }
}
