/*
 * Copyright 2024-2024 the original author or authors.
 */

package io.modelcontextprotocol.server;

import java.util.List;
import java.util.UUID;

import io.modelcontextprotocol.MockMcpServerTransport;
import io.modelcontextprotocol.MockMcpServerTransportProvider;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.ProtocolVersions;
import org.junit.jupiter.api.Test;

import static org.assertj.core.api.Assertions.assertThat;

/**
 * Tests for MCP server protocol version negotiation and compatibility.
 */
class McpServerProtocolVersionTests {

	private static final McpSchema.Implementation SERVER_INFO = McpSchema.Implementation.builder("test-server", "1.0.0")
		.build();

	private static final McpSchema.Implementation CLIENT_INFO = McpSchema.Implementation.builder("test-client", "1.0.0")
		.build();

	private McpSchema.JSONRPCRequest jsonRpcInitializeRequest(String requestId, String protocolVersion) {
		return new McpSchema.JSONRPCRequest(McpSchema.METHOD_INITIALIZE, requestId,
				McpSchema.InitializeRequest
					.builder(protocolVersion, McpSchema.ClientCapabilities.builder().build(), CLIENT_INFO)
					.build());
	}

	@Test
	void shouldUseLatestVersionByDefault() {
		MockMcpServerTransport serverTransport = new MockMcpServerTransport();
		var transportProvider = new MockMcpServerTransportProvider(serverTransport);
		McpAsyncServer server = McpServer.async(transportProvider).serverInfo(SERVER_INFO).build();

		String requestId = UUID.randomUUID().toString();

		transportProvider.simulateIncomingMessage(jsonRpcInitializeRequest(requestId, ProtocolVersions.MCP_2025_11_25));

		McpSchema.JSONRPCMessage response = serverTransport.getLastSentMessage();
		assertThat(response).isInstanceOf(McpSchema.JSONRPCResponse.class);
		McpSchema.JSONRPCResponse jsonResponse = (McpSchema.JSONRPCResponse) response;
		assertThat(jsonResponse.id()).isEqualTo(requestId);
		assertThat(jsonResponse.result()).isInstanceOf(McpSchema.InitializeResult.class);
		McpSchema.InitializeResult result = (McpSchema.InitializeResult) jsonResponse.result();

		var protocolVersion = transportProvider.protocolVersions().get(transportProvider.protocolVersions().size() - 1);
		assertThat(result.protocolVersion()).isEqualTo(protocolVersion);

		server.closeGracefully().subscribe();
	}

	@Test
	void shouldNegotiateSpecificVersion() {
		String oldVersion = "0.1.0";
		MockMcpServerTransport serverTransport = new MockMcpServerTransport();
		var transportProvider = new MockMcpServerTransportProvider(serverTransport);

		McpAsyncServer server = McpServer.async(transportProvider).serverInfo(SERVER_INFO).build();

		server.setProtocolVersions(List.of(oldVersion, ProtocolVersions.MCP_2025_11_25));

		String requestId = UUID.randomUUID().toString();

		transportProvider.simulateIncomingMessage(jsonRpcInitializeRequest(requestId, oldVersion));

		McpSchema.JSONRPCMessage response = serverTransport.getLastSentMessage();
		assertThat(response).isInstanceOf(McpSchema.JSONRPCResponse.class);
		McpSchema.JSONRPCResponse jsonResponse = (McpSchema.JSONRPCResponse) response;
		assertThat(jsonResponse.id()).isEqualTo(requestId);
		assertThat(jsonResponse.result()).isInstanceOf(McpSchema.InitializeResult.class);
		McpSchema.InitializeResult result = (McpSchema.InitializeResult) jsonResponse.result();
		assertThat(result.protocolVersion()).isEqualTo(oldVersion);

		server.closeGracefully().subscribe();
	}

	@Test
	void shouldSuggestLatestVersionForUnsupportedVersion() {
		String unsupportedVersion = "999.999.999";
		MockMcpServerTransport serverTransport = new MockMcpServerTransport();
		var transportProvider = new MockMcpServerTransportProvider(serverTransport);

		McpAsyncServer server = McpServer.async(transportProvider).serverInfo(SERVER_INFO).build();

		String requestId = UUID.randomUUID().toString();

		transportProvider.simulateIncomingMessage(jsonRpcInitializeRequest(requestId, unsupportedVersion));

		McpSchema.JSONRPCMessage response = serverTransport.getLastSentMessage();
		assertThat(response).isInstanceOf(McpSchema.JSONRPCResponse.class);
		McpSchema.JSONRPCResponse jsonResponse = (McpSchema.JSONRPCResponse) response;
		assertThat(jsonResponse.id()).isEqualTo(requestId);
		assertThat(jsonResponse.result()).isInstanceOf(McpSchema.InitializeResult.class);
		McpSchema.InitializeResult result = (McpSchema.InitializeResult) jsonResponse.result();
		var protocolVersion = transportProvider.protocolVersions().get(transportProvider.protocolVersions().size() - 1);
		assertThat(result.protocolVersion()).isEqualTo(protocolVersion);

		server.closeGracefully().subscribe();
	}

	@Test
	void shouldUseHighestVersionWhenMultipleSupported() {
		String oldVersion = "0.1.0";
		String middleVersion = "0.2.0";
		String latestVersion = ProtocolVersions.MCP_2025_11_25;

		MockMcpServerTransport serverTransport = new MockMcpServerTransport();
		var transportProvider = new MockMcpServerTransportProvider(serverTransport);

		McpAsyncServer server = McpServer.async(transportProvider).serverInfo(SERVER_INFO).build();

		server.setProtocolVersions(List.of(oldVersion, middleVersion, latestVersion));

		String requestId = UUID.randomUUID().toString();
		transportProvider.simulateIncomingMessage(jsonRpcInitializeRequest(requestId, latestVersion));

		McpSchema.JSONRPCMessage response = serverTransport.getLastSentMessage();
		assertThat(response).isInstanceOf(McpSchema.JSONRPCResponse.class);
		McpSchema.JSONRPCResponse jsonResponse = (McpSchema.JSONRPCResponse) response;
		assertThat(jsonResponse.id()).isEqualTo(requestId);
		assertThat(jsonResponse.result()).isInstanceOf(McpSchema.InitializeResult.class);
		McpSchema.InitializeResult result = (McpSchema.InitializeResult) jsonResponse.result();
		assertThat(result.protocolVersion()).isEqualTo(latestVersion);

		server.closeGracefully().subscribe();
	}

}
