package com.smtaiserver.smtaiserver.javaai.qwen.agent;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.smtaiserver.smtaiserver.core.SMTAIServerApp;
import com.smtaiserver.smtaiserver.core.SMTAIServerRequest;
import com.smtaiserver.smtaiserver.database.SMTDatabase.DBRecord;
import com.smtaiserver.smtaiserver.javaai.SMTJavaAIError;
import com.smtaiserver.smtaiserver.javaai.llm.core.SMTLLMConnect;
import com.smtservlet.util.Json;
import com.smtservlet.util.SMTJsonWriter;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.TimeUnit;
import javax.servlet.http.HttpServletResponse;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;

public class SMTQwenAgentLightRAG extends SMTQwenAgent
{

	@Override
	public void initInstance(DBRecord rec) throws Exception
	{
		super.initInstance(rec);

	}

	@Override
	public SMTJavaAIError callAgents(String jsonPath, Json jsonArgs, SMTLLMConnect llm, String question, SMTAIServerRequest tranReq) throws Exception
	{
		SMTJsonWriter jsonWrResult = tranReq.getResultJsonWr();

		StringBuilder sbAnswer = new StringBuilder();
		SMTJavaAIError error = callRAGServer(question, tranReq, sbAnswer);
		if(error != null)
			return error;
		jsonWrResult.addKeyValue("json_ok", true);
		jsonWrResult.addKeyValue("answer_type", "knowledge");
		jsonWrResult.beginArray("knowledge");
		{
			jsonWrResult.beginMap(null);
			jsonWrResult.addKeyValue("answer", sbAnswer.toString());
			jsonWrResult.endMap();
		}
		jsonWrResult.endArray();
		return null;
	}



	private SMTJavaAIError callRAGServer(String question, SMTAIServerRequest tranReq, StringBuilder sbAnswer) throws Exception {
		OkHttpClient okHttpClient = new OkHttpClient.Builder()
				.readTimeout(0, TimeUnit.SECONDS) // 不超时，支持流
				.build();
		HttpServletResponse response = tranReq.getResponse();
		response.setContentType("application/json");
		response.setCharacterEncoding("UTF-8");

		SMTJsonWriter jsonWriter = new SMTJsonWriter(false);
		jsonWriter.addKeyValue("query",question);
		String json = jsonWriter.getRootJson().toString();

		RequestBody body = RequestBody.create(MediaType.parse("application/json"), json);
		String lightragQueyStream = (String)SMTAIServerApp.getApp().getGlobalConfig("lightrag_quey_stream");
		Request request = new Request.Builder()
				.url(lightragQueyStream)
				.post(body)
				.build();
		tranReq.sendChunkedBlock("begin_stream", "");
		ObjectMapper objectMapper = new ObjectMapper(); // 用于解析 JSON
		try (Response lightRagResp = okHttpClient.newCall(request).execute()) {
			if (lightRagResp.body() != null) {
				BufferedReader reader = new BufferedReader(new InputStreamReader(lightRagResp.body().byteStream(), StandardCharsets.UTF_8));
				String line;
				while ((line = reader.readLine()) != null) {
					try {
						// 每行是 {"response": "xxx"}，我们只取出 "xxx"
						JsonNode node = objectMapper.readTree(line);
						String content = node.get("response").asText(); // 自动处理编码问题，返回的是中文

						// 逐字符推送
						for (char ch : content.toCharArray()) {
							tranReq.sendChunkedStreamBlock(com.smtservlet.util.SMTStatic.toString(ch));
							sbAnswer.append(ch);
							Thread.sleep(50);
						}

					} catch (Exception e) {
						System.err.println("解析失败：" + line);
					}
				}
			}
		} catch (Exception e) {
			String errMsg = "调用 LightRAG 失败: " + e.getMessage();
			for (char ch : errMsg.toCharArray()) {
				tranReq.sendChunkedStreamBlock(com.smtservlet.util.SMTStatic.toString(ch));
				sbAnswer.append(ch);
			}
		} finally {
			tranReq.sendChunkedBlock("end_stream", "");
		}

		return null;
	}

}