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

package io.modelcontextprotocol.client.transport;

import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.net.http.HttpResponse.BodyHandler;
import java.time.Duration;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletionException;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import java.util.function.Function;

import io.modelcontextprotocol.client.McpAsyncClient;
import io.modelcontextprotocol.client.transport.ResponseSubscribers.ResponseEvent;
import io.modelcontextprotocol.client.transport.customizer.McpAsyncHttpClientRequestCustomizer;
import io.modelcontextprotocol.client.transport.customizer.McpHttpClientAuthorizationErrorHandler;
import io.modelcontextprotocol.client.transport.customizer.McpHttpClientTransportAuthorizationErrorHandler;
import io.modelcontextprotocol.client.transport.customizer.McpSyncHttpClientRequestCustomizer;
import io.modelcontextprotocol.common.McpTransportContext;
import io.modelcontextprotocol.json.McpJsonDefaults;
import io.modelcontextprotocol.json.McpJsonMapper;
import io.modelcontextprotocol.json.TypeRef;
import io.modelcontextprotocol.spec.ClosedMcpTransportSession;
import io.modelcontextprotocol.spec.DefaultMcpTransportSession;
import io.modelcontextprotocol.spec.DefaultMcpTransportStream;
import io.modelcontextprotocol.spec.HttpHeaders;
import io.modelcontextprotocol.spec.McpClientTransport;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpTransportException;
import io.modelcontextprotocol.spec.McpTransportSession;
import io.modelcontextprotocol.spec.McpTransportSessionClosedException;
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
import io.modelcontextprotocol.spec.McpTransportStream;
import io.modelcontextprotocol.spec.ProtocolVersions;
import io.modelcontextprotocol.util.Assert;
import io.modelcontextprotocol.util.Utils;
import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;
import reactor.core.publisher.Mono;
import reactor.util.function.Tuple2;
import reactor.util.function.Tuples;
import reactor.util.retry.Retry;

/**
 * An implementation of the Streamable HTTP protocol as defined by the
 * <code>2025-03-26</code> version of the MCP specification.
 *
 * <p>
 * The transport is capable of resumability and reconnects. It reacts to transport-level
 * session invalidation and will propagate {@link McpTransportSessionNotFoundException
 * appropriate exceptions} to the higher level abstraction layer when needed in order to
 * allow proper state management. The implementation handles servers that are stateful and
 * provide session meta information, but can also communicate with stateless servers that
 * do not provide a session identifier and do not support SSE streams.
 * </p>
 * <p>
 * This implementation does not handle backwards compatibility with the <a href=
 * "https://modelcontextprotocol.io/specification/2024-11-05/basic/transports#http-with-sse">"HTTP
 * with SSE" transport</a>. In order to communicate over the phased-out
 * <code>2024-11-05</code> protocol, use {@link HttpClientSseClientTransport} or
 * {@code WebFluxSseClientTransport}.
 * </p>
 *
 * @author Christian Tzolov
 * @author Daniel Garnier-Moiroux
 * @see <a href=
 * "https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#streamable-http">Streamable
 * HTTP transport specification</a>
 */
public class HttpClientStreamableHttpTransport implements McpClientTransport {

	private static final Logger logger = LoggerFactory.getLogger(HttpClientStreamableHttpTransport.class);

	private static final String DEFAULT_ENDPOINT = "/mcp";

	/**
	 * HTTP client for sending messages to the server. Uses HTTP POST over the message
	 * endpoint
	 */
	private final HttpClient httpClient;

	/** HTTP request builder for building requests to send messages to the server */
	private final HttpRequest.Builder requestBuilder;

	/**
	 * Event type for JSON-RPC messages received through the SSE connection. The server
	 * sends messages with this event type to transmit JSON-RPC protocol data.
	 */
	private static final String MESSAGE_EVENT_TYPE = "message";

	private static final String APPLICATION_JSON = "application/json";

	private static final String APPLICATION_JSON_UTF8 = "application/json; charset=utf-8";

	private static final String TEXT_EVENT_STREAM = "text/event-stream";

	public static int NOT_FOUND = 404;

	public static int METHOD_NOT_ALLOWED = 405;

	public static int BAD_REQUEST = 400;

	private final McpJsonMapper jsonMapper;

	private final URI baseUri;

	private final String endpoint;

	private final boolean openConnectionOnStartup;

	private final McpHttpClientTransportAuthorizationErrorHandler authorizationErrorHandler;

	private final boolean resumableStreams;

	private final McpAsyncHttpClientRequestCustomizer httpRequestCustomizer;

	private final AtomicReference<McpTransportSession<Disposable>> activeSession = new AtomicReference<>();

	private final AtomicReference<Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>>> handler = new AtomicReference<>();

	private final AtomicReference<Consumer<Throwable>> exceptionHandler = new AtomicReference<>();

	private final List<String> supportedProtocolVersions;

	private final String latestSupportedProtocolVersion;

	private HttpClientStreamableHttpTransport(McpJsonMapper jsonMapper, HttpClient httpClient,
			HttpRequest.Builder requestBuilder, String baseUri, String endpoint, boolean resumableStreams,
			boolean openConnectionOnStartup, McpAsyncHttpClientRequestCustomizer httpRequestCustomizer,
			McpHttpClientTransportAuthorizationErrorHandler authorizationErrorHandler,
			List<String> supportedProtocolVersions) {
		this.jsonMapper = jsonMapper;
		this.httpClient = httpClient;
		this.requestBuilder = requestBuilder;
		this.baseUri = URI.create(baseUri);
		this.endpoint = endpoint;
		this.resumableStreams = resumableStreams;
		this.openConnectionOnStartup = openConnectionOnStartup;
		this.authorizationErrorHandler = authorizationErrorHandler;
		this.activeSession.set(createTransportSession());
		this.httpRequestCustomizer = httpRequestCustomizer;
		this.supportedProtocolVersions = Collections.unmodifiableList(supportedProtocolVersions);
		this.latestSupportedProtocolVersion = this.supportedProtocolVersions.stream()
			.sorted(Comparator.reverseOrder())
			.findFirst()
			.get();
	}

	@Override
	public List<String> protocolVersions() {
		return supportedProtocolVersions;
	}

	public static Builder builder(String baseUri) {
		return new Builder(baseUri);
	}

	@Override
	public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
		return Mono.deferContextual(ctx -> {
			this.handler.set(handler);
			if (this.openConnectionOnStartup) {
				logger.debug("Eagerly opening connection on startup");
				return this.reconnect(null).onErrorComplete(t -> {
					logger.warn("Eager connect failed ", t);
					return true;
				}).then();
			}
			return Mono.empty();
		});
	}

	private McpTransportSession<Disposable> createTransportSession() {
		Function<String, Publisher<Void>> onClose = sessionId -> sessionId == null ? Mono.empty()
				: createDelete(sessionId);
		return new DefaultMcpTransportSession(onClose);
	}

	private Publisher<Void> createDelete(String sessionId) {

		var uri = Utils.resolveUri(this.baseUri, this.endpoint);
		return Mono.deferContextual(ctx -> {
			var builder = this.requestBuilder.copy()
				.uri(uri)
				.header("Cache-Control", "no-cache")
				.header(HttpHeaders.MCP_SESSION_ID, sessionId)
				.header(HttpHeaders.PROTOCOL_VERSION,
						ctx.getOrDefault(McpAsyncClient.NEGOTIATED_PROTOCOL_VERSION,
								this.latestSupportedProtocolVersion))
				.DELETE();
			var transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY);
			return Mono.from(this.httpRequestCustomizer.customize(builder, "DELETE", uri, null, transportContext));
		}).flatMap(requestBuilder -> {
			var request = requestBuilder.build();
			return Mono.fromFuture(() -> this.httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofString()));
		}).then();
	}

	@Override
	public void setExceptionHandler(Consumer<Throwable> handler) {
		logger.debug("Exception handler registered");
		this.exceptionHandler.set(handler);
	}

	private void handleException(Throwable t) {
		logger.debug("Handling exception for session {}", sessionIdOrPlaceholder(this.activeSession.get()), t);
		if (t instanceof McpTransportSessionNotFoundException) {
			McpTransportSession<?> invalidSession = this.activeSession.getAndSet(createTransportSession());
			logger.warn("Server does not recognize session {}. Invalidating.", invalidSession.sessionId());
			invalidSession.close();
		}
		Consumer<Throwable> handler = this.exceptionHandler.get();
		if (handler != null) {
			handler.accept(t);
		}
	}

	@Override
	public Mono<Void> closeGracefully() {
		return Mono.defer(() -> {
			logger.debug("Graceful close triggered");
			McpTransportSession<Disposable> currentSession = this.activeSession
				.getAndSet(ClosedMcpTransportSession.INSTANCE);
			if (currentSession != null) {
				return Mono.from(currentSession.closeGracefully());
			}
			return Mono.empty();
		});
	}

	private Mono<Disposable> reconnect(McpTransportStream<Disposable> stream) {
		return Mono.deferContextual(ctx -> {
			var rh = this.handler.get();
			if (rh == null) {
				logger.warn("Transport has no request handler registered. Remember to call connect!");
			}

			final Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> requestHandler = rh != null
					? rh : msg -> Mono.error(new IllegalStateException("No request handler"));

			final McpTransportSession<Disposable> transportSession = this.activeSession.get();

			if (ClosedMcpTransportSession.INSTANCE.equals(transportSession)) {
				throw new McpTransportSessionClosedException();
			}

			if (stream != null) {
				logger.debug("Reconnecting stream {} with lastId {}", stream.streamId(), stream.lastId());
			}
			else {
				logger.debug("Reconnecting with no prior stream");
			}

			final AtomicReference<Disposable> disposableRef = new AtomicReference<>();

			var uri = Utils.resolveUri(this.baseUri, this.endpoint);

			Disposable connection = Mono.deferContextual(connectionCtx -> {
				HttpRequest.Builder requestBuilder = this.requestBuilder.copy();

				if (transportSession != null && transportSession.sessionId().isPresent()) {
					requestBuilder = requestBuilder.header(HttpHeaders.MCP_SESSION_ID,
							transportSession.sessionId().get());
				}

				if (stream != null && stream.lastId().isPresent()) {
					requestBuilder = requestBuilder.header(HttpHeaders.LAST_EVENT_ID, stream.lastId().get());
				}

				var builder = requestBuilder.uri(uri)
					.header(HttpHeaders.ACCEPT, TEXT_EVENT_STREAM)
					.header("Cache-Control", "no-cache")
					.header(HttpHeaders.PROTOCOL_VERSION,
							connectionCtx.getOrDefault(McpAsyncClient.NEGOTIATED_PROTOCOL_VERSION,
									this.latestSupportedProtocolVersion))
					.GET();
				var transportContext = connectionCtx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY);
				return Mono.from(this.httpRequestCustomizer.customize(builder, "GET", uri, null, transportContext));
			})
				.flatMapMany(requestBuilder -> Flux.<ResponseEvent>create(sseSink -> this.httpClient
					.sendAsync(requestBuilder.build(), this.toSendMessageBodySubscriber(sseSink))
					.whenComplete((response, throwable) -> {
						if (throwable != null) {
							sseSink.error(throwable);
						}
						else {
							logger.debug("SSE connection established successfully");
						}
					})).flatMap(responseEvent -> {
						int statusCode = responseEvent.responseInfo().statusCode();
						if (statusCode == 401 || statusCode == 403) {
							logger.debug("Authorization error in reconnect with code {}", statusCode);
							var request = requestBuilder.build();
							var requestSnapshot = new HttpRequestSnapshot(request.uri(), request.method(),
									request.headers());
							return Mono.<McpSchema.JSONRPCMessage>error(
									new McpHttpClientTransportAuthorizationException(
											"Authorization error connecting to SSE stream", requestSnapshot,
											responseEvent.responseInfo()));
						}
						else if (statusCode == METHOD_NOT_ALLOWED) {
							logger.debug("The server does not support SSE streams, using request-response mode.");
							return Flux.empty();
						}

						if (!(responseEvent instanceof ResponseSubscribers.SseResponseEvent sseResponseEvent)) {
							return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
									"Unrecognized server error when connecting to SSE stream, status code: "
											+ statusCode));
						}
						else if (statusCode >= 200 && statusCode < 300) {
							if (MESSAGE_EVENT_TYPE.equals(sseResponseEvent.sseEvent().event())) {
								String data = sseResponseEvent.sseEvent().data();
								// Per 2025-11-25 spec (SEP-1699), servers may
								// send SSE events
								// with empty data to prime the client for
								// reconnection.
								// Skip these events as they contain no JSON-RPC
								// message.
								if (data == null || data.isBlank()) {
									logger.debug("Skipping SSE event with empty data (stream primer)");
									return Flux.empty();
								}
								try {
									// We don't support batching ATM and probably
									// won't since the next version considers
									// removing it.
									McpSchema.JSONRPCMessage message = McpSchema
										.deserializeJsonRpcMessage(this.jsonMapper, data);

									Tuple2<Optional<String>, Iterable<McpSchema.JSONRPCMessage>> idWithMessages = Tuples
										.of(Optional.ofNullable(sseResponseEvent.sseEvent().id()), List.of(message));

									McpTransportStream<Disposable> sessionStream = stream != null ? stream
											: new DefaultMcpTransportStream<>(this.resumableStreams, this::reconnect);
									logger.debug("Connected stream {}", sessionStream.streamId());

									return Flux.from(sessionStream.consumeSseStream(Flux.just(idWithMessages)));

								}
								catch (IOException ioException) {
									return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
											"Error parsing JSON-RPC message: " + responseEvent, ioException));
								}
							}
							else {
								logger.debug("Received SSE event with type: {}", sseResponseEvent.sseEvent());
								return Flux.empty();
							}
						}
						else if (statusCode == NOT_FOUND) {

							if (transportSession != null && transportSession.sessionId().isPresent()) {
								// only if the request was sent with a session id
								// and the response is 404, we consider it a
								// session not found error.
								logger.debug("Session not found for session ID: {}",
										transportSession.sessionId().get());
								String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession);
								McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException(
										"Session not found for session ID: " + sessionIdRepresentation);
								return Flux.<McpSchema.JSONRPCMessage>error(exception);
							}
							return Flux.<McpSchema.JSONRPCMessage>error(
									new McpTransportException("Server Not Found. Status code:" + statusCode
											+ ", response-event:" + responseEvent));
						}
						else if (statusCode == BAD_REQUEST) {
							if (transportSession != null && transportSession.sessionId().isPresent()) {
								// only if the request was sent with a session id
								// and thre response is 404, we consider it a
								// session not found error.
								String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession);
								McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException(
										"Session not found for session ID: " + sessionIdRepresentation);
								return Flux.<McpSchema.JSONRPCMessage>error(exception);
							}
							return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
									"Bad Request. Status code:" + statusCode + ", response-event:" + responseEvent));
						}
						return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
								"Received unrecognized SSE event type: " + sseResponseEvent.sseEvent().event()));
					})
					.retryWhen(authorizationErrorRetrySpec())
					.flatMap(jsonrpcMessage -> requestHandler.apply(Mono.just(jsonrpcMessage)))
					.onErrorMap(CompletionException.class, t -> t.getCause())
					.doFinally(s -> {
						Disposable ref = disposableRef.getAndSet(null);
						if (ref != null) {
							transportSession.removeConnection(ref);
						}
					}))
				.onErrorComplete(t -> {
					this.handleException(t);
					return true;
				})
				.contextWrite(ctx)
				.subscribe();

			disposableRef.set(connection);
			transportSession.addConnection(connection);
			return Mono.just(connection);
		});

	}

	private Retry authorizationErrorRetrySpec() {
		return Retry.from(companion -> companion.flatMap(retrySignal -> {
			if (!(retrySignal.failure() instanceof McpHttpClientTransportAuthorizationException authException)) {
				return Mono.error(retrySignal.failure());
			}
			if (retrySignal.totalRetriesInARow() >= this.authorizationErrorHandler.maxRetries()) {
				return Mono.error(retrySignal.failure());
			}
			return Mono.deferContextual(ctx -> {
				var transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY);
				return Mono
					.from(this.authorizationErrorHandler.handle(authException.getRequestSnapshot(),
							authException.getResponseInfo(), transportContext))
					.switchIfEmpty(Mono.just(false))
					.flatMap(shouldRetry -> shouldRetry ? Mono.just(retrySignal.totalRetries())
							: Mono.error(retrySignal.failure()));
			});
		}));
	}

	private BodyHandler<Void> toSendMessageBodySubscriber(FluxSink<ResponseEvent> sink) {

		BodyHandler<Void> responseBodyHandler = responseInfo -> {

			String contentType = responseInfo.headers().firstValue(HttpHeaders.CONTENT_TYPE).orElse("").toLowerCase();

			if (contentType.contains(TEXT_EVENT_STREAM)) {
				// For SSE streams, use line subscriber that returns Void
				logger.debug("Received SSE stream response, using line subscriber");
				return ResponseSubscribers.sseToBodySubscriber(responseInfo, sink);
			}
			else if (contentType.contains(APPLICATION_JSON)) {
				// For JSON responses and others, use string subscriber
				logger.debug("Received response, using string subscriber");
				return ResponseSubscribers.aggregateBodySubscriber(responseInfo, sink);
			}

			logger.debug("Received Bodyless response, using discarding subscriber");
			return ResponseSubscribers.bodilessBodySubscriber(responseInfo, sink);
		};

		return responseBodyHandler;

	}

	public String toString(McpSchema.JSONRPCMessage message) {
		try {
			return this.jsonMapper.writeValueAsString(message);
		}
		catch (IOException e) {
			throw new RuntimeException("Failed to serialize JSON-RPC message", e);
		}
	}

	public Mono<Void> sendMessage(McpSchema.JSONRPCMessage sentMessage) {
		return Mono.create(deliveredSink -> {
			var rh = this.handler.get();
			if (rh == null) {
				logger.warn("Transport has no request handler registered. Remember to call connect!");
			}

			final Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> requestHandler = rh != null
					? rh : msg -> Mono.error(new IllegalStateException("No request handler"));

			var transportSession = this.activeSession.get();

			if (ClosedMcpTransportSession.INSTANCE.equals(transportSession)) {
				throw new McpTransportSessionClosedException();
			}

			logger.debug("Sending message {}", sentMessage);

			final AtomicReference<Disposable> disposableRef = new AtomicReference<>();

			var uri = Utils.resolveUri(this.baseUri, this.endpoint);
			String jsonBody = this.toString(sentMessage);

			Disposable connection = Mono.deferContextual(ctx -> {
				HttpRequest.Builder requestBuilder = this.requestBuilder.copy();

				if (transportSession != null && transportSession.sessionId().isPresent()) {
					requestBuilder = requestBuilder.header(HttpHeaders.MCP_SESSION_ID,
							transportSession.sessionId().get());
				}

				var builder = requestBuilder.uri(uri)
					.header(HttpHeaders.ACCEPT, APPLICATION_JSON + ", " + TEXT_EVENT_STREAM)
					.header(HttpHeaders.CONTENT_TYPE, APPLICATION_JSON_UTF8)
					.header(HttpHeaders.CACHE_CONTROL, "no-cache")
					.header(HttpHeaders.PROTOCOL_VERSION,
							ctx.getOrDefault(McpAsyncClient.NEGOTIATED_PROTOCOL_VERSION,
									this.latestSupportedProtocolVersion))
					.POST(HttpRequest.BodyPublishers.ofString(jsonBody));
				var transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY);
				return Mono
					.from(this.httpRequestCustomizer.customize(builder, "POST", uri, jsonBody, transportContext));
			}).flatMapMany(requestBuilder -> Flux.<ResponseEvent>create(responseEventSink -> {
				// Create the async request with proper body subscriber selection
				Mono.fromFuture(this.httpClient
					.sendAsync(requestBuilder.build(), this.toSendMessageBodySubscriber(responseEventSink))
					.whenComplete((response, throwable) -> {
						if (throwable != null) {
							responseEventSink.error(throwable);
						}
						else {
							logger.debug("SSE connection established successfully");
						}
					})).onErrorMap(CompletionException.class, t -> t.getCause()).onErrorComplete().subscribe();

			}).flatMap(responseEvent -> {
				int statusCode = responseEvent.responseInfo().statusCode();
				if (statusCode == 401 || statusCode == 403) {
					var request = requestBuilder.build();
					var requestSnapshot = new HttpRequestSnapshot(request.uri(), request.method(), request.headers());
					logger.debug("Authorization error in sendMessage with code {}", statusCode);
					return Mono.<McpSchema.JSONRPCMessage>error(new McpHttpClientTransportAuthorizationException(
							"Authorization error when sending message", requestSnapshot, responseEvent.responseInfo()));
				}

				if (transportSession.markInitialized(
						responseEvent.responseInfo().headers().firstValue("mcp-session-id").orElseGet(() -> null))) {
					// Once we have a session, we try to open an async stream for
					// the server to send notifications and requests out-of-band.

					reconnect(null).contextWrite(deliveredSink.contextView()).subscribe();
				}

				String sessionRepresentation = sessionIdOrPlaceholder(transportSession);

				if (statusCode >= 200 && statusCode < 300) {

					String contentType = responseEvent.responseInfo()
						.headers()
						.firstValue(HttpHeaders.CONTENT_TYPE)
						.orElse("")
						.toLowerCase();

					String contentLength = responseEvent.responseInfo()
						.headers()
						.firstValue(HttpHeaders.CONTENT_LENGTH)
						.orElse(null);

					// For empty content or HTTP code 202 (ACCEPTED), assume success
					if (contentType.isBlank() || "0".equals(contentLength) || statusCode == 202) {
						// if (contentType.isBlank() || "0".equals(contentLength)) {
						logger.debug("No body returned for POST in session {}", sessionRepresentation);
						// No content type means no response body, so we can just
						// return an empty stream
						deliveredSink.success();
						return Flux.empty();
					}
					else if (contentType.contains(TEXT_EVENT_STREAM)) {
						return Flux.just(((ResponseSubscribers.SseResponseEvent) responseEvent).sseEvent())
							.flatMap(sseEvent -> {
								String data = sseEvent.data();
								// Per 2025-11-25 spec (SEP-1699), servers may send SSE
								// events
								// with empty data to prime the client for reconnection.
								// Skip these events as they contain no JSON-RPC message.
								if (data == null || data.isBlank()) {
									logger.debug("Skipping SSE event with empty data (stream primer)");
									return Flux.empty();
								}
								try {
									// We don't support batching ATM and probably
									// won't
									// since the
									// next version considers removing it.
									McpSchema.JSONRPCMessage message = McpSchema
										.deserializeJsonRpcMessage(this.jsonMapper, data);

									Tuple2<Optional<String>, Iterable<McpSchema.JSONRPCMessage>> idWithMessages = Tuples
										.of(Optional.ofNullable(sseEvent.id()), List.of(message));

									McpTransportStream<Disposable> sessionStream = new DefaultMcpTransportStream<>(
											this.resumableStreams, this::reconnect);

									logger.debug("Connected stream {}", sessionStream.streamId());

									deliveredSink.success();

									return Flux.from(sessionStream.consumeSseStream(Flux.just(idWithMessages)));
								}
								catch (IOException ioException) {
									return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
											"Error parsing JSON-RPC message: " + responseEvent, ioException));
								}
							});
					}
					else if (contentType.contains(APPLICATION_JSON)) {
						deliveredSink.success();
						String data = ((ResponseSubscribers.AggregateResponseEvent) responseEvent).data();
						if (sentMessage instanceof McpSchema.JSONRPCNotification) {
							logger.warn("Notification: {} received non-compliant response: {}", sentMessage,
									Utils.hasText(data) ? data : "[empty]");
							return Mono.empty();
						}

						try {
							return Mono.just(McpSchema.deserializeJsonRpcMessage(jsonMapper, data));
						}
						catch (IOException e) {
							return Mono.error(new McpTransportException(
									"Error deserializing JSON-RPC message: " + responseEvent, e));
						}
					}
					logger.warn("Unknown media type {} returned for POST in session {}", contentType,
							sessionRepresentation);

					return Flux.<McpSchema.JSONRPCMessage>error(
							new RuntimeException("Unknown media type returned: " + contentType));
				}
				else if (statusCode == NOT_FOUND) {
					if (transportSession != null && transportSession.sessionId().isPresent()) {
						// only if the request was sent with a session id and the
						// response is 404, we consider it a session not found error.
						logger.debug("Session not found for session ID: {}", transportSession.sessionId().get());
						McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException(
								"Session not found for session ID: " + sessionRepresentation);
						return Flux.<McpSchema.JSONRPCMessage>error(exception);
					}
					return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
							"Server Not Found. Status code:" + statusCode + ", response-event:" + responseEvent));
				}
				else if (statusCode == BAD_REQUEST) {
					// Some implementations can return 400 when presented with a
					// session id that it doesn't know about, so we will
					// invalidate the session
					// https://github.com/modelcontextprotocol/typescript-sdk/issues/389

					if (transportSession != null && transportSession.sessionId().isPresent()) {
						// only if the request was sent with a session id and the
						// response is 404, we consider it a session not found error.
						McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException(
								"Session not found for session ID: " + sessionRepresentation);
						return Flux.<McpSchema.JSONRPCMessage>error(exception);
					}
					return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
							"Bad Request. Status code:" + statusCode + ", response-event:" + responseEvent));
				}

				return Flux.<McpSchema.JSONRPCMessage>error(
						new RuntimeException("Failed to send message: " + responseEvent));
			})
				.retryWhen(authorizationErrorRetrySpec())
				.flatMap(jsonRpcMessage -> requestHandler.apply(Mono.just(jsonRpcMessage)))
				.onErrorMap(CompletionException.class, t -> t.getCause())
				.doFinally(s -> {
					logger.debug("SendMessage finally: {}", s);
					Disposable ref = disposableRef.getAndSet(null);
					if (ref != null) {
						transportSession.removeConnection(ref);
					}
				})).onErrorComplete(t -> {
					// handle the error first
					try {
						this.handleException(t);
					}
					catch (Exception e) {
						logger.error("Error handling exception {}", t.getMessage(), e);
					}
					// inform the caller of sendMessage
					deliveredSink.error(t);
					return true;
				}).contextWrite(deliveredSink.contextView()).subscribe();

			disposableRef.set(connection);
			transportSession.addConnection(connection);
		});

	}

	private static String sessionIdOrPlaceholder(McpTransportSession<?> transportSession) {
		return transportSession.sessionId().orElse("[missing_session_id]");
	}

	@Override
	public <T> T unmarshalFrom(Object data, TypeRef<T> typeRef) {
		return this.jsonMapper.convertValue(data, typeRef);
	}

	/**
	 * Builder for {@link HttpClientStreamableHttpTransport}.
	 */
	public static class Builder {

		private final String baseUri;

		private McpJsonMapper jsonMapper;

		private HttpClient.Builder clientBuilder = HttpClient.newBuilder().version(HttpClient.Version.HTTP_1_1);

		private String endpoint = DEFAULT_ENDPOINT;

		private boolean resumableStreams = true;

		private boolean openConnectionOnStartup = false;

		private HttpRequest.Builder requestBuilder = HttpRequest.newBuilder();

		private McpAsyncHttpClientRequestCustomizer httpRequestCustomizer = McpAsyncHttpClientRequestCustomizer.NOOP;

		private Duration connectTimeout = Duration.ofSeconds(10);

		private List<String> supportedProtocolVersions = List.of(ProtocolVersions.MCP_2024_11_05,
				ProtocolVersions.MCP_2025_03_26, ProtocolVersions.MCP_2025_06_18, ProtocolVersions.MCP_2025_11_25);

		private McpHttpClientTransportAuthorizationErrorHandler authorizationErrorHandler = McpHttpClientTransportAuthorizationErrorHandler.NOOP;

		/**
		 * Creates a new builder with the specified base URI.
		 * @param baseUri the base URI of the MCP server
		 */
		private Builder(String baseUri) {
			Assert.hasText(baseUri, "baseUri must not be empty");
			this.baseUri = baseUri;
		}

		/**
		 * Sets the HTTP client builder.
		 * @param clientBuilder the HTTP client builder
		 * @return this builder
		 */
		public Builder clientBuilder(HttpClient.Builder clientBuilder) {
			Assert.notNull(clientBuilder, "clientBuilder must not be null");
			this.clientBuilder = clientBuilder;
			return this;
		}

		/**
		 * Customizes the HTTP client builder.
		 * @param clientCustomizer the consumer to customize the HTTP client builder
		 * @return this builder
		 */
		public Builder customizeClient(final Consumer<HttpClient.Builder> clientCustomizer) {
			Assert.notNull(clientCustomizer, "clientCustomizer must not be null");
			clientCustomizer.accept(clientBuilder);
			return this;
		}

		/**
		 * Sets the HTTP request builder.
		 * @param requestBuilder the HTTP request builder
		 * @return this builder
		 */
		public Builder requestBuilder(HttpRequest.Builder requestBuilder) {
			Assert.notNull(requestBuilder, "requestBuilder must not be null");
			this.requestBuilder = requestBuilder;
			return this;
		}

		/**
		 * Configure a custom {@link McpJsonMapper} implementation to use.
		 * @param jsonMapper instance to use
		 * @return the builder instance
		 */
		public Builder jsonMapper(McpJsonMapper jsonMapper) {
			Assert.notNull(jsonMapper, "jsonMapper must not be null");
			this.jsonMapper = jsonMapper;
			return this;
		}

		/**
		 * Configure the endpoint to make HTTP requests against.
		 * @param endpoint endpoint to use
		 * @return the builder instance
		 */
		public Builder endpoint(String endpoint) {
			Assert.hasText(endpoint, "endpoint must be a non-empty String");
			this.endpoint = endpoint;
			return this;
		}

		/**
		 * Configure whether to use the stream resumability feature by keeping track of
		 * SSE event ids.
		 * @param resumableStreams if {@code true} event ids will be tracked and upon
		 * disconnection, the last seen id will be used upon reconnection as a header to
		 * resume consuming messages.
		 * @return the builder instance
		 */
		public Builder resumableStreams(boolean resumableStreams) {
			this.resumableStreams = resumableStreams;
			return this;
		}

		/**
		 * Configure whether the client should open an SSE connection upon startup. Not
		 * all servers support this (although it is in theory possible with the current
		 * specification), so use with caution. By default, this value is {@code false}.
		 * @param openConnectionOnStartup if {@code true} the {@link #connect(Function)}
		 * method call will try to open an SSE connection before sending any JSON-RPC
		 * request
		 * @return the builder instance
		 */
		public Builder openConnectionOnStartup(boolean openConnectionOnStartup) {
			this.openConnectionOnStartup = openConnectionOnStartup;
			return this;
		}

		/**
		 * Sets the customizer for {@link HttpRequest.Builder}, to modify requests before
		 * executing them.
		 * <p>
		 * This overrides the customizer from
		 * {@link #asyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer)}.
		 * <p>
		 * Do NOT use a blocking {@link McpSyncHttpClientRequestCustomizer} in a
		 * non-blocking context. Use
		 * {@link #asyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer)}
		 * instead.
		 * @param syncHttpRequestCustomizer the request customizer
		 * @return this builder
		 */
		public Builder httpRequestCustomizer(McpSyncHttpClientRequestCustomizer syncHttpRequestCustomizer) {
			this.httpRequestCustomizer = McpAsyncHttpClientRequestCustomizer.fromSync(syncHttpRequestCustomizer);
			return this;
		}

		/**
		 * Sets the customizer for {@link HttpRequest.Builder}, to modify requests before
		 * executing them.
		 * <p>
		 * This overrides the customizer from
		 * {@link #httpRequestCustomizer(McpSyncHttpClientRequestCustomizer)}.
		 * <p>
		 * Do NOT use a blocking implementation in a non-blocking context.
		 * @param asyncHttpRequestCustomizer the request customizer
		 * @return this builder
		 */
		public Builder asyncHttpRequestCustomizer(McpAsyncHttpClientRequestCustomizer asyncHttpRequestCustomizer) {
			this.httpRequestCustomizer = asyncHttpRequestCustomizer;
			return this;
		}

		/**
		 * Sets the handler to be used when the server responds with HTTP 401 or HTTP 403
		 * when sending a message.
		 * @param authorizationErrorHandler the handler
		 * @return this builder
		 * @deprecated in favor of
		 * {@link #authorizationErrorHandler(McpHttpClientTransportAuthorizationErrorHandler)}
		 */
		@Deprecated(forRemoval = true, since = "2.0.0")
		public Builder authorizationErrorHandler(McpHttpClientAuthorizationErrorHandler authorizationErrorHandler) {
			this.authorizationErrorHandler = new McpHttpClientTransportAuthorizationErrorHandler() {
				@Override
				public Publisher<Boolean> handle(HttpRequestSnapshot requestSnapshot,
						HttpResponse.ResponseInfo responseInfo, McpTransportContext context) {
					return authorizationErrorHandler.handle(responseInfo, context);
				}

				@Override
				public int maxRetries() {
					return authorizationErrorHandler.maxRetries();
				}
			};
			return this;
		}

		/**
		 * Sets the handler to be used when the server responds with HTTP 401 or HTTP 403
		 * when sending a message.
		 * @param authorizationErrorHandler the handler
		 * @return this builder
		 */
		public Builder authorizationErrorHandler(
				McpHttpClientTransportAuthorizationErrorHandler authorizationErrorHandler) {
			this.authorizationErrorHandler = authorizationErrorHandler;
			return this;
		}

		/**
		 * Sets the connection timeout for the HTTP client.
		 * @param connectTimeout the connection timeout duration
		 * @return this builder
		 */
		public Builder connectTimeout(Duration connectTimeout) {
			Assert.notNull(connectTimeout, "connectTimeout must not be null");
			this.connectTimeout = connectTimeout;
			return this;
		}

		/**
		 * Sets the list of supported protocol versions used in version negotiation. By
		 * default, the client will send the latest of those versions in the
		 * {@code MCP-Protocol-Version} header.
		 * <p>
		 * Setting this value only updates the values used in version negotiation, and
		 * does NOT impact the actual capabilities of the transport. It should only be
		 * used for compatibility with servers having strict requirements around the
		 * {@code MCP-Protocol-Version} header.
		 * @param supportedProtocolVersions protocol versions supported by this transport
		 * @return this builder
		 * @see <a href=
		 * "https://modelcontextprotocol.io/specification/2024-11-05/basic/lifecycle#version-negotiation">version
		 * negotiation specification</a>
		 * @see <a href=
		 * "https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#protocol-version-header">Protocol
		 * Version Header</a>
		 */
		public Builder supportedProtocolVersions(List<String> supportedProtocolVersions) {
			Assert.notEmpty(supportedProtocolVersions, "supportedProtocolVersions must not be empty");
			this.supportedProtocolVersions = Collections.unmodifiableList(supportedProtocolVersions);
			return this;
		}

		/**
		 * Construct a fresh instance of {@link HttpClientStreamableHttpTransport} using
		 * the current builder configuration.
		 * @return a new instance of {@link HttpClientStreamableHttpTransport}
		 */
		public HttpClientStreamableHttpTransport build() {
			HttpClient httpClient = this.clientBuilder.connectTimeout(this.connectTimeout).build();
			return new HttpClientStreamableHttpTransport(jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper,
					httpClient, requestBuilder, baseUri, endpoint, resumableStreams, openConnectionOnStartup,
					httpRequestCustomizer, authorizationErrorHandler, supportedProtocolVersions);
		}

	}

}
