package repro;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.io.InputStream;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.file.Files;
import java.util.concurrent.TimeUnit;

import org.apache.catalina.Context;
import org.apache.catalina.startup.Tomcat;
import org.junit.jupiter.api.Test;

import io.modelcontextprotocol.server.McpAsyncServer;
import io.modelcontextprotocol.server.McpServer;
import io.modelcontextprotocol.server.transport.HttpServletStreamableServerTransportProvider;

class StreamableGetStreamHeadersTest {

	@Test
	void getStreamSendsHeadersBeforeFirstEvent() throws Exception {
		HttpServletStreamableServerTransportProvider transport = HttpServletStreamableServerTransportProvider.builder()
			.build();
		McpAsyncServer server = McpServer.async(transport).serverInfo("repro", "1").build();

		Tomcat tomcat = new Tomcat();
		tomcat.setBaseDir(Files.createTempDirectory("tomcat").toString());
		tomcat.setPort(0);
		Context context = tomcat.addContext("", null);
		Tomcat.addServlet(context, "mcp", transport).setAsyncSupported(true);
		context.addServletMappingDecoded("/*", "mcp");
		tomcat.start();

		try {
			URI endpoint = URI.create("http://127.0.0.1:" + tomcat.getConnector().getLocalPort() + "/mcp");
			HttpClient client = HttpClient.newHttpClient();

			String sessionId = client.send(post(endpoint, null, """
					{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25",\
					"capabilities":{},"clientInfo":{"name":"repro","version":"1"}}}"""),
					HttpResponse.BodyHandlers.discarding())
				.headers()
				.firstValue("Mcp-Session-Id")
				.orElseThrow();
			assertEquals(202, client.send(post(endpoint, sessionId, """
					{"jsonrpc":"2.0","method":"notifications/initialized"}"""),
					HttpResponse.BodyHandlers.discarding())
				.statusCode());

			HttpRequest get = HttpRequest.newBuilder(endpoint)
				.header("Accept", "text/event-stream")
				.header("Mcp-Session-Id", sessionId)
				.GET()
				.build();

			// Completes once the status line and headers arrive. Fails with TimeoutException.
			HttpResponse<InputStream> response = client.sendAsync(get, HttpResponse.BodyHandlers.ofInputStream())
				.get(5, TimeUnit.SECONDS);

			assertEquals(200, response.statusCode());
			assertTrue(response.headers().firstValue("Content-Type").orElse("").startsWith("text/event-stream"));
			response.body().close();
		}
		finally {
			server.closeGracefully().block();
			tomcat.stop();
			tomcat.destroy();
		}
	}

	private static HttpRequest post(URI endpoint, String sessionId, String body) {
		HttpRequest.Builder builder = HttpRequest.newBuilder(endpoint)
			.header("Accept", "application/json, text/event-stream")
			.header("Content-Type", "application/json")
			.POST(HttpRequest.BodyPublishers.ofString(body));
		if (sessionId != null) {
			builder.header("Mcp-Session-Id", sessionId);
		}
		return builder.build();
	}

}