package com.easyagents.agent.runtime.mcp;

import com.easyagents.agent.runtime.AgentRuntimeException;
import com.easyagents.agent.runtime.agentscope.AgentScopeSkillAdapter;
import com.easyagents.agent.runtime.skill.AgentSkillBoxSpec;
import com.easyagents.agent.runtime.skill.AgentSkillSpec;
import com.easyagents.agent.runtime.tool.AgentToolSpec;
import io.agentscope.core.message.ToolResultBlock;
import io.agentscope.core.tool.Toolkit;
import io.agentscope.core.tool.mcp.McpClientWrapper;
import io.modelcontextprotocol.spec.McpSchema;
import org.junit.Assert;
import org.junit.Test;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 测试 MCP Toolkit 适配器。
 */
public class McpToolkitAdapterTest {

    @Test
    public void shouldRegisterEnabledMcpToolsAndBuildRuntimeToolSpecs() {
        FakeMcpClientWrapper client = new FakeMcpClientWrapper("demo",
                List.of(tool("search"), tool("write_file")));
        McpToolkitAdapter adapter = new McpToolkitAdapter(new FakeMcpClientFactory(client));
        McpSpec spec = stdioSpec();
        spec.setEnableTools(List.of("search"));
        spec.setApprovalRequired(true);
        spec.setMetadata(Map.of("owner", "runtime"));
        Toolkit toolkit = new Toolkit();

        McpRegistration registration = adapter.register(List.of(spec), toolkit);

        Assert.assertEquals(1, registration.getClients().size());
        Assert.assertEquals(1, registration.getToolSpecs().size());
        Assert.assertNotNull(toolkit.getTool("search"));
        Assert.assertNull(toolkit.getTool("write_file"));
        AgentToolSpec toolSpec = registration.getToolSpecs().get(0);
        Assert.assertEquals("search", toolSpec.getName());
        Assert.assertTrue(toolSpec.isApprovalRequired());
        Assert.assertEquals("MCP", toolSpec.getMetadata().get("source"));
        Assert.assertEquals("demo", toolSpec.getMetadata().get("mcpName"));
        Assert.assertEquals("stdio", toolSpec.getMetadata().get("transportType"));
        Assert.assertEquals("runtime", toolSpec.getMetadata().get("owner"));
    }

    @Test
    public void shouldFilterSensitiveMetadataKeysFromRuntimeToolSpec() {
        FakeMcpClientWrapper client = new FakeMcpClientWrapper("demo", List.of(tool("search")));
        McpToolkitAdapter adapter = new McpToolkitAdapter(new FakeMcpClientFactory(client));
        McpSpec spec = stdioSpec();
        spec.setMetadata(Map.of(
                "owner", "runtime",
                "apiKey", "secret-key",
                "Authorization", "Bearer secret-token",
                "password", "secret-password"));

        McpRegistration registration = adapter.register(List.of(spec), new Toolkit());

        Map<String, Object> metadata = registration.getToolSpecs().get(0).getMetadata();
        Assert.assertEquals("runtime", metadata.get("owner"));
        Assert.assertFalse(metadata.containsKey("apiKey"));
        Assert.assertFalse(metadata.containsKey("Authorization"));
        Assert.assertFalse(metadata.containsKey("password"));
    }

    @Test
    public void shouldRegisterAliasedMcpToolAndCallRawToolName() {
        FakeMcpClientWrapper client = new FakeMcpClientWrapper("demo", List.of(tool("search")));
        McpToolkitAdapter adapter = new McpToolkitAdapter(new FakeMcpClientFactory(client));
        McpSpec spec = stdioSpec();
        spec.setEnableTools(List.of("mcp_1_search"));
        spec.setToolAliases(Map.of("search", "mcp_1_search"));
        Toolkit toolkit = new Toolkit();

        McpRegistration registration = adapter.register(List.of(spec), toolkit);

        Assert.assertNotNull(toolkit.getTool("mcp_1_search"));
        Assert.assertNull(toolkit.getTool("search"));
        AgentToolSpec toolSpec = registration.getToolSpecs().get(0);
        Assert.assertEquals("mcp_1_search", toolSpec.getName());
        Assert.assertEquals("search", toolSpec.getMetadata().get("rawMcpToolName"));
        registration.getClients().get(0).callTool("mcp_1_search", Map.of("q", "hello")).block();
        Assert.assertEquals("search", client.lastCalledToolName.get());
    }

    /**
     * 验证未配置工具白名单时会为 MCP 下全部工具动态生成运行时别名。
     */
    @Test
    public void shouldRegisterAllToolsWithDynamicPrefixAndCallRawToolName() {
        FakeMcpClientWrapper client = new FakeMcpClientWrapper("demo",
                List.of(tool("search.tool"), tool("write-file")));
        McpToolkitAdapter adapter = new McpToolkitAdapter(new FakeMcpClientFactory(client));
        McpSpec spec = stdioSpec();
        spec.setDescription("Demo MCP");
        spec.setToolNamePrefix("mcp_20_");
        Toolkit toolkit = new Toolkit();

        McpRegistration registration = adapter.register(List.of(spec), toolkit);

        Assert.assertNotNull(toolkit.getTool("mcp_20_search_tool"));
        Assert.assertNotNull(toolkit.getTool("mcp_20_write-file"));
        Assert.assertNull(toolkit.getTool("search.tool"));
        Assert.assertEquals(2, registration.getToolSpecs().size());
        AgentToolSpec toolSpec = registration.getToolSpecs().get(0);
        Assert.assertEquals("mcp_20_search_tool", toolSpec.getName());
        Assert.assertEquals("search.tool", toolSpec.getMetadata().get("rawMcpToolName"));
        Assert.assertEquals("Demo MCP - search.tool", toolSpec.getMetadata().get("toolDisplayName"));
        registration.getClients().get(0).callTool("mcp_20_search_tool", Map.of("q", "hello")).block();
        Assert.assertEquals("search.tool", client.lastCalledToolName.get());
    }

    /**
     * 验证动态别名安全化后发生重名时会自动追加序号。
     */
    @Test
    public void shouldDeduplicateDynamicAliasesAfterSanitizingToolNames() {
        FakeMcpClientWrapper client = new FakeMcpClientWrapper("demo",
                List.of(tool("search.tool"), tool("search_tool")));
        McpToolkitAdapter adapter = new McpToolkitAdapter(new FakeMcpClientFactory(client));
        McpSpec spec = stdioSpec();
        spec.setToolNamePrefix("mcp_20_");
        Toolkit toolkit = new Toolkit();

        McpRegistration registration = adapter.register(List.of(spec), toolkit);

        Assert.assertNotNull(toolkit.getTool("mcp_20_search_tool"));
        Assert.assertNotNull(toolkit.getTool("mcp_20_search_tool_2"));
        Assert.assertEquals("search.tool", registration.getToolSpecs().get(0).getMetadata().get("rawMcpToolName"));
        Assert.assertEquals("search_tool", registration.getToolSpecs().get(1).getMetadata().get("rawMcpToolName"));
        registration.getClients().get(0).callTool("mcp_20_search_tool_2", Map.of("q", "hello")).block();
        Assert.assertEquals("search_tool", client.lastCalledToolName.get());
    }

    @Test
    public void shouldReturnEmptyToolSpecsWhenMcpServerHasNoTools() {
        FakeMcpClientWrapper client = new FakeMcpClientWrapper("demo", List.of());
        McpToolkitAdapter adapter = new McpToolkitAdapter(new FakeMcpClientFactory(client));

        McpRegistration registration = adapter.register(List.of(stdioSpec()), new Toolkit());

        Assert.assertEquals(1, registration.getClients().size());
        Assert.assertTrue(registration.getToolSpecs().isEmpty());
    }

    @Test
    public void shouldCloseCreatedClientWhenRegistrationFails() {
        FakeMcpClientWrapper client = new FakeMcpClientWrapper("demo", List.of(tool("search")));
        client.failOnListTools = true;
        McpToolkitAdapter adapter = new McpToolkitAdapter(new FakeMcpClientFactory(client));

        try {
            adapter.register(List.of(stdioSpec()), new Toolkit());
            Assert.fail("Expected MCP registration failure.");
        } catch (AgentRuntimeException | IllegalStateException ignored) {
            Assert.assertTrue(client.closed.get());
        }
    }

    /**
     * 验证 Skill MCP 只注册到禁用的 Skill Tool Group，加载前不向模型暴露。
     */
    @Test
    public void shouldRegisterSkillMcpAsInactiveSkillToolGroup() {
        List<McpSchema.Tool> frozenTools = List.of(tool("search"));
        FakeMcpClientWrapper client = new FakeMcpClientWrapper("demo",
                List.of(tool("search"), tool("new_remote_tool")));
        McpToolkitAdapter adapter = new McpToolkitAdapter(new FakeMcpClientFactory(client));
        McpSpec spec = stdioSpec();
        spec.setSkillId("skill-1");
        spec.setToolAliases(Map.of("search", "skill_1_mcp_search"));
        spec.setFrozenToolManifest(McpToolManifest.fromTools(frozenTools));
        spec.setFrozenToolManifestHash(McpToolManifest.hash(spec.getFrozenToolManifest()));
        Toolkit toolkit = new Toolkit();

        McpRegistration registration = adapter.register(List.of(spec), toolkit);

        Assert.assertNull(toolkit.getTool("skill_1_mcp_search"));
        Assert.assertEquals(1, registration.getSkillRegistrations().size());
        Assert.assertEquals(List.of("skill_1_mcp_search"),
                registration.getSkillRegistrations().get(0).getEnableTools());
        AgentSkillSpec skill = new AgentSkillSpec();
        skill.setSkillId("skill-1");
        skill.setName("Search Skill");
        skill.setDescription("Search through MCP.");
        skill.setSkillContent("Load this skill before searching.");
        AgentSkillBoxSpec skillBoxSpec = new AgentSkillBoxSpec();
        skillBoxSpec.setSkills(List.of(skill));

        new AgentScopeSkillAdapter().createSkillBox(skillBoxSpec, toolkit, Map.of(),
                registration.getSkillRegistrations());

        Assert.assertNotNull(toolkit.getTool("skill_1_mcp_search"));
        Assert.assertFalse(toolkit.getActiveGroups().contains("skill-1_skill_tools"));
        Assert.assertTrue(toolkit.getToolSchemas().stream()
                .noneMatch(schema -> "skill_1_mcp_search".equals(schema.getName())));
        Assert.assertNull(toolkit.getTool("skill_1_mcp_new_remote_tool"));
        Assert.assertEquals(1, client.remoteListCalls.get());
    }

    /**
     * 验证 Skill MCP 会在读取 Tool 清单前完成异步初始化。
     */
    @Test
    public void shouldInitializeSkillMcpBeforeListingTools() {
        FakeMcpClientWrapper client = new FakeMcpClientWrapper("demo", List.of(tool("search")));
        client.deferInitialization = true;
        McpToolkitAdapter adapter = new McpToolkitAdapter(new FakeMcpClientFactory(client));
        McpSpec spec = stdioSpec();
        spec.setSkillId("skill-1");
        spec.setFrozenToolManifest(McpToolManifest.fromTools(List.of(tool("search"))));
        spec.setFrozenToolManifestHash(McpToolManifest.hash(spec.getFrozenToolManifest()));

        McpRegistration registration = adapter.register(List.of(spec), new Toolkit());

        Assert.assertTrue(client.isInitialized());
        Assert.assertEquals(1, client.remoteListCalls.get());
        Assert.assertEquals(1, registration.getSkillRegistrations().size());
    }

    /**
     * 验证冻结 Tool 缺失时拒绝注册并关闭 client。
     */
    @Test
    public void shouldRejectMissingFrozenSkillMcpToolAndCloseClient() {
        FakeMcpClientWrapper client = new FakeMcpClientWrapper("demo", List.of(tool("other")));
        McpToolkitAdapter adapter = new McpToolkitAdapter(new FakeMcpClientFactory(client));
        McpSpec spec = stdioSpec();
        spec.setSkillId("skill-1");
        spec.setFrozenToolManifest(McpToolManifest.fromTools(List.of(tool("search"))));
        spec.setFrozenToolManifestHash(McpToolManifest.hash(spec.getFrozenToolManifest()));

        try {
            adapter.register(List.of(spec), new Toolkit());
            Assert.fail("Expected frozen MCP tool validation failure.");
        } catch (AgentRuntimeException expected) {
            Assert.assertTrue(expected.getMessage().contains("missing"));
            Assert.assertTrue(client.closed.get());
        }
    }

    /**
     * 验证冻结 Tool Schema 漂移时拒绝注册。
     */
    @Test(expected = AgentRuntimeException.class)
    public void shouldRejectChangedFrozenSkillMcpSchema() {
        McpSchema.Tool expectedTool = tool("search");
        McpSchema.JsonSchema changedSchema = new McpSchema.JsonSchema("object",
                Map.of("keyword", Map.of("type", "string")), List.of("keyword"), null, null, null);
        McpSchema.Tool actualTool = new McpSchema.Tool("search", "search", "search description",
                changedSchema, null, null, null);
        McpToolkitAdapter adapter = new McpToolkitAdapter(
                new FakeMcpClientFactory(new FakeMcpClientWrapper("demo", List.of(actualTool))));
        McpSpec spec = stdioSpec();
        spec.setSkillId("skill-1");
        spec.setFrozenToolManifest(McpToolManifest.fromTools(List.of(expectedTool)));
        spec.setFrozenToolManifestHash(McpToolManifest.hash(spec.getFrozenToolManifest()));

        adapter.register(List.of(spec), new Toolkit());
    }

    /**
     * 验证远端仅调整 Tool 描述时不破坏已发布 Skill 的运行兼容性。
     */
    @Test
    public void shouldAllowChangedDescriptionWhenFrozenSchemaIsStable() {
        McpSchema.Tool expectedTool = tool("search");
        McpSchema.Tool actualTool = new McpSchema.Tool(
                "search", "search", "updated description",
                expectedTool.inputSchema(), expectedTool.outputSchema(), null, null);
        McpToolkitAdapter adapter = new McpToolkitAdapter(
                new FakeMcpClientFactory(new FakeMcpClientWrapper("demo", List.of(actualTool))));
        McpSpec spec = stdioSpec();
        spec.setSkillId("skill-1");
        spec.setFrozenToolManifest(McpToolManifest.fromTools(List.of(expectedTool)));
        spec.setFrozenToolManifestHash(McpToolManifest.hash(spec.getFrozenToolManifest()));

        McpRegistration registration = adapter.register(List.of(spec), new Toolkit());

        Assert.assertEquals(1, registration.getSkillRegistrations().size());
        Assert.assertEquals("search description", registration.getToolSpecs().get(0).getDescription());
    }

    private McpSpec stdioSpec() {
        McpSpec spec = new McpSpec();
        spec.setName("demo");
        spec.setTransportType(McpTransportType.STDIO);
        spec.setCommand("node");
        spec.setTimeout(Duration.ofSeconds(10));
        spec.setInitializationTimeout(Duration.ofSeconds(3));
        return spec;
    }

    private McpSchema.Tool tool(String name) {
        McpSchema.JsonSchema schema = new McpSchema.JsonSchema("object",
                Map.of("q", Map.of("type", "string", "description", "query")),
                List.of("q"), null, null, null);
        return new McpSchema.Tool(name, name, name + " description", schema, null, null, null);
    }

    private static class FakeMcpClientFactory extends McpClientFactory {

        private final McpClientWrapper client;

        private FakeMcpClientFactory(McpClientWrapper client) {
            this.client = client;
        }

        @Override
        public McpClientWrapper create(McpSpec spec) {
            return client;
        }
    }

    private static class FakeMcpClientWrapper extends McpClientWrapper {

        private final List<McpSchema.Tool> tools;
        private final AtomicBoolean closed = new AtomicBoolean(false);
        private final AtomicReference<String> lastCalledToolName = new AtomicReference<>();
        private final java.util.concurrent.atomic.AtomicInteger remoteListCalls =
                new java.util.concurrent.atomic.AtomicInteger();
        private boolean failOnListTools;
        private boolean deferInitialization;

        private FakeMcpClientWrapper(String name, List<McpSchema.Tool> tools) {
            super(name);
            this.tools = tools;
        }

        @Override
        public Mono<Void> initialize() {
            if (deferInitialization) {
                return Mono.fromRunnable(() -> initialized = true);
            }
            initialized = true;
            return Mono.empty();
        }

        @Override
        public Mono<List<McpSchema.Tool>> listTools() {
            if (!initialized) {
                return Mono.error(new IllegalStateException("client is not initialized"));
            }
            remoteListCalls.incrementAndGet();
            if (failOnListTools) {
                return Mono.error(new IllegalStateException("list tools failed"));
            }
            return Mono.just(tools);
        }

        @Override
        public Mono<McpSchema.CallToolResult> callTool(String toolName, Map<String, Object> arguments) {
            lastCalledToolName.set(toolName);
            return Mono.just(new McpSchema.CallToolResult(List.of(), false));
        }

        @Override
        public void close() {
            closed.set(true);
            initialized = false;
        }
    }
}
