/*
 * Copyright 2024-2026 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.cloud.ai.a2a.core.server;

import org.springframework.web.servlet.function.ServerRequest;

import java.time.Duration;
import java.util.Map;
import java.util.concurrent.Flow;
import java.util.function.Function;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import io.a2a.server.requesthandlers.JSONRPCHandler;
import io.a2a.spec.AgentCard;
import io.a2a.spec.CancelTaskRequest;
import io.a2a.spec.DeleteTaskPushNotificationConfigRequest;
import io.a2a.spec.GetTaskPushNotificationConfigRequest;
import io.a2a.spec.GetTaskRequest;
import io.a2a.spec.JSONParseError;
import io.a2a.spec.JSONRPCError;
import io.a2a.spec.JSONRPCErrorResponse;
import io.a2a.spec.JSONRPCRequest;
import io.a2a.spec.JSONRPCResponse;
import io.a2a.spec.ListTaskPushNotificationConfigRequest;
import io.a2a.spec.MessageSendParams;
import io.a2a.spec.NonStreamingJSONRPCRequest;
import io.a2a.spec.SendMessageRequest;
import io.a2a.spec.SendStreamingMessageRequest;
import io.a2a.spec.SetTaskPushNotificationConfigRequest;
import io.a2a.spec.StreamingJSONRPCRequest;
import io.a2a.spec.TaskResubscriptionRequest;
import io.a2a.spec.UnsupportedOperationError;
import io.a2a.util.Utils;
import org.reactivestreams.FlowAdapters;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;

/**
 * The request handler for A2A protocol request by JSON RPC 2.0.
 *
 * @author xiweng.yy
 */
public class JsonRpcA2aRequestHandler implements A2aRequestHandler {

	private static final Logger LOGGER = LoggerFactory.getLogger(JsonRpcA2aRequestHandler.class);

	private final JSONRPCHandler jsonRpcHandler;

	public JsonRpcA2aRequestHandler(JSONRPCHandler jsonRpcHandler) {
		this.jsonRpcHandler = jsonRpcHandler;
	}

	@Override
	public AgentCard getAgentCard() {
		return jsonRpcHandler.getAgentCard();
	}

	@Override
	public Object onHandler(String body, ServerRequest.Headers headers) {
		boolean streaming = isStreamingRequest(body);
		Object result = null;
		try {
			result = streaming ? handleStreamRequest(body) : handleNonStreamRequest(body);
		}
		catch (JsonProcessingException e) {
			result = new JSONRPCErrorResponse(null, new JSONParseError());
		}
		return result;
	}

	private static boolean isStreamingRequest(String requestBody) {
		try {
			JsonNode node = Utils.OBJECT_MAPPER.readTree(requestBody);
			JsonNode method = node != null ? node.get("method") : null;
			return method != null && (SendStreamingMessageRequest.METHOD.equals(method.asText())
					|| TaskResubscriptionRequest.METHOD.equals(method.asText()));
		}
		catch (Exception e) {
			return false;
		}
	}

	private Flux<?> handleStreamRequest(String body) throws JsonProcessingException {
		StreamingJSONRPCRequest<?> request = Utils.OBJECT_MAPPER.readValue(body, StreamingJSONRPCRequest.class);
		Flow.Publisher<? extends JSONRPCResponse<?>> publisher;
		if (request instanceof SendStreamingMessageRequest req) {
			SendStreamingMessageRequest.Builder newReqBuilder = new SendStreamingMessageRequest.Builder()
				.id(req.getId())
				.jsonrpc(req.getJsonrpc())
				.method(req.getMethod())
				.params(injectStreamMetadata(req.getParams(), true));
			publisher = jsonRpcHandler.onMessageSendStream(newReqBuilder.build());
			LOGGER.info("get Stream publisher {}", publisher);
		}
		else if (request instanceof TaskResubscriptionRequest req) {
			publisher = jsonRpcHandler.onResubscribeToTask(req);
		}
		else {
			return Flux.just(generateErrorResponse(request, new UnsupportedOperationError()));
		}
		return Flux.from(FlowAdapters.toPublisher(publisher))
			.map((Function<JSONRPCResponse<?>, JSONRPCResponse<?>>) jsonrpcResponse -> jsonrpcResponse)
			.delaySubscription(Duration.ofMillis(10));
	}

	private JSONRPCResponse<?> handleNonStreamRequest(String body) throws JsonProcessingException {
		NonStreamingJSONRPCRequest<?> request = Utils.OBJECT_MAPPER.readValue(body, NonStreamingJSONRPCRequest.class);
		if (request instanceof GetTaskRequest req) {
			return jsonRpcHandler.onGetTask(req);
		}
		else if (request instanceof SendMessageRequest req) {
			SendMessageRequest.Builder newReqBuilder = new SendMessageRequest.Builder().id(req.getId())
				.jsonrpc(req.getJsonrpc())
				.method(req.getMethod())
				.params(injectStreamMetadata(req.getParams(), false));
			return jsonRpcHandler.onMessageSend(newReqBuilder.build());
		}
		else if (request instanceof CancelTaskRequest req) {
			return jsonRpcHandler.onCancelTask(req);
		}
		else if (request instanceof GetTaskPushNotificationConfigRequest req) {
			return jsonRpcHandler.getPushNotificationConfig(req);
		}
		else if (request instanceof SetTaskPushNotificationConfigRequest req) {
			return jsonRpcHandler.setPushNotificationConfig(req);
		}
		else if (request instanceof ListTaskPushNotificationConfigRequest req) {
			return jsonRpcHandler.listPushNotificationConfig(req);
		}
		else if (request instanceof DeleteTaskPushNotificationConfigRequest req) {
			return jsonRpcHandler.deletePushNotificationConfig(req);
		}
		else {
			return generateErrorResponse(request, new UnsupportedOperationError());
		}
	}

	private static JSONRPCErrorResponse generateErrorResponse(JSONRPCRequest<?> request, JSONRPCError error) {
		return new JSONRPCErrorResponse(request.getId(), error);
	}

	private MessageSendParams injectStreamMetadata(MessageSendParams original, boolean isStreaming) {
		if (null == original.metadata()) {
			MessageSendParams.Builder newBuilder = new MessageSendParams.Builder();
			newBuilder.configuration(original.configuration());
			newBuilder.metadata(Map.of(GraphAgentExecutor.STREAMING_METADATA_KEY, isStreaming));
			newBuilder.message(original.message());
			return newBuilder.build();
		}
		else {
			original.metadata().put(GraphAgentExecutor.STREAMING_METADATA_KEY, isStreaming);
			return original;
		}
	}

}
