content: /*
   * Copyright 2024-2024 the original author or authors.
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *      https://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package com.alibaba.langengine.mcp;
  
  import com.alibaba.fastjson.JSON;
  import com.alibaba.fastjson.TypeReference;
  import com.alibaba.langengine.mcp.spec.*;
  import reactor.core.publisher.Flux;
  import reactor.core.publisher.Mono;
  import reactor.core.publisher.Sinks;
  import reactor.core.scheduler.Schedulers;
  
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.function.Function;
  
  @SuppressWarnings("unused")
  public class MockMcpTransport implements ClientMcpTransport, ServerMcpTransport {
  
  	private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
  
  	private final Sinks.Many<JSONRPCMessage> outgoing = Sinks.many().multicast().onBackpressureBuffer();
  
  	private final Sinks.Many<JSONRPCMessage> inbound = Sinks.many().unicast().onBackpressureBuffer();
  
  	private final Flux<JSONRPCMessage> outboundView = outgoing.asFlux().cache(1);
  
  	public void simulateIncomingMessage(JSONRPCMessage message) {
  		if (inbound.tryEmitNext(message).isFailure()) {
  			throw new RuntimeException("Failed to emit message " + message);
  		}
  		inboundMessageCount.incrementAndGet();
  	}
  
  	@Override
  	public Mono<Void> sendMessage(JSONRPCMessage message) {
  		if (outgoing.tryEmitNext(message).isFailure()) {
  			return Mono.error(new RuntimeException("Can't emit outgoing message " + message));
  		}
  		return Mono.empty();
  	}
  
  	public JSONRPCRequest getLastSentMessageAsRequest() {
  		return (JSONRPCRequest) outboundView.blockFirst();
  	}
  
  	public JSONRPCNotification getLastSentMessageAsNotifiation() {
  		return (JSONRPCNotification) outboundView.blockFirst();
  	}
  
  	public JSONRPCMessage getLastSentMessage() {
  		return outboundView.blockFirst();
  	}
  
  	private volatile boolean connected = false;
  
  	@Override
  	public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
  		if (connected) {
  			return Mono.error(new IllegalStateException("Already connected"));
  		}
  		connected = true;
  		return inbound.asFlux()
  			.publishOn(Schedulers.boundedElastic())
  			.flatMap(message -> Mono.just(message).transform(handler))
  			.doFinally(signal -> connected = false)
  			.then();
  	}
  
  	@Override
  	public Mono<Void> closeGracefully() {
  		return Mono.defer(() -> {
  			connected = false;
  			outgoing.tryEmitComplete();
  			inbound.tryEmitComplete();
  			return Mono.empty();
  		});
  	}
  
  	@Override
  	public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
  //		return new ObjectMapper().convertValue(data, typeRef);
  		String json = JSON.toJSONString(data);
  		return JSON.parseObject(json, typeRef);
  	}
  }