package tech.aiflowy.ai.config;

import com.agentsflex.core.llm.response.AiMessageResponse;
import com.agentsflex.core.message.AiMessage;
import com.agentsflex.core.message.Message;
import com.agentsflex.core.prompt.HistoriesPrompt;
import com.agentsflex.core.prompt.Prompt;
import com.alibaba.fastjson.JSON;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.JsonArray;
import com.google.gson.JsonObject;
import com.mybatisflex.core.query.QueryWrapper;
import okhttp3.*;
import okio.BufferedSource;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import tech.aiflowy.ai.entity.AiBotConversationMessage;
import tech.aiflowy.ai.entity.AiBotMessage;
import tech.aiflowy.ai.service.AiBotConversationMessageService;
import tech.aiflowy.ai.service.AiBotMessageService;
import tech.aiflowy.ai.service.impl.AiBotMessageServiceImpl;
import tech.aiflowy.common.ai.MySseEmitter;
import tech.aiflowy.common.util.StringUtil;
import tech.aiflowy.common.web.controller.BaseCurdController;

import javax.annotation.Resource;
import java.io.IOException;
import java.math.BigInteger;
import java.util.*;
import java.util.concurrent.CompletableFuture;
public class DifyStreamClient {
    private final OkHttpClient client;
    private final String apiUrl;
    private final String apiKey;
    private final Gson gson;
    private String prompt;
    private AiBotMessageService aiBotMessageService;

    public DifyStreamClient(String apiUrl, String apiKey, AiBotMessageService aiBotMessageService) {
        this.apiUrl = apiUrl;
        this.apiKey = apiKey;
        this.gson = new GsonBuilder().setPrettyPrinting().create();
        this.client = new OkHttpClient.Builder().build();
        this.aiBotMessageService = aiBotMessageService;
    }

    // 流式聊天方法 - 直接集成SseEmitter
    public CompletableFuture<Void> chatStream(String message, String userId, MySseEmitter emitter, String sessionId, BigInteger botId) {
        prompt = message;
        QueryWrapper qw = new QueryWrapper();
        qw.eq(AiBotMessage::getSessionId, sessionId)
                .orderBy(AiBotMessage::getId,false)
                .limit(6);

        List<AiBotMessage> history = aiBotMessageService.list(qw);
//        System.out.println("======================history==================\n");
//        for (AiBotMessage aiBotMessage : history) {
//            System.out.println(aiBotMessage.toMessage());
//        }
//        System.out.println("\n======================history==================");

        // 构建请求JSON
        JsonObject requestBody = new JsonObject();
        requestBody.add("inputs", new JsonObject());
        requestBody.addProperty("query", message);
        requestBody.addProperty("response_mode", "streaming");
        requestBody.addProperty("conversation_id", "");
        requestBody.addProperty("user", userId);
        requestBody.add("files", new JsonArray());
        // 添加历史对话信息
        JsonArray historyArray = new JsonArray();
        for (AiBotMessage msg : history) {
            historyArray.add(String.valueOf(msg));
        }
        requestBody.add("history", historyArray);

        RequestBody body = RequestBody.create(
                gson.toJson(requestBody),
                MediaType.parse("application/json; charset=utf-8")
        );

        Request request = new Request.Builder()
                .url(apiUrl)
                .post(body)
                .header("Authorization", apiKey)
                .header("Content-Type", "application/json")
                .build();

        CompletableFuture<Void> future = new CompletableFuture<>();

        // 设置SseEmitter生命周期回调
        emitter.onTimeout(() -> {
            System.out.println("SSE连接超时");
            emitter.complete();
            future.complete(null);
        });

        emitter.onCompletion(() -> {
            System.out.println("SSE连接已完成");
            future.complete(null);
        });

        emitter.onError(e -> {
            System.out.println("SSE连接错误: " + e.getMessage());
            emitter.completeWithError(e);
            future.completeExceptionally(e);
        });

        // 发送异步请求
        client.newCall(request).enqueue(new Callback() {
            @Override
            public void onFailure(Call call, IOException e) {
                emitter.completeWithError(e);
                future.completeExceptionally(e);
            }

            @Override
            public void onResponse(Call call, Response response) {
                AiMessage aiMessage = new AiMessage();
                com.agentsflex.core.llm.response.AiMessageResponse aiMessageResponse
                        = new com.agentsflex.core.llm.response.AiMessageResponse(new HistoriesPrompt(), response.message(), aiMessage);
                try (ResponseBody responseBody = response.body()) {
                    if (!response.isSuccessful()) {
                        emitter.completeWithError(new IOException("API错误: " + response.code()));
                        return;
                    }

                    // 使用BufferedSource逐行读取响应内容
                    BufferedSource source = responseBody.source();
                    String line;
                    StringBuffer sb = new StringBuffer();
                    // 标记是否为第一条有效数据（用于处理某些API的特殊格式）
                    boolean isFirstData = true;

                    while ((line = source.readUtf8Line()) != null) {
                        if (line.startsWith("data: ")) {
                            String data = line.substring(6).trim();

                            // 忽略空数据或结束标记
                            if (data.isEmpty() || data.equals("[DONE]")) {
                                continue;
                            }

                            try {
                                // 根据实际API响应结构提取消息内容
                                // 这里需要根据具体API调整路径
                                JsonObject messageObj = gson.fromJson(data, JsonObject.class);
//                                System.out.println(messageObj);
                                if(messageObj.get("event").getAsString().equals("message_end")){
                                    AiBotMessage aiBotMessage = new AiBotMessage();
                                    aiBotMessage.setBotId(botId);
                                    aiBotMessage.setSessionId(sessionId);
                                    aiBotMessage.setAccountId(new BigInteger(userId));
                                    aiBotMessage.setRole("assistant");
                                    aiBotMessage.setContent(sb.toString());
                                    aiBotMessage.setCreated(new Date());
                                    aiBotMessage.setIsExternalMsg(1);
                                    aiBotMessageService.save(aiBotMessage);
//                                    System.out.println("end");
                                }else{
                                    String context = messageObj.get("answer").getAsString();
                                    if (context != null) {
//                                    // 只移除HTML标签，保留Markdown特殊字符
//                                        context = context.replaceAll("(?i)<[^>]*>", "");
                                        context = context.replaceFirst("Thinking...", "");
                                    }
                                    sb.append(context);
                                    aiMessage.setContent(context);
                                    aiMessageResponse.setMessage(aiMessage);
                                    if (StringUtil.hasText(messageObj.get("answer").getAsString())) {
//                                        System.out.println(aiMessage);
                                        if(aiMessage.getContent().startsWith("</details>")){
                                            aiMessage.setContent(aiMessage.getContent().replaceAll("(?i)<[^>]*>", "</details>"+"\n\n"));
                                        }
                                        // 发送消息片段给前端
                                        emitter.send(JSON.toJSONString(aiMessage));
                                    }
                                }

                                // 重置第一条数据标记
                                isFirstData = false;

                            } catch (Exception e) {
                                // 记录解析错误但继续处理后续数据
                                emitter.completeWithError(e);
                            }
                        }
                    }

                    // 所有数据处理完毕，发送完成信号
                    emitter.send(SseEmitter.event().name("complete"));
                    emitter.complete();

                } catch (IOException e) {
                    emitter.completeWithError(e);
                    emitter.complete();
                } catch (Exception e) {
                    // 处理其他异常
                    emitter.completeWithError(e);
                }
            }

//            @Override
            public void onResponses(Call call, Response response) {
//                try (ResponseBody responseBody = response.body()) {
//                    if (!response.isSuccessful() || responseBody == null) {
//                        IOException e = new IOException("Unexpected code " + response);
//                        emitter.completeWithError(e);
//                        future.completeExceptionally(e);
//                        return;
//                    }
////                    String rawResponse = responseBody.string(); // 打印原始响应
////                    System.out.println("Dify原始响应: " + rawResponse);
//
//                    BufferedSource source = responseBody.source();
//                    String line;
//
//                    // 处理流式响应
//                    while ((line = source.readUtf8Line()) != null) {
//                        if (line.startsWith("data: ")) {
//                            String data = line.substring(6); // 移除"data: "前缀
//
//                            // 跳过空数据
//                            if (data.trim().equals("[DONE]")) {
//                                emitter.complete();
//                                future.complete(null);
//                                break;
//                            }
//
//                            try {
//                                // 解析JSON响应
//                                JsonObject jsonResponse = gson.fromJson(data, JsonObject.class);
//                                // 发送消息内容到客户端
//                                if (jsonResponse.has("message") && jsonResponse.get("message").isJsonObject()) {
//                                    String content = jsonResponse.getAsJsonObject("message").get("content").getAsString();
//                                    emitter.send(SseEmitter.event()
//                                            .data(content)
//                                            .name("message")
//                                    );
//                                }
//
//                                // 处理函数调用
//                                if (jsonResponse.has("function_call") && jsonResponse.get("function_call").isJsonObject()) {
//                                    emitter.send(SseEmitter.event()
//                                            .data(jsonResponse.get("function_call"))
//                                            .name("function_call")
//                                    );
//                                }
//
//                            } catch (Exception e) {
//                                emitter.completeWithError(e);
//                                future.completeExceptionally(e);
//                                break;
//                            }
//                        }
//                    }
                try (ResponseBody responseBody = response.body()) {
                    if (!response.isSuccessful()) {
                        emitter.completeWithError(new IOException("API错误"));
                        return;
                    }
                    StringBuilder fullAnswer = new StringBuilder();
                    BufferedSource source = responseBody.source();
                    String line;
                    while ((line = source.readUtf8Line()) != null) {
                        if (line.startsWith("data: ")) {
                            String data = line.substring(6).trim();

                            if (data.equals("[DONE]")) {
                                // 发送完整回答并结束流
                                emitter.send(fullAnswer.toString());
                                emitter.complete();
                                break;
                            }

                            try {
                                JsonObject json = gson.fromJson(data, JsonObject.class);
                                System.out.println(json);
                                String answerFragment = json.get("answer").getAsString();
                                // 过滤无效字符（如️⃣、\n）
                                String cleanedFragment = answerFragment.replaceAll("[^\u4e00-\u9fa50-9a-zA-Z\\s\\p{Punct}]", "");
                                fullAnswer.append(cleanedFragment); // 拼接碎片

                                // 可选：每拼接一部分发送给前端（需前端支持增量渲染）
                                // emitter.send(SseEmitter.event().data(cleanedFragment));

                            } catch (Exception e) {
                                emitter.sendAndComplete("解析错误: " + e.getMessage());
                            }
                        }
                    }
                } catch (Exception e) {
                    emitter.completeWithError(e);
                    future.completeExceptionally(e);
                }
            }
        });

        return future;
    }

    // 数据模型类
    public interface StreamResponseListener {
        void onMessage(ChatContext context, AiMessageResponse response);
        void onStop(ChatContext context);
        void onFailure(ChatContext context, Throwable throwable);
    }

    public static class ChatContext {
        // 上下文信息
    }

    public static class AiMessageResponse {
        private Message message;
        private JsonArray functionCallers;

        public Message getMessage() { return message; }
        public void setMessage(Message message) { this.message = message; }
        public JsonArray getFunctionCallers() { return functionCallers; }
        public void setFunctionCallers(JsonArray functionCallers) { this.functionCallers = functionCallers; }
    }

    public static class Message {
        private String content;

        public String getContent() { return content; }
        public void setContent(String content) { this.content = content; }
    }
}