package com.zy.ai.mcp.config;

import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.json.jackson.JacksonMcpJsonMapper;
import io.modelcontextprotocol.server.McpServer;
import io.modelcontextprotocol.server.McpServerFeatures;
import io.modelcontextprotocol.server.McpSyncServer;
import io.modelcontextprotocol.server.transport.WebMvcSseServerTransportProvider;
import io.modelcontextprotocol.spec.McpSchema;
import org.springframework.ai.mcp.server.common.autoconfigure.properties.McpServerChangeNotificationProperties;
import org.springframework.ai.mcp.server.common.autoconfigure.properties.McpServerProperties;
import org.springframework.ai.mcp.server.common.autoconfigure.properties.McpServerSseProperties;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.env.Environment;
import org.springframework.web.context.support.StandardServletEnvironment;
import org.springframework.web.servlet.function.RouterFunction;
import org.springframework.web.servlet.function.ServerResponse;

import java.util.Collections;
import java.util.List;

@Configuration
@EnableConfigurationProperties(McpServerSseProperties.class)
public class SpringAiMcpTransportConfig {

    @Bean("wcsOfficialSseMcpSupport")
    public OfficialSseMcpSupport wcsOfficialSseMcpSupport(
            @Qualifier("mcpServerObjectMapper") ObjectMapper objectMapper,
            McpServerSseProperties sseProperties,
            McpServerProperties serverProperties,
            McpServerChangeNotificationProperties changeNotificationProperties,
            @Qualifier("syncTools") ObjectProvider<List<McpServerFeatures.SyncToolSpecification>> syncToolsProvider,
            Environment environment) {

        WebMvcSseServerTransportProvider transportProvider = WebMvcSseServerTransportProvider.builder()
                .jsonMapper(new JacksonMcpJsonMapper(objectMapper))
                .baseUrl(sseProperties.getBaseUrl())
                .sseEndpoint(sseProperties.getSseEndpoint())
                .messageEndpoint(sseProperties.getSseMessageEndpoint())
                .keepAliveInterval(sseProperties.getKeepAliveInterval())
                .build();

        List<McpServerFeatures.SyncToolSpecification> syncTools = syncToolsProvider.getIfAvailable(Collections::emptyList);
        McpSyncServer mcpSyncServer = buildSseSyncServer(
                transportProvider,
                serverProperties,
                changeNotificationProperties,
                syncTools,
                environment
        );

        return new OfficialSseMcpSupport(transportProvider, mcpSyncServer);
    }

    @Bean("webMvcSseServerRouterFunction")
    public RouterFunction<ServerResponse> webMvcSseServerRouterFunction(
            @Qualifier("wcsOfficialSseMcpSupport") OfficialSseMcpSupport support) {
        return support.routerFunction();
    }

    private McpSyncServer buildSseSyncServer(
            WebMvcSseServerTransportProvider transportProvider,
            McpServerProperties serverProperties,
            McpServerChangeNotificationProperties changeNotificationProperties,
            List<McpServerFeatures.SyncToolSpecification> syncTools,
            Environment environment) {

        McpServer.SingleSessionSyncSpecification specification = McpServer.sync(transportProvider);
        specification.serverInfo(new McpSchema.Implementation(serverProperties.getName(), serverProperties.getVersion()));

        McpSchema.ServerCapabilities.Builder capabilitiesBuilder = McpSchema.ServerCapabilities.builder();
        if (serverProperties.getCapabilities().isTool()) {
            capabilitiesBuilder.tools(changeNotificationProperties.isToolChangeNotification());
            if (syncTools != null && !syncTools.isEmpty()) {
                specification.tools(syncTools);
            }
        }

        specification.capabilities(capabilitiesBuilder.build());
        specification.instructions(serverProperties.getInstructions());
        specification.requestTimeout(serverProperties.getRequestTimeout());
        if (environment instanceof StandardServletEnvironment) {
            specification.immediateExecution(true);
        }

        return specification.build();
    }

    public static final class OfficialSseMcpSupport implements AutoCloseable {

        private final WebMvcSseServerTransportProvider transportProvider;
        private final McpSyncServer mcpSyncServer;

        public OfficialSseMcpSupport(WebMvcSseServerTransportProvider transportProvider, McpSyncServer mcpSyncServer) {
            this.transportProvider = transportProvider;
            this.mcpSyncServer = mcpSyncServer;
        }

        public RouterFunction<ServerResponse> routerFunction() {
            return transportProvider.getRouterFunction();
        }

        @Override
        public void close() {
            mcpSyncServer.closeGracefully();
            mcpSyncServer.close();
        }
    }
}