/*
 * 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 com.pavis.admin.aigc.service.impl;

import cn.hutool.core.util.StrUtil;
import com.pavis.admin.aigc.common.model.req.ChatReq;
import com.pavis.admin.aigc.common.model.resp.ChatResp;
import com.pavis.admin.aigc.common.utils.StreamEmitter;
import com.pavis.admin.aigc.enums.RoleEnum;
import com.pavis.admin.aigc.mapper.AigcMessageMapper;
import com.pavis.admin.aigc.mapper.AppKnowledgeMapper;
import com.pavis.admin.aigc.model.entity.AigcAppDO;
import com.pavis.admin.aigc.model.entity.AigcMessageDO;
import com.pavis.admin.aigc.service.AigcAppService;
import com.pavis.admin.aigc.service.AigcMessageService;
import com.pavis.admin.aigc.service.PvsChatService;
import dev.langchain4j.model.output.TokenUsage;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import org.springframework.stereotype.Service;

import java.util.List;

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

    private final AigcAppService aigcAppService;
    private final AppKnowledgeMapper appKnowledgeMapper;
    private final AigcMessageMapper aigcMessageMapper;

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

        if (StrUtil.isNotBlank(req.getAppId())) {
            AigcAppDO app = aigcAppService.get(req.getAppId());
            if (app != null) {
                req.setModelId(app.getModelId().toString());
                req.setPromptText(app.getPrompt());
                List<Long> longs = appKnowledgeMapper.selectByKnowledgeIds(app.getId());
                req.setKnowledgeIds(longs);
            }
        }

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

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

                        // save assistant message
                        req.setMessage(text.toString());
                        req.setRole(RoleEnum.ASSISTANT.getName());

                    }).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) {
            AigcMessageDO message = new AigcMessageDO();
            BeanUtils.copyProperties(req, message);
            // message.setIp(ServletUtil.getIpAddr());
            // message.setPromptTokens(inputToken);
            message.setContent(req.getMessage());
            message.setUsageToken(outputToken.toString());
            // aigcMessageService.addMessage(message);
            aigcMessageMapper.insert(message);
        }
    }

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

    // @Override
    // public AigcOss image(ImageReq req) {
    //     Response<Image> res = aigcAppService.image(req);
    //
    //     String path = res.content().url().toString();
    //     AigcOss oss = new AigcOss();
    //     oss.setUrl(path);
    //     return oss;
    // }
}
