/*
 * Copyright 2025 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.studio.core.base.manager;

import com.alibaba.cloud.ai.studio.runtime.domain.mcp.Content;
import com.alibaba.cloud.ai.studio.runtime.domain.mcp.McpServerCallToolRequest;
import com.alibaba.cloud.ai.studio.runtime.domain.mcp.McpServerCallToolResponse;
import com.alibaba.cloud.ai.studio.runtime.domain.mcp.McpServerDeployConfig;
import com.alibaba.cloud.ai.studio.runtime.domain.mcp.McpTool;
import com.alibaba.cloud.ai.studio.runtime.domain.mcp.TextContent;
import com.alibaba.cloud.ai.studio.runtime.enums.McpInstallTypeEnum;
import com.alibaba.cloud.ai.studio.runtime.enums.McpServerStatusEnum;
import com.alibaba.cloud.ai.studio.runtime.domain.Result;
import com.alibaba.cloud.ai.studio.runtime.domain.tool.InputSchema;
import com.alibaba.cloud.ai.studio.runtime.utils.JsonUtils;
import com.alibaba.cloud.ai.studio.core.base.constants.CacheConstants;
import com.alibaba.cloud.ai.studio.core.base.entity.McpServerEntity;
import com.alibaba.cloud.ai.studio.core.utils.LogUtils;
import com.alibaba.cloud.ai.studio.core.utils.concurrent.ThreadPoolUtils;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
import io.modelcontextprotocol.spec.McpClientTransport;
import io.modelcontextprotocol.spec.McpSchema;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;

import java.net.URL;
import java.net.http.HttpClient;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

import static com.alibaba.cloud.ai.studio.runtime.enums.ErrorCode.MCP_PARSE_CONFIG_ERROR;
import static com.alibaba.cloud.ai.studio.runtime.enums.ErrorCode.MCP_PARSE_URL_ERROR;
import static com.alibaba.cloud.ai.studio.core.utils.LogUtils.FAIL;
import static com.alibaba.cloud.ai.studio.core.utils.LogUtils.SUCCESS;

/**
 * Manager class for handling Model Context Protocol (MCP) operations. Provides
 * functionality for client creation, tool management, and server configuration.
 */
@Slf4j
@Component
public class MCPManager {

	/** Redis manager for caching operations */
	private final RedisManager redisManager;

	public MCPManager(RedisManager redisManager) {
		this.redisManager = redisManager;
	}

	/**
	 * Creates a synchronous MCP client for the given server entity.
	 * @param entity Server entity containing deployment configuration
	 * @return Synchronous MCP client instance
	 */
	public McpSyncClient getMcpSyncClient(McpServerEntity entity) {
		McpServerDeployConfig deployConfig = JsonUtils.fromJson(entity.getDeployConfig(), McpServerDeployConfig.class);

		String remoteEndpoint = deployConfig.getRemoteEndpoint();
		if (StringUtils.isBlank(remoteEndpoint)) {
			remoteEndpoint = "/sse";
		}
		String host = entity.getHost();
		LogUtils.error("Client host", host, entity.getServerCode());
		HttpClient.Builder builder1 = HttpClient.newBuilder()
			.version(HttpClient.Version.HTTP_1_1)
			.connectTimeout(Duration.ofSeconds(60));
		McpClientTransport transport = HttpClientSseClientTransport.builder(host)
			.clientBuilder(builder1)
			.sseEndpoint(remoteEndpoint)
			.build();
		McpClient.SyncSpec builder = McpClient.sync(transport)
			.requestTimeout(Duration.ofSeconds(60))
			.initializationTimeout(Duration.ofSeconds(60))
			.capabilities(McpSchema.ClientCapabilities.builder().roots(true).build());
		return builder.build();
	}

	/**
	 * Retrieves available tools from the specified MCP server. Implements caching
	 * mechanism for non-SSE installations.
	 * @param entity Server entity to query tools from
	 * @return List of available tools
	 */
	public List<McpTool> getTools(McpServerEntity entity) {
		long start = System.currentTimeMillis();
		if (entity == null || !entity.getStatus().equals(McpServerStatusEnum.Normal.getCode())) {
			return new ArrayList<>();
		}
		String cacheKey = "mcp_tools_cache_" + entity.getServerCode();
		boolean needCache = !"SSE".equals(entity.getInstallType());
		if (needCache) {
			List<McpTool> toolList = redisManager.get(cacheKey);
			if (toolList != null) {
				return toolList;
			}
		}
		Future<Object> future = ThreadPoolUtils.TOOL_TASK_EXECUTOR.submit(new Callable<Object>() {
			@Override
			public Object call() throws Exception {
				McpSyncClient client = getMcpSyncClient(entity);
				try {
					client.initialize();
					McpSchema.ListToolsResult toolsResult = client.listTools();
					List<McpSchema.Tool> toolsList = toolsResult.tools();
					List<McpTool> tools = new ArrayList<>();
					if (CollectionUtils.isEmpty(toolsList)) {
						return tools;
					}
					toolsList.forEach(e -> {
						McpTool tool = new McpTool();
						tool.setName(e.name());
						tool.setDescription(e.description());
						McpSchema.JsonSchema jsonSchema = e.inputSchema();
						InputSchema inputSchema = new InputSchema();
						inputSchema.setType(jsonSchema.type());
						inputSchema.setProperties(jsonSchema.properties());
						inputSchema.setRequired(jsonSchema.required());
						inputSchema.setAdditionalProperties(jsonSchema.additionalProperties());
						tool.setInputSchema(inputSchema);
						tools.add(tool);
					});
					LogUtils.error("FromServer GetTool", entity.getServerCode(), System.currentTimeMillis() - start,
							tools.size());
					return tools;
				}
				catch (Exception ex) {
					LogUtils.error("getTools error", ex, entity.getServerCode(), entity.getHost());
				}
				finally {
					if (client != null) {
						client.close();
					}
				}
				return new ArrayList<>();
			}
		});
		try {
			List<McpTool> result = (List<McpTool>) future.get(3, TimeUnit.SECONDS);
			if (result != null && !result.isEmpty() && needCache) {
				redisManager.put(cacheKey, result, CacheConstants.CACHE_EMPTY_TTL);
			}
			LogUtils.error("FromServer Last", entity.getServerCode(), System.currentTimeMillis() - start, result);
			return result;
		}
		catch (Exception ex) {
			LogUtils.error("toolsGetFuture Exception", ex, entity.getServerCode());
			return new ArrayList<>();
		}
	}

	/**
	 * Executes a specific tool on the MCP server.
	 * @param request Tool execution request parameters
	 * @param entity Server entity where the tool resides
	 * @return Response from the tool execution
	 */
	public McpServerCallToolResponse callTool(McpServerCallToolRequest request, McpServerEntity entity) {
		Long start = System.currentTimeMillis();
		McpSyncClient client = getMcpSyncClient(entity);
		McpServerCallToolResponse response = new McpServerCallToolResponse();
		try {
			client.initialize();
			McpSchema.CallToolRequest callToolRequest = new McpSchema.CallToolRequest(request.getToolName(),
					request.getToolParams());
			McpSchema.CallToolResult callToolResult = client.callTool(callToolRequest);
			response.setIsError(callToolResult.isError());
			List<McpSchema.Content> contentList = callToolResult.content();
			List<Content> content = new ArrayList<>();
			contentList.forEach(e -> {
				if (e.type().equals("text")) {
					TextContent textContent = new TextContent();
					textContent.setType(e.type());
					McpSchema.TextContent tContent = (McpSchema.TextContent) e;
					textContent.setText(tContent.text());
					content.add(textContent);
				}
			});
			response.setContent(content);
		}
		catch (Exception ex) {
			LogUtils.monitor("McpService", "callTool", start, FAIL, request, ex.getMessage(), ex);
			LogUtils.error("McpServerManager callTool exception", ex);
			throw ex;
		}
		finally {
			if (client != null) {
				client.close();
			}
		}
		LogUtils.monitor("McpService", "callTool", start, SUCCESS, request, response);
		return response;
	}

	/**
	 * Processes and validates installation configuration. Handles URL validation and
	 * configuration transformation based on installation type.
	 * @param originDeployConfig Original deployment configuration string
	 * @param installType Installation type (e.g., SSE)
	 * @return Result containing processed configuration or error
	 */
	public Result<String> processInstallConfig(String originDeployConfig, String installType) {
		try {
			McpInstallTypeEnum installTypeEnum = McpInstallTypeEnum.of(installType);
			HashMap<String, Object> targetDeployConfig = new HashMap<>();
			Map<String, Object> originConfig = JsonUtils.fromJsonToMap(originDeployConfig);
			Map<String, Object> mcpServers = JsonUtils.fromJsonToMap(JsonUtils.toJson(originConfig.get("mcpServers")));

			if (mcpServers == null || mcpServers.keySet().size() != 1) {
				LogUtils.error("ParseConfigError", "mcpServers must be configured and only one is supported");
				return Result.error(MCP_PARSE_CONFIG_ERROR);
			}
			targetDeployConfig.put("install_config", originDeployConfig);
			if (installTypeEnum == McpInstallTypeEnum.SSE) {
				// check request header and host address
				for (String singleServer : mcpServers.keySet()) {
					Map<String, Object> singleServerConfig = JsonUtils
						.fromJsonToMap(JsonUtils.toJson(mcpServers.get(singleServer)));
					String url = (String) singleServerConfig.get("url");
					try {
						URL urlObj = new URL(url);
//						if (!urlObj.getPath().endsWith("/sse")) {
//							return Result.error(MCP_PARSE_URL_ERROR);
//						}

						String remoteAddress = urlObj.getProtocol() + "://" + urlObj.getHost();
						if (urlObj.getPort() != -1) {
							remoteAddress += ":" + urlObj.getPort();
						}
						targetDeployConfig.put("remote_address", remoteAddress);
						targetDeployConfig.put("remote_endpoint", urlObj.getPath());
						String query = urlObj.getQuery();
						if (StringUtils.isNotBlank(query)) {
							targetDeployConfig.put("remote_endpoint", urlObj.getPath() + "?" + query);
						}
						Map<Object, Object> headers = JsonUtils
							.fromJsonToMap(JsonUtils.toJson(singleServerConfig.get("headers")));
						targetDeployConfig.put("remote_header", headers);
					}
					catch (Exception urlCheckEx) {
						LogUtils.error("processInstallConfig", url, urlCheckEx);
						return Result.error(MCP_PARSE_URL_ERROR);
					}
					break;
				}
			}
			return Result.success(JsonUtils.toJson(targetDeployConfig));
		}
		catch (Exception ex) {
			LogUtils.error("processInstallConfig exception", ex, originDeployConfig);
			throw ex;
		}
	}

}
