package com.ruoyi.aicall.llm.impl;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.coze.openapi.client.auth.OAuthToken;
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.JWTOAuthClient;
import com.coze.openapi.service.auth.TokenAuth;
import com.coze.openapi.service.service.CozeAPI;
import com.ruoyi.aicall.domain.CcLlmAgentAccount;
import com.ruoyi.aicall.llm.AbstractLlmCapability;
import com.ruoyi.aicall.llm.model.CozeAccount;
import com.ruoyi.aicall.model.ChatRequest;
import com.ruoyi.common.utils.CommonUtils;
import com.ruoyi.common.utils.uuid.UuidGenerator;
import io.reactivex.Flowable;
import io.reactivex.schedulers.Schedulers;
import link.thingscloud.freeswitch.esl.util.CurrentTimeMillisClock;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.util.CollectionUtils;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

@Slf4j
public class Coze extends AbstractLlmCapability {

    private String conversationId = "";
    private static final String COZE_TOKEN_TYPE_PAT = "pat";
    private static final String COZE_TOKEN_TYPE_OAUTH = "oauth";
    private String token = "";
    private int expireTime = 0;

    @Override
    public String getIntentionByTips(CcLlmAgentAccount ccLlmAgentAccount, JSONArray dialogueContents) {
        String intentionTips = ccLlmAgentAccount.getIntentionTips();
        return null;
    }

    @Override
    public void getStreamingChatContent(SseEmitter emitter, ChatRequest chatRequest) throws IOException {

        String requestId = chatRequest.getRequestId();
        String question = chatRequest.getQuestion();
        if (CollectionUtils.isEmpty(chatRequest.getMessages())) {
            chatRequest.setMessages(new ArrayList<>());
        }
        String url = getAccount().serverUrl;
        if(!url.endsWith("/")){
            url = url + "/";
        }
        TokenAuth authCli = new TokenAuth(getToken(requestId));
        // Init the Coze client through the access_token.
        CozeAPI coze =
                new CozeAPI.Builder()
                        .baseURL(url)
                        .auth(authCli)
                        .readTimeout(10000)
                        .build();

        /*
         * Step one, create chat
         * Call the coze.chat().stream() method to create a chat. The create method is a streaming
         * chat and will return a Flowable ChatEvent. Developers should iterate the iterator to get
         * chat event and handle them.
         * */
        String botID = ((CozeAccount)getAccount()).getBotId();
        CreateChatReq req =
                CreateChatReq.builder()
                        .botID(botID)
                        .userID(requestId)
                        .messages(Collections.singletonList(Message.buildUserQuestionText(question)))
                        .build();
        if(!StringUtils.isEmpty(conversationId)) {
            req.setConversationID(conversationId);
        }

        JSONObject userMessage = new JSONObject();
        userMessage.put("role",  "user");
        userMessage.put("content",  question);
        userMessage.put("content_type", "text");
        chatRequest.getMessages().add(userMessage);
        Flowable<ChatEvent> resp = coze.chat().stream(req);

        StringBuilder responseBuilder = new StringBuilder();
        resp.subscribeOn(Schedulers.io())

                .subscribe(
                        event -> {
                            if (ChatEventType.CONVERSATION_MESSAGE_DELTA.equals(event.getEvent())) {
                                System.out.print(event.getMessage().getContent());

                                if(StringUtils.isEmpty(conversationId)){
                                    conversationId = event.getMessage().getConversationId();
                                    log.info("{} coze chat conversation_id = {}", requestId, conversationId);
                                }
                                String tmpText = event.getMessage().getContent().trim();
                                responseBuilder.append(tmpText);
                                JSONObject rsponseData = new JSONObject();
                                rsponseData.put("requestId", requestId);
                                rsponseData.put("code", "200");
                                rsponseData.put("content", tmpText);
                                chatRequest.setConversationId(conversationId);
                                rsponseData.put("messages", chatRequest.getMessages());
                                rsponseData.put("conversationId", chatRequest.getConversationId());
                                rsponseData.put("msgType", "streamMsgSlice");
                                log.info(JSONObject.toJSONString(rsponseData));
                                emitter.send(SseEmitter.event()
                                        .name("streamMsgSlice")
                                        .data(JSONObject.toJSONString(rsponseData)));

                            }
                        },
                        throwable -> {
                            System.err.println(": " + throwable.getMessage());
                            log.error("{} coze error occurred {} {}", requestId, throwable.toString(),
                                    CommonUtils.getStackTraceString(throwable.getStackTrace()));
                            release();
                        },
                        () -> {
                            log.info("{} coze talk done.", requestId);
                            release();
                        });

        acquire();
        coze.shutdownExecutor();


        JSONObject rsponseData = new JSONObject();
        rsponseData.put("requestId", requestId);
        rsponseData.put("code", "200");
        rsponseData.put("content", responseBuilder.toString());
        JSONObject assistantMessage = new JSONObject();
        assistantMessage.put("role",  "assistant");
        assistantMessage.put("content",  responseBuilder.toString());
        assistantMessage.put("content_type", "text");
        chatRequest.getMessages().add(assistantMessage);
        chatRequest.setConversationId(conversationId);
        rsponseData.put("messages", chatRequest.getMessages());
        rsponseData.put("conversationId", chatRequest.getConversationId());
        rsponseData.put("msgType", "message");
        log.info(JSONObject.toJSONString(rsponseData));
        emitter.send(SseEmitter.event()
                .name("message")
                .data(JSONObject.toJSONString(rsponseData)));
    }


    private String getToken(String requestId){
        String cozeTokenType =  ((CozeAccount)getAccount()).getTokenType();
        if(COZE_TOKEN_TYPE_PAT.equalsIgnoreCase(cozeTokenType)){
            return ((CozeAccount)getAccount()).getPatToken();
        }

        String  cozeAPIBase = "https://api.coze.cn";
        String jwtOauthClientID = ((CozeAccount)getAccount()).getOauthClientId();
        String jwtOauthPrivateKey = ((CozeAccount)getAccount()).getOauthPrivateKey();
        String jwtOauthPublicKeyID = ((CozeAccount)getAccount()).getOauthPublicKeyId();
        JWTOAuthClient oauth = null;

        try {
            oauth = new JWTOAuthClient.JWTOAuthBuilder()
                    .clientID(jwtOauthClientID)
                    .privateKey(jwtOauthPrivateKey)
                    .publicKey(jwtOauthPublicKeyID)
                    .baseURL(cozeAPIBase).ttl(24 * 3600)
                    .build();
        } catch (Throwable e) {
            log.error("{} coze getToken error: {} {} ", requestId, e.toString(), CommonUtils.getStackTraceString(e.getStackTrace()));
            return "";
        }

        Long expiredMillsLeft = expireTime - (CurrentTimeMillisClock.now()/1000);
        //过期前1小时
        if (StringUtils.isEmpty(token) || expiredMillsLeft < 3600 || expireTime == 0) {
            synchronized (jwtOauthClientID.intern()) {
                if (StringUtils.isEmpty(token) || expiredMillsLeft < 3600 || expireTime == 0) {
                    OAuthToken aAuthToken = null;
                    try {
                        aAuthToken = oauth.getAccessToken();
                        expireTime = aAuthToken.getExpiresIn();
                        token = aAuthToken.getAccessToken();
                        log.info("{} successfully getAccessToken={} , expireTime={} ",
                                requestId, aAuthToken.getAccessToken().substring(0, 30) + "***", aAuthToken.getExpiresIn());
                    } catch (Throwable e) {
                        log.error("{} coze getAccessToken error: {} ", requestId, CommonUtils.getStackTraceString(e.getStackTrace()));
                        return "";
                    }
                }
            }
        }
        return token;
    }


    private  final Object waitHandle = new Object();
    private  void release(){
        synchronized (waitHandle) {
            waitHandle.notify();
        }
    }
    private  void acquire(){
        try {
            synchronized (waitHandle) {
                waitHandle.wait();
            }
        }
        catch (Throwable throwable){
            log.error("acquire error: {} {} ", throwable.toString() ,
                    CommonUtils.getStackTraceString(throwable.getStackTrace()));
        }
    }
}
