package com.lemon.testng;

import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.alibaba.fastjson.JSON;
import com.google.common.collect.Lists;
import com.lemon.utils.AuthorizationUtils;
import com.pcloud.llm.cockpit.api.constants.Constants;
import com.pcloud.llm.cockpit.api.domain.dto.request.WorkflowRequestBody;
import com.pcloud.llm.cockpit.api.domain.dto.response.CompletionResponse;
import com.pcloud.llm.cockpit.api.domain.dto.response.CompletionResponseData;
import org.apache.log4j.Logger;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;

import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;

public class WebSocketManager {


    // 日志
    public static Logger log = Logger.getLogger(WebSocketManager.class);
    
    public WebSocketClient webSocketClient;
    public CountDownLatch latch4Event3;
    public CountDownLatch latch4Event5;
    public String msgRemoteId;
    public String receivedRequestId;
    public String expectedRequestId;
    // 创建定时任务发送心跳消息
    public ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
    public List<String> event5Messages = new ArrayList<>();
    public List<String> event5Messages4Stream = new ArrayList<>();
    public String mergeStreamMessage;
    public Integer messageCount;



    public WebSocketManager() throws URISyntaxException {
        String requestWsUrl = AuthorizationUtils.env.get("request_ws_url");
        requestWsUrl = StrUtil.isBlank(requestWsUrl) ? "wss://wss.5rs.me/chat/ws" : requestWsUrl;
        webSocketClient = new WebSocketClient(new URI(requestWsUrl)) {
            @Override
            public void onOpen(ServerHandshake handshake) {
                log.info("WebSocket连接已打开");
            }

            @Override
            public void onMessage(String message) {
                log.info("收到消息: " + message);

                // 检查消息是否为空或是数组
                if (StrUtil.isBlank(message) || message.startsWith("[\"{")) {
                    log.info("接收的消息是一个数组，不是普通对话消息，不处理");
                    return;
                }


                // 解析消息为 JSON 对象
                JSONObject receiveJson = JSONUtil.parseObj(message);
                Integer event = receiveJson.getInt("event");

                if (event == null || Lists.newArrayList(1).contains(event)) {
                    log.info("event为null或心跳信息不处理");
                    return;
                }

                Optional.ofNullable(receiveJson.getJSONObject("messageBody")).ifPresentOrElse(messageBody -> {
                    if (event == 3) {
                        // 提取并验证 requestId
                        receivedRequestId = receiveJson.getStr("requestId");
                        log.info("receivedRequestId是" + receivedRequestId + "、expectedRequestId是" + expectedRequestId);

                        if (receivedRequestId.equals(expectedRequestId)) {
                            msgRemoteId = messageBody.getStr("msgRemoteId");
                            latch4Event3.countDown();
                        }
                    } else if (event == 5) {
                        // 处理 event-5 消息
                        handleEvent5(messageBody);
                    }
                }, () -> log.info("messageBody为null，消息异常不处理"));
            }

            private void handleEvent5(JSONObject messageBody) {
                Boolean stream = messageBody.getBool("stream");
                String streamState = messageBody.getStr("streamState");
                JSONObject messageBodyJSONObject = messageBody.getJSONObject("message");
                if  (null == messageBodyJSONObject) {
                    log.info("messageBodyJSONObject为null，不继续处理");
                }
                String msgRemoteId4Event5 = messageBodyJSONObject.getStr("msgRemoteId");
                String contentType = messageBodyJSONObject.getStr("contentType");
                JSONObject json4Option = messageBodyJSONObject.getJSONObject("optionMsgDTO");
                String subType = json4Option != null ? json4Option.getStr("subType") : null;

                if (!msgRemoteId.equals(msgRemoteId4Event5)) {
                    log.info("收到的msgRemoteId不匹配，忽略");
                    return;
                }

                if (Lists.newArrayList("Option").contains(contentType) &&
                        Lists.newArrayList("CHAT_RECOMMEND").contains(subType)) {
                    log.info("推荐问题，不记录");
                    return;
                }

                if (!stream) {
                    if (msgRemoteId.equals(msgRemoteId4Event5)) {
                        event5Messages.add(messageBody.getStr("message"));
                        latch4Event5.countDown();
                    }
                } else {
                    handleStreamMessages(messageBody, streamState);
                }
            }

            private void handleStreamMessages(JSONObject messageBody, String streamState) {
                if ("START".equals(streamState)) {
                    mergeStreamMessage = messageBody.getStr("message");
                } else if ("END".equals(streamState)) {
                    latch4Event5.countDown();
                    event5Messages.add(mergeStreamMessage);
                } else {
                    JSONObject json4Message = JSONUtil.parseObj(mergeStreamMessage);
                    JSONObject json4Content = Optional.ofNullable(json4Message.getJSONObject("content"))
                            .orElseGet(() -> {
                                log.warn("mergeStreamMessage 中 content 为 null");
                                return new JSONObject();
                            });

                    JSONObject json4Text = Optional.ofNullable(json4Content.getJSONObject("textMsgDTO"))
                            .orElseGet(() -> {
                                log.warn("mergeStreamMessage 中 textMsgDTO 为 null");
                                return new JSONObject();
                            });

                    String content = json4Text.getStr("content", "");

                    JSONObject json4Message4Process = JSONUtil.parseObj(messageBody.getStr("message"));
                    JSONObject json4Content4Process = Optional.ofNullable(json4Message4Process.getJSONObject("content"))
                            .orElseGet(() -> {
                                log.warn("messageBody 中 content 为 null");
                                return new JSONObject();
                            });

                    JSONObject json4Text4Process = Optional.ofNullable(json4Content4Process.getJSONObject("textMsgDTO"))
                            .orElseGet(() -> {
                                log.warn("messageBody 中 textMsgDTO 为 null");
                                return new JSONObject();
                            });
                    String content4Process = json4Text4Process.getStr("content", "");
                    content = content + content4Process;
                    json4Text.set("content", content);
                    mergeStreamMessage = JSONUtil.toJsonStr(json4Message);
                    log.info("当前的流逝消息累积是：" + mergeStreamMessage);
                }
            }



            @Override
            public void onClose(int code, String reason, boolean remote) {
                log.info("WebSocket连接已关闭，原因: " + reason);
            }

            @Override
            public void onError(Exception ex) {
                log.info("WebSocket出错: " + ex.getMessage());
            }
        };
    }

    public void connect() {
        webSocketClient.connect();
        // 等待连接成功逻辑可以添加在这里
    }

    public void sendMessage(String message) {
        webSocketClient.send(message);
    }

    public void close() {
        webSocketClient.close();
    }

    /**
     * 工作流同步接口
     * @param callCode  Dify调用code
     * @param dto       参数
     * @return    同步返回
     */
    public CompletionResponse workflowsRunBlocking(String callCode, WorkflowRequestBody dto) {
        dto.setResponse_mode(Constants.RESPONSE_MODE_BLOCKING);
        HttpHeaders headers = new HttpHeaders();
        headers.add("Content-Type", MediaType.APPLICATION_JSON_VALUE);
        headers.add(Constants.CALL_CODE_HEADER, callCode);
         try {
            String server = "http://192.168.92.235:8313";
            HttpRequest post= HttpUtil.createPost(server + "/v1/workflows/run");
            post.addHeaders(headers.toSingleValueMap());
            post.body(JSONUtil.toJsonStr(dto));
            HttpResponse httpResponse = post.execute();
            CompletionResponse result = JSONUtil.toBean(httpResponse.body(), CompletionResponse.class);
            log.info("workflow:" + callCode + "->" + JSON.toJSONString(dto) + "=>" + JSON.toJSONString(result) + ",server=" + server);
            return result;
        }  catch (Throwable t) {
            log.error("workflow:" + callCode + "->" + JSON.toJSONString(dto), t);
            CompletionResponse result = new CompletionResponse();
            CompletionResponseData data = new CompletionResponseData();
            data.setStatus("failed");
            data.setError(t.getMessage());
            return result;
        }
    }
}