package cn.teammodel.ai.listener;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import cn.teammodel.ai.SseHelper;
import cn.teammodel.ai.domain.SparkChatResponse;
import cn.teammodel.common.ErrorCode;
import cn.teammodel.config.exception.ServiceException;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.annotation.Nullable;
import javax.validation.constraints.NotNull;
import java.util.function.Consumer;

/**
 * okhttp 调用 ws 接口时的 listener
 * @author winter
 * @create 2023-12-15 16:17
 */
@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
@RequiredArgsConstructor
public class SparkGptStreamListener extends WebSocketListener {
    private final SseEmitter sseEmitter;
    private String requestJson;
    /**
    * 模型响应的完整回复
    */
    private String answer = "";

    @Override
    public void onOpen(WebSocket webSocket, @NotNull Response response) {
        // ws 建立连接后发送请求的数据, 在 onMessage 中接受 server 的响应数据
        webSocket.send(requestJson);
        // 执行成功回调
        try {
            onOpen.accept(webSocket);
        } catch (Exception e) {
            this.onFailure(webSocket, e, response);
        }
    }

    @Override
    public void onMessage(WebSocket webSocket, String text) {
        // 向 sse 推送消息(一次 onMessage 事件推送一次)
        SparkChatResponse sparkChatResponse = JSONUtil.toBean(text, SparkChatResponse.class);
        // 请求是否异常
        if (sparkChatResponse.getHeader().getCode() != 0) {
            this.onFailure(
                    webSocket,
                    new ServiceException(ErrorCode.SYSTEM_ERROR.getCode(),
                    sparkChatResponse.getHeader().getMessage()),
            null
            );
        }
        // 推送回答 segment
        String msgSegment = sparkChatResponse.getPayload().getChoices().getText().get(0).getContent();
//         处理消息格式(空格和换行符)
//        msgSegment = StrUtil.replace(msgSegment, " ", "&#32;").replaceAll("\n", "&#92n");
        msgSegment = StrUtil.replace(msgSegment, " ", "&#32;").replaceAll("\n", "<br>");
        answer += msgSegment;
        SseHelper.send(sseEmitter, msgSegment);

        // 处理模型的最终响应
        if (sparkChatResponse.getHeader().getStatus() == 2) {
            // 其实 spark 会主动断开连接
            webSocket.close(1000, "done");
            onComplete.accept(answer);
            SseHelper.complete(sseEmitter);
        }
    }

    @Override
    public void onFailure(WebSocket webSocket, Throwable t, @Nullable Response response) {
        webSocket.close(1000, t.getMessage());
        this.onError.accept(t);
        // 失败时结束 sse 连接
        SseHelper.send(sseEmitter,t.getMessage() + "[DONE]");
        SseHelper.complete(sseEmitter);
    }


    // 这几个 function 可以在 listener 被调用时设置, 实现类似事件的回调(可以使用模板方法模式实现)
    protected Consumer<WebSocket> onOpen = (s) -> {};
    protected Consumer<Throwable> onError = (s) -> {};
    protected Consumer<String> onComplete = (s) -> {};
}
