/*
 * Copyright (c) 2024 LangChat. TyCoding All Rights Reserved.
 *
 * Licensed under the GNU Affero General Public License, Version 3 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.gnu.org/licenses/agpl-3.0.html
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package cn.iocoder.yudao.module.langchat.service.core.impl;

import cn.hutool.core.util.StrUtil;
import cn.iocoder.yudao.module.langchat.dal.dto.ChatReq;
import cn.iocoder.yudao.module.langchat.dal.dto.ChatRes;
import cn.iocoder.yudao.module.langchat.dal.dto.ImageR;
import cn.iocoder.yudao.module.langchat.entity.AigcApp;
import cn.iocoder.yudao.module.langchat.entity.AigcMessage;
import cn.iocoder.yudao.module.langchat.entity.AigcOss;
import cn.iocoder.yudao.module.langchat.enums.RoleEnum;
import cn.iocoder.yudao.module.langchat.service.aigc.AigcMessageService;
import cn.iocoder.yudao.module.langchat.service.core.ChatService;
import cn.iocoder.yudao.module.langchat.service.core.LangChatService;
import cn.iocoder.yudao.module.langchat.store.AppStore;
import cn.iocoder.yudao.module.langchat.utils.ServletUtil;
import cn.iocoder.yudao.module.langchat.utils.StreamEmitter;
import dev.langchain4j.data.image.Image;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.model.output.TokenUsage;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import org.springframework.stereotype.Service;

/**
 * @author tycoding
 * @since 2024/1/4
 */
@Slf4j
@Service
@AllArgsConstructor
public class ChatServiceImpl implements ChatService {

    private final LangChatService langChatService;
    private final AigcMessageService aigcMessageService;
    private final AppStore appStore;

    @Override
    public void chat(ChatReq req) {
        StreamEmitter emitter = req.getEmitter();
        long startTime = System.currentTimeMillis();
        StringBuilder text = new StringBuilder();

        if (StrUtil.isNotBlank(req.getAppId())) {
            AigcApp app = appStore.get(req.getAppId());
            if (app != null) {
                req.setModelId(app.getModelId());
                req.setPromptText(app.getPrompt());
                req.setKnowledgeIds(app.getKnowledgeIds());
            }
        }

        // save user message
        req.setRole(RoleEnum.USER.getName());
        saveMessage(req, 0, 0);

        try {
            langChatService
                    .chat(req)
                    .onNext(e -> {
                        text.append(e);
                        emitter.send(new ChatRes(e));
                    })
                    .onComplete((e) -> {
                        TokenUsage tokenUsage = e.tokenUsage();
                        ChatRes res = new ChatRes(tokenUsage.totalTokenCount(), startTime);
                        emitter.send(res);
                        emitter.complete();

                        // save assistant message
                        req.setMessage(text.toString());
                        req.setRole(RoleEnum.ASSISTANT.getName());
                        saveMessage(req, tokenUsage.inputTokenCount(), tokenUsage.outputTokenCount());
                    })
                    .onError((e) -> {
                        emitter.error(e.getMessage());
                        throw new RuntimeException(e.getMessage());
                    })
                    .start();
        } catch (Exception e) {
            e.printStackTrace();
            emitter.error(e.getMessage());
            throw new RuntimeException(e.getMessage());
        }
    }

    private void saveMessage(ChatReq req, Integer inputToken, Integer outputToken) {
        if (req.getConversationId() != null) {
            AigcMessage message = new AigcMessage();
            BeanUtils.copyProperties(req, message);
            message.setIp(ServletUtil.getIpAddr());
            message.setPromptTokens(inputToken);
            message.setTokens(outputToken);
            aigcMessageService.addMessage(message);
        }
    }

    @Override
    public String text(ChatReq req) {
        String text;
        try {
            text = langChatService.text(req);
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(e.getMessage());
        }
        return text;
    }

    @Override
    public AigcOss image(ImageR req) {
        Response<Image> res = langChatService.image(req);

        String path = res.content().url().toString();
        AigcOss oss = new AigcOss();
        oss.setUrl(path);
        return oss;
    }
}
