package cn.iocoder.yudao.module.llm.service.async;

import cn.hutool.json.JSONObject;
import cn.iocoder.yudao.framework.common.util.collection.CollectionUtils;
import cn.iocoder.yudao.module.llm.controller.admin.dataset.vo.DatasetQuestionRespVO;
import cn.iocoder.yudao.module.llm.dal.dataobject.dataset.DatasetDO;
import cn.iocoder.yudao.module.llm.dal.mysql.dataset.DatasetMapper;
import cn.iocoder.yudao.module.llm.service.dataset.vo.AigcDatasetVo;
import cn.iocoder.yudao.module.llm.service.http.TrainHttpService;
import cn.iocoder.yudao.module.llm.service.http.vo.AigcDatasetFileRespV0;
import cn.iocoder.yudao.module.llm.service.http.vo.ModelCompletionsReqVO;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

@Service
public class AsyncDataSetService {
    @Resource
    private TrainHttpService trainHttpService;
    @Resource
    private DatasetMapper datasetMapper;
    private static final Logger log = LoggerFactory.getLogger(AsyncKnowledgeBase.class);

    @Async
    public void JsonFileWrite(DatasetDO datasetDO,List<DatasetQuestionRespVO> datasetQuestionList) {
        StringBuilder sb = new StringBuilder();
        datasetQuestionList.forEach(dataSource -> {
            List<ModelCompletionsReqVO.ModelCompletionsMessage> messages = new ArrayList<>();
            ModelCompletionsReqVO.ModelCompletionsMessage questionMessage = new ModelCompletionsReqVO.ModelCompletionsMessage();
            questionMessage.setRole("user");
            questionMessage.setContent(StringUtils.isNotBlank(dataSource.getQuestion())?dataSource.getQuestion():"");
            messages.add(questionMessage);
            ModelCompletionsReqVO.ModelCompletionsMessage answerMessage = new ModelCompletionsReqVO.ModelCompletionsMessage();
            answerMessage.setRole("assistant");
            if (!CollectionUtils.isAnyEmpty(dataSource.getDatasetAnswerRespVO())){
                answerMessage.setContent(StringUtils.isNotBlank(dataSource.getDatasetAnswerRespVO().get(0).getAnswer())?dataSource.getDatasetAnswerRespVO().get(0).getAnswer():"");
            }else {
                answerMessage.setContent("");
            }
            messages.add(answerMessage);
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("messages", messages);
            sb.append(jsonObject.toString()).append("\n");

        });
        ObjectMapper mapper = new ObjectMapper();
        try (InputStream inputStream = new ByteArrayInputStream(sb.toString().getBytes())){
            AigcDatasetFileRespV0 aigcDatasetFileRespV0 = trainHttpService.AigcUploadFile(new HashMap<>(), inputStream, datasetDO.getDatasetName() + datasetDO.getId() + ".json");
            if (aigcDatasetFileRespV0 != null){
                datasetMapper.setJobid(datasetDO.getId(),aigcDatasetFileRespV0.getFileId());

                String s3Url = aigcDatasetFileRespV0.getS3Url();
                int lastIndex = s3Url.lastIndexOf("/storage");

                //todo 1111
                String url = s3Url.substring(lastIndex + 1);
                datasetMapper.setUrl(datasetDO.getId(),url);

                log.info("[JsonFileWrite][写入文件成功]");
            }
        } catch (IOException e) {
            // 记录异常信息
            log.error("[JsonFileWrite][写入文件失败] {}", e.getMessage());
        }
    }
}
