/*
 * 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.graph.node;

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
import io.modelcontextprotocol.spec.McpSchema.TextContent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * MCP Node: Node for calling MCP Server
 */
public class McpNode implements NodeAction {

	private static final Pattern VARIABLE_PATTERN = Pattern.compile("\\$\\{(.+?)\\}");

	private static final Logger log = LoggerFactory.getLogger(McpNode.class);

	private final String url;

	private final String tool;

	private final Map<String, String> headers;

	private final Map<String, Object> params;

	private final String outputKey;

	private final List<String> inputParamKeys;

	private HttpClientSseClientTransport transport;

	private McpSyncClient client;

	private McpNode(Builder builder) {
		this.url = builder.url;
		this.tool = builder.tool;
		this.headers = builder.headers;
		this.params = builder.params;
		this.outputKey = builder.outputKey;
		this.inputParamKeys = builder.inputParamKeys;
	}

	@Override
	public Map<String, Object> apply(OverAllState state) throws Exception {
		log.info(
				"[McpNode] Start executing apply, original configuration: url={}, tool={}, headers={}, inputParamKeys={}",
				url, tool, headers, inputParamKeys);

		// Build transport and client
		String baseUrl = this.url;
		String sseEndpoint = "/sse";
		if (this.url.contains("/sse?")) {
			int idx = this.url.indexOf("/sse?");
			baseUrl = this.url.substring(0, idx);
			sseEndpoint = this.url.substring(idx); // e.g. /sse?key=xxx
		}
		HttpClientSseClientTransport.Builder transportBuilder = HttpClientSseClientTransport.builder(baseUrl)
			.sseEndpoint(sseEndpoint);
		if (this.headers != null && !this.headers.isEmpty()) {
			transportBuilder.customizeRequest(req -> this.headers.forEach(req::header));
		}
		this.transport = transportBuilder.build();
		this.client = McpClient.sync(this.transport).build();
		InitializeResult initializeResult = this.client.initialize();
		log.info("[McpNode] MCP Client initialized: {}", initializeResult);
		// Variable replacement
		String finalTool = replaceVariables(tool, state);
		Map<String, Object> finalParams = new HashMap<>();
		// 1. First read from inputParamKeys
		if (inputParamKeys != null) {
			for (String key : inputParamKeys) {
				Object value = state.value(key).orElse(null);
				if (value != null) {
					finalParams.put(key, value);
				}
			}
		}
		// 2. Then use params (after variable replacement) to overwrite
		Map<String, Object> replacedParams = replaceVariablesObj(params, state);
		if (replacedParams != null) {
			finalParams.putAll(replacedParams);
		}
		log.info("[McpNode] after replace params: url={}, tool={}, headers={}, params={}", url, finalTool, headers,
				finalParams);

		// Directly use the already initialized client
		CallToolResult result;
		try {
			McpSchema.CallToolRequest request = new McpSchema.CallToolRequest(finalTool, finalParams);
			log.info("[McpNode] CallToolRequest: {}", request);
			result = client.callTool(request);
			log.info("[McpNode] tool call result: {}", result);
		}
		catch (Exception e) {
			log.error("[McpNode] MCP call fail:", e);
			throw new McpNodeException("MCP call fail: " + e.getMessage(), e);
		}

		// Result handling
		Map<String, Object> updatedState = new HashMap<>();
		// updatedState.put("mcp_result", result.content());
		updatedState.put("messages", result.content());
		if (StringUtils.hasLength(this.outputKey)) {
			Object content = result.content();
			if (content instanceof List<?> list && !CollectionUtils.isEmpty(list)) {
				Object first = list.get(0);
				// Compatible with the text field of TextContent
				if (first instanceof TextContent textContent) {
					updatedState.put(this.outputKey, textContent.text());
				}
				else if (first instanceof Map<?, ?> map && map.containsKey("text")) {
					updatedState.put(this.outputKey, map.get("text"));
				}
				else {
					updatedState.put(this.outputKey, first);
				}
			}
			else {
				updatedState.put(this.outputKey, content);
			}
		}
		log.info("[McpNode] update state: {}", updatedState);
		return updatedState;
	}

	private String replaceVariables(String template, OverAllState state) {
		if (template == null)
			return null;
		Matcher matcher = VARIABLE_PATTERN.matcher(template);
		StringBuilder result = new StringBuilder();
		while (matcher.find()) {
			String key = matcher.group(1);
			Object value = state.value(key).orElse("");
			log.info("[McpNode] replace param: {} -> {}", key, value);
			matcher.appendReplacement(result, value.toString());
		}
		matcher.appendTail(result);
		return result.toString();
	}

	private Map<String, Object> replaceVariablesObj(Map<String, Object> map, OverAllState state) {
		if (map == null)
			return null;
		Map<String, Object> result = new HashMap<>();
		map.forEach((k, v) -> {
			if (v instanceof String) {
				result.put(k, replaceVariables((String) v, state));
			}
			else {
				result.put(k, v);
			}
		});
		return result;
	}

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

	public static class Builder {

		private String url;

		private String tool;

		private Map<String, String> headers = new HashMap<>();

		private Map<String, Object> params = new HashMap<>();

		private String outputKey;

		private List<String> inputParamKeys;

		public Builder url(String url) {
			this.url = url;
			return this;
		}

		public Builder tool(String tool) {
			this.tool = tool;
			return this;
		}

		public Builder header(String name, String value) {
			this.headers.put(name, value);
			return this;
		}

		public Builder param(String name, Object value) {
			this.params.put(name, value);
			return this;
		}

		public Builder outputKey(String outputKey) {
			this.outputKey = outputKey;
			return this;
		}

		public Builder inputParamKeys(List<String> inputParamKeys) {
			this.inputParamKeys = inputParamKeys;
			return this;
		}

		public McpNode build() {
			return new McpNode(this);
		}

	}

	public static class McpNodeException extends RuntimeException {

		public McpNodeException(String message, Throwable cause) {
			super(message, cause);
		}

	}

}
