package org.dromara.system.service.impl;

import com.coze.openapi.client.chat.CreateChatReq;
import com.coze.openapi.client.chat.model.ChatEvent;
import com.coze.openapi.client.chat.model.ChatEventType;
import com.coze.openapi.client.connversations.message.model.Message;
import com.coze.openapi.service.auth.TokenAuth;
import com.coze.openapi.service.config.Consts;
import com.coze.openapi.service.service.CozeAPI;
import io.reactivex.Flowable;
import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.dromara.system.domain.bo.CozeChatBo;
import org.dromara.system.service.ICozeChatService;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import jakarta.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Collections;

@Slf4j
@Service
@RequiredArgsConstructor
public class CozeChatServiceImpl implements ICozeChatService {

    @Value("${coze.apiKey}")
    private String apiKey;

    @Value("${coze.botId}")
    private String defaultBotId;

    private CozeAPI coze;

    @PostConstruct
    public void init() {
        TokenAuth authCli = new TokenAuth(apiKey);
        coze =
            new CozeAPI.Builder()
                .baseURL(Consts.COZE_CN_BASE_URL)
                .auth(authCli)
                .readTimeout(10000)
                .build();
    }

    @Override
    public void chatStream(CozeChatBo bo, HttpServletResponse response) {
        response.setContentType("text/event-stream");
        response.setCharacterEncoding("UTF-8");
        response.setHeader("Cache-Control", "no-cache");
        response.setHeader("Connection", "keep-alive");

        try (PrintWriter writer = response.getWriter()) {
            String botId = bo.getBotId() != null ? bo.getBotId() : defaultBotId;
            String userId = bo.getUserId() != null ? bo.getUserId() : "user_123";

            // 构建创建聊天请求
            CreateChatReq req = CreateChatReq.builder()
                .botID(botId)
                .userID(userId)
                .messages(Collections.singletonList(Message.buildUserQuestionText(bo.getContent())))
                .build();

            // 如果有会话ID，设置会话ID
            if (bo.getConversationId() != null && !bo.getConversationId().isEmpty()) {
                req.setConversationID(bo.getConversationId());
            }

            // 调用流式聊天接口
            Flowable<ChatEvent> resp = coze.chat().stream(req);

            // 处理流式响应
            resp.blockingForEach(
                event -> {
                    try {
                        if (ChatEventType.CONVERSATION_MESSAGE_DELTA.equals(event.getEvent())) {
                            // 处理消息增量
                            String content = event.getMessage().getContent();
                            if (content != null && !content.isEmpty()) {
                                writer.write("data: " + content + "\n\n");
                                writer.flush();
                            }
                        } else if (ChatEventType.CONVERSATION_CHAT_COMPLETED.equals(event.getEvent())) {
                            // 处理聊天完成事件，输出token使用量
                            if (event.getChat() != null && event.getChat().getUsage() != null) {
                                long tokenCount = event.getChat().getUsage().getTokenCount();
                                log.info("Token usage: {}", tokenCount);
                                writer.write("data: [DONE] Token usage: " + tokenCount + "\n\n");
                                writer.flush();
                            } else {
                                writer.write("data: [DONE]\n\n");
                                writer.flush();
                            }
                        }
                    } catch (Exception e) {
                        log.error("写入流数据失败", e);
                    }
                });

        } catch (Exception e) {
            log.error("流式对话失败", e);
            try (PrintWriter writer = response.getWriter()) {
                writer.write("data: [ERROR] " + e.getMessage() + "\n\n");
                writer.flush();
            } catch (IOException ioException) {
                log.error("写入错误信息失败", ioException);
            }
        }
    }
}
