package com.moyz.adi.common.interfaces;

import com.moyz.adi.common.exception.BaseException;
import com.moyz.adi.common.util.JsonUtil;
import com.moyz.adi.common.util.LocalCache;
import com.moyz.adi.common.vo.AnswerMeta;
import com.moyz.adi.common.vo.ChatMeta;
import com.moyz.adi.common.vo.PromptMeta;
import com.moyz.adi.common.vo.SseAskParams;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.TokenStream;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.net.Proxy;
import java.util.UUID;

import static com.moyz.adi.common.enums.ErrorEnum.B_LLM_SERVICE_DISABLED;

@Slf4j
public abstract class AbstractLLMService<T> {

    protected Proxy proxy;

    protected String modelName;

    protected T setting;

    protected StreamingChatLanguageModel streamingChatLanguageModel;
    protected ChatLanguageModel chatLanguageModel;

    public AbstractLLMService(String modelName, String settingName, Class<T> clazz, Proxy proxy) {
        this.modelName = modelName;
        this.proxy = proxy;
        String st = LocalCache.CONFIGS.get(settingName);
        setting = JsonUtil.fromJson(st, clazz);
    }

    /**
     * 检测该service是否可用（不可用的情况通常是没有配置key）
     *
     * @return
     */
    public abstract boolean isEnabled();

    public ChatLanguageModel getChatLLM() {
        if (null != chatLanguageModel) {
            return chatLanguageModel;
        }
        chatLanguageModel = buildChatLLM();
        return chatLanguageModel;
    }

    public StreamingChatLanguageModel getStreamingChatLLM() {
        if (null != streamingChatLanguageModel) {
            return streamingChatLanguageModel;
        }
        streamingChatLanguageModel = buildStreamingChatLLM();
        return streamingChatLanguageModel;
    }

    protected abstract ChatLanguageModel buildChatLLM();

    protected abstract StreamingChatLanguageModel buildStreamingChatLLM();

    protected abstract String parseError(Object error);

    public Response<AiMessage> chat(ChatMessage chatMessage) {
        if(!isEnabled()){
            log.error("llm service is disabled");
            throw new BaseException(B_LLM_SERVICE_DISABLED);
        }
        return getChatLLM().generate(chatMessage);
    }

    public void sseChat(SseAskParams params, TriConsumer<String, PromptMeta, AnswerMeta> consumer) {
        if(!isEnabled()){
            log.error("llm service is disabled");
            throw new BaseException(B_LLM_SERVICE_DISABLED);
        }
        //create chat assistant
        AiServices<IChatAssistant> serviceBuilder = AiServices.builder(IChatAssistant.class)
                .streamingChatLanguageModel(getStreamingChatLLM());
        if (null != params.getChatMemory()) {
            serviceBuilder.chatMemory(params.getChatMemory());
        }
        IChatAssistant chatAssistant = serviceBuilder.build();

        TokenStream tokenStream;
        if (StringUtils.isNotBlank(params.getSystemMessage())) {
            tokenStream = chatAssistant.chat(params.getSystemMessage(), params.getUserMessage());
        } else {
            tokenStream = chatAssistant.chat(params.getUserMessage());
        }
        tokenStream.onNext((content) -> {
                    log.info("get content:{}", content);
                    //加空格配合前端的fetchEventSource进行解析，见https://github.com/Azure/fetch-event-source/blob/45ac3cfffd30b05b79fbf95c21e67d4ef59aa56a/src/parse.ts#L129-L133
                    try {
                        params.getSseEmitter().send(" " + content);
                    } catch (IOException e) {
                        log.error("stream onNext error", e);
                    }
                })
                .onComplete((response) -> {
                    log.info("返回数据结束了:{}", response);
                    String questionUuid = StringUtils.isNotBlank(params.getRegenerateQuestionUuid()) ? params.getRegenerateQuestionUuid() : UUID.randomUUID().toString().replace("-", "");
                    PromptMeta questionMeta = new PromptMeta(response.tokenUsage().inputTokenCount(), questionUuid);
                    AnswerMeta answerMeta = new AnswerMeta(response.tokenUsage().outputTokenCount(), UUID.randomUUID().toString().replace("-", ""));
                    ChatMeta chatMeta = new ChatMeta(questionMeta, answerMeta);
                    String meta = JsonUtil.toJson(chatMeta).replaceAll("\r\n", "");
                    log.info("meta:" + meta);
                    try {
                        params.getSseEmitter().send(" [META]" + meta);
                    } catch (IOException e) {
                        log.error("stream onComplete error", e);
                        throw new RuntimeException(e);
                    }
                    // close eventSourceEmitter after tokens was calculated
                    params.getSseEmitter().complete();
                    consumer.accept(response.content().text(), questionMeta, answerMeta);
                })
                .onError((error) -> {
                    log.error("stream error", error);
                    try {
                        String errorMsg = parseError(error);
                        if(StringUtils.isBlank(errorMsg)){
                            errorMsg = error.getMessage();
                        }
                        params.getSseEmitter().send(SseEmitter.event().name("error").data(errorMsg));
                    } catch (IOException e) {
                        log.error("sse error", e);
                    }
                    params.getSseEmitter().complete();
                })
                .start();
    }

}

dev.langchain4j.service.AiServices.builder
[((1439, 1474), 'com.moyz.adi.common.util.LocalCache.CONFIGS.get'), ((2985, 3092), 'dev.langchain4j.service.AiServices.builder'), ((4357, 4402), 'java.util.UUID.randomUUID'), ((4357, 4385), 'java.util.UUID.randomUUID'), ((4622, 4667), 'java.util.UUID.randomUUID'), ((4622, 4650), 'java.util.UUID.randomUUID'), ((4784, 4832), 'com.moyz.adi.common.util.JsonUtil.toJson'), ((5780, 5827), 'org.springframework.web.servlet.mvc.method.annotation.SseEmitter.event'), ((5780, 5812), 'org.springframework.web.servlet.mvc.method.annotation.SseEmitter.event')]