package com.easyagents.mcp.client;

import io.modelcontextprotocol.json.TypeRef;
import io.modelcontextprotocol.spec.McpClientTransport;
import io.modelcontextprotocol.spec.McpSchema;
import org.junit.Test;
import reactor.core.publisher.Mono;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;

/**
 * MCP 环境检测测试。
 */
public class McpEnvironmentCheckerTest {

    @Test
    public void checkValidStdioConfigWithoutProbe() {
        String json = """
            {
              "mcpServers": {
                "test": {
                  "transport": "stdio",
                  "command": "java",
                  "args": ["-version"]
                }
              }
            }
            """;

        McpEnvironmentCheckResult result = new McpEnvironmentChecker(false).check(json);

        assertEquals(McpCheckStatus.SUCCESS, result.getOverallStatus());
        assertEquals("test", result.getServers().get(0).getServerName());
    }

    @Test
    public void checkMissingCommand() {
        String json = """
            {
              "mcpServers": {
                "test": {
                  "transport": "stdio"
                }
              }
            }
            """;

        McpEnvironmentCheckResult result = new McpEnvironmentChecker(false).check(json);

        assertEquals(McpCheckStatus.FAILED, result.getOverallStatus());
    }

    @Test
    public void checkMissingHttpUrl() {
        String json = """
            {
              "mcpServers": {
                "test": {
                  "transport": "http-sse"
                }
              }
            }
            """;

        McpEnvironmentCheckResult result = new McpEnvironmentChecker(false).check(json);

        assertEquals(McpCheckStatus.FAILED, result.getOverallStatus());
    }

    @Test
    public void checkUnresolvedInputEnv() {
        String json = """
            {
              "mcpServers": {
                "test": {
                  "transport": "stdio",
                  "command": "java",
                  "env": {
                    "API_KEY": "${input:api_key}"
                  }
                }
              }
            }
            """;

        McpEnvironmentCheckResult result = new McpEnvironmentChecker(false).check(json);

        assertEquals(McpCheckStatus.FAILED, result.getOverallStatus());
        assertFalse(result.getServers().get(0).getChecks().toString().contains("secret"));
    }

    @Test
    public void checkInvalidJson() {
        McpEnvironmentCheckResult result = new McpEnvironmentChecker(false).check("{ invalid json }");

        assertEquals(McpCheckStatus.FAILED, result.getOverallStatus());
    }

    @Test
    public void closeTransportWhenProbeFailed() {
        TrackingCloseableTransport closeableTransport = new TrackingCloseableTransport();
        McpEnvironmentChecker checker = new McpEnvironmentChecker(true, transport -> (spec, resolvedEnv) -> closeableTransport);
        String json = """
            {
              "mcpServers": {
                "test": {
                  "transport": "stdio",
                  "command": "java"
                }
              }
            }
            """;

        McpEnvironmentCheckResult result = checker.check(json);

        assertEquals(McpCheckStatus.FAILED, result.getOverallStatus());
        assertTrue(closeableTransport.closed);
    }

    private static class TrackingCloseableTransport implements CloseableTransport {
        private boolean closed;

        @Override
        public McpClientTransport getTransport() {
            return new FailingClientTransport();
        }

        @Override
        public void close() {
            closed = true;
        }
    }

    private static class FailingClientTransport implements McpClientTransport {
        @Override
        public Mono<Void> connect(java.util.function.Function<Mono<McpSchema.JSONRPCMessage>,
            Mono<McpSchema.JSONRPCMessage>> handler) {
            return Mono.error(new IllegalStateException("probe failed"));
        }

        @Override
        public Mono<Void> closeGracefully() {
            return Mono.empty();
        }

        @Override
        public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
            return Mono.empty();
        }

        @Override
        public <T> T unmarshalFrom(Object data, TypeRef<T> typeRef) {
            return null;
        }
    }
}
