package com.gxwebsoft.ai.util;

import com.aliyun.bailian20231229.models.RetrieveRequest;
import com.aliyun.bailian20231229.models.RetrieveResponse;
import com.aliyun.teautil.models.RuntimeOptions;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.util.*;

/**
 * 检索知识库
 * @author GIIT-YC
 *
 */
public class KnowledgeBaseRetrieve {

	static String ALIBABA_CLOUD_ACCESS_KEY_ID = "LTAI5tD5YRKuxWz6Eg7qrM4P";
	static String ALIBABA_CLOUD_ACCESS_KEY_SECRET = "bO8TBDXflOwbtSKimPpG8XrJnyzgTk";
	static String WORKSPACE_ID = "llm-4pf5auwewoz34zqu";
	
	
	/**
	 * 检查并提示设置必要的环境变量。
	 *
	 * @return true 如果所有必需的环境变量都已设置，否则 false
	 */
	public static boolean checkEnvironmentVariables() {
		Map<String, String> requiredVars = new HashMap<>();
        requiredVars.put("ALIBABA_CLOUD_ACCESS_KEY_ID", "阿里云访问密钥ID");
        requiredVars.put("ALIBABA_CLOUD_ACCESS_KEY_SECRET", "阿里云访问密钥密码");
        requiredVars.put("WORKSPACE_ID", "阿里云百炼业务空间ID");

		List<String> missingVars = new ArrayList<>();
		for (Map.Entry<String, String> entry : requiredVars.entrySet()) {
			String value = System.getenv(entry.getKey());
			if (value == null || value.isEmpty()) {
				missingVars.add(entry.getKey());
				System.out.println("错误：请设置 " + entry.getKey() + " 环境变量 (" + entry.getValue() + ")");
			}
		}

		return missingVars.isEmpty();
	}

	/**
	 * 初始化客户端（Client）。
	 *
	 * @return 配置好的客户端对象
	 */
	public static com.aliyun.bailian20231229.Client createClient() throws Exception {
		com.aliyun.teaopenapi.models.Config config = new com.aliyun.teaopenapi.models.Config()
				.setAccessKeyId(ALIBABA_CLOUD_ACCESS_KEY_ID)
				.setAccessKeySecret(ALIBABA_CLOUD_ACCESS_KEY_SECRET);
		// 下方接入地址以公有云的公网接入地址为例，可按需更换接入地址。
		config.endpoint = "bailian.cn-beijing.aliyuncs.com";
		return new com.aliyun.bailian20231229.Client(config);
	}

	/**
	 * 在指定的知识库中检索信息。
	 *
	 * @param client      客户端对象（bailian20231229Client）
	 * @param workspaceId 业务空间ID
	 * @param indexId     知识库ID
	 * @param query       检索查询语句
	 * @return 阿里云百炼服务的响应
	 */
	public static RetrieveResponse retrieveIndex(com.aliyun.bailian20231229.Client client, String workspaceId,
			String indexId, String query) throws Exception {
		RetrieveRequest retrieveRequest = new RetrieveRequest();
		retrieveRequest.setIndexId(indexId);
		retrieveRequest.setQuery(query);
		retrieveRequest.setDenseSimilarityTopK(null);
		RuntimeOptions runtime = new RuntimeOptions();
		return client.retrieveWithOptions(workspaceId, retrieveRequest, null, runtime);
	}

	/**
	 * 使用阿里云百炼服务检索知识库。
	 */
	public static void main(String[] args) {
//		if (!checkEnvironmentVariables()) {
//			System.out.println("环境变量校验未通过。");
//			return;
//		}

		try {
			// 步骤1：初始化客户端（Client）
			System.out.println("步骤1：创建Client");
			com.aliyun.bailian20231229.Client client = createClient();

			// 步骤2：检索知识库
			System.out.println("步骤2：检索知识库");
			Scanner scanner = new Scanner(System.in);
			System.out.print("请输入知识库ID："); // 即 CreateIndex 接口返回的 Data.Id，您也可以在阿里云百炼控制台的知识库页面获取。
			String indexId = scanner.nextLine();
			System.out.print("请输入检索query：");
			String query = scanner.nextLine();
			String workspaceId = WORKSPACE_ID;
			RetrieveResponse resp = retrieveIndex(client, workspaceId, indexId, query);

			// 请自行安装jackson-databind。将响应体responsebody转换为 JSON 字符串
			ObjectMapper mapper = new ObjectMapper();
			String result = mapper.writeValueAsString(resp.getBody());
			System.out.println(result);
		} catch (Exception e) {
			System.out.println("发生错误：" + e.getMessage());
		}
	}
}