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


import cn.iocoder.yudao.framework.common.pojo.PageResult;
import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
import cn.iocoder.yudao.module.llm.constant.DataConstants;
import cn.iocoder.yudao.module.llm.controller.admin.dataset.dto.DataJsonTemplate;
import cn.iocoder.yudao.module.llm.controller.admin.dataset.vo.*;
import cn.iocoder.yudao.module.llm.dal.dataobject.dataset.DatasetAnswerDO;
import cn.iocoder.yudao.module.llm.dal.dataobject.dataset.DatasetDO;
import cn.iocoder.yudao.module.llm.dal.dataobject.dataset.DatasetFilesDO;
import cn.iocoder.yudao.module.llm.dal.dataobject.dataset.DatasetQuestionDO;
import cn.iocoder.yudao.module.llm.dal.mysql.dataset.DatasetAnswerMapper;
import cn.iocoder.yudao.module.llm.dal.mysql.dataset.DatasetFilesMapper;
import cn.iocoder.yudao.module.llm.dal.mysql.dataset.DatasetMapper;
import cn.iocoder.yudao.module.llm.dal.mysql.dataset.DatasetQuestionMapper;
import cn.iocoder.yudao.module.llm.utils.DataSetReadFileUtils;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.stereotype.Service;
import org.springframework.validation.annotation.Validated;

import javax.annotation.Resource;

import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
import static cn.iocoder.yudao.module.llm.enums.ErrorCodeConstants.DATASET_NOT_EXISTS;

/**
 * 数据集 Service 实现类
 *
 * @author 华大大模型
 */
@Service
@Validated
public class DatasetServiceImpl implements DatasetService {

    @Resource
    private DatasetMapper datasetMapper;
    @Resource
    private DatasetFilesMapper datasetFilesMapper;
    @Resource
    private DatasetQuestionMapper datasetQuestionMapper;
    @Resource
    private DatasetAnswerMapper datasetAnswerMapper;

    @Override
    public Long createDataset(DatasetSaveReqVO createReqVO) {
        if(createReqVO.getType() == null){
            createReqVO.setType(0);
        }
        // 插入
        DatasetDO dataset = BeanUtils.toBean(createReqVO, DatasetDO.class);
        datasetMapper.insert(dataset);
        List<DatasetFilesSaveReqVO> datasetFiles = createReqVO.getDatasetFiles();
        if (CollectionUtils.isNotEmpty(datasetFiles)){
            datasetFiles.stream().forEach(
                    datasetFilesSaveReqVO -> {
                        datasetFilesSaveReqVO.setDatasetId(dataset.getId());
                    }
            );
            parseFile(datasetFiles);
            Long count = datasetQuestionMapper.selectCount(new LambdaQueryWrapper<DatasetQuestionDO>()
                    .eq(DatasetQuestionDO::getDatasetId, dataset.getId()));
            dataset.setDataLength(count);
            datasetMapper.updateById(dataset);
        }
        return dataset.getId();
    }


    private static long getFileContentLength(File file)throws IOException {
        FileInputStream fis = new FileInputStream(file);
        byte[] buffer = new byte[1024];
        long charCount =0;
        while(fis.read(buffer)!=-1) {
            charCount += new String(buffer).length();
        }
        fis.close();
        return charCount;
    }

    @Override
    public void updateDataset(DatasetSaveReqVO updateReqVO) {
        // 校验存在
        validateDatasetExists(updateReqVO.getId());
        // 更新
        DatasetDO updateObj = BeanUtils.toBean(updateReqVO, DatasetDO.class);
        List<DatasetFilesSaveReqVO> datasetFiles = updateReqVO.getDatasetFiles();
        if (CollectionUtils.isNotEmpty(datasetFiles)){
            datasetFiles.stream().forEach(
                    datasetFilesSaveReqVO -> {
                        datasetFilesSaveReqVO.setDatasetId(updateObj.getId());
                    }
            );
            parseFile(datasetFiles);
            Long count = datasetQuestionMapper.selectCount(new LambdaQueryWrapper<DatasetQuestionDO>()
                    .eq(DatasetQuestionDO::getDatasetId, updateObj.getId()));
            updateObj.setDataLength(count);
            Long annoCount = datasetQuestionMapper.selectCount(new LambdaQueryWrapper<DatasetQuestionDO>()
                    .eq(DatasetQuestionDO::getDatasetId, updateObj.getId())
                    .eq(DatasetQuestionDO::getStatus,2));
            double ratio = count == 0 ? 0 : ((double) annoCount / count) *100;
            Integer formattedRatio = ratio == 0 ? 0 : (int) ratio;
            if (formattedRatio != null){
                updateObj.setAnnotateProgress(formattedRatio);
            }
        }
        datasetMapper.updateById(updateObj);
    }

    @Override
    public void deleteDataset(Long id) {
        // 校验存在
        validateDatasetExists(id);
        // 删除
        datasetMapper.deleteById(id);
    }

    private void validateDatasetExists(Long id) {
        if (datasetMapper.selectById(id) == null) {
            throw exception(DATASET_NOT_EXISTS);
        }
    }

    @Override
    public DatasetRespVO getDataset(Long id) {
        DatasetDO datasetDO = datasetMapper.selectById(id);
        DatasetRespVO datasetRespVO = BeanUtils.toBean(datasetDO, DatasetRespVO.class);
        /*List<DatasetQuestionDO> datasetQuestionDO = datasetQuestionMapper.selectList(new LambdaQueryWrapper<DatasetQuestionDO>().eq(DatasetQuestionDO::getDatasetId, id));
        List<DatasetQuestionRespVO> datasetQuestionRespVOS = BeanUtils.toBean(datasetQuestionDO, DatasetQuestionRespVO.class);
        datasetRespVO.setDatasetQuestionRespVOS(datasetQuestionRespVOS);*/
        return datasetRespVO;
    }

    @Override
    public PageResult<DatasetDO> getDatasetPage(DatasetPageReqVO pageReqVO) {
        return datasetMapper.selectPage(pageReqVO);
    }

    @Override
    public List<DatasetTreeNode>  queryAll() {
        /*List<DatasetDO> datasetDOS0 = datasetMapper.selectList(new LambdaQueryWrapper<DatasetDO>().eq(DatasetDO::getType, DataConstants.dataTypePrivate));
        List<DatasetRespVO> datasetRespVOS0 = BeanUtils.toBean(datasetDOS0, DatasetRespVO.class);
        List<DatasetDO> datasetDOS1 = datasetMapper.selectList(new LambdaQueryWrapper<DatasetDO>().eq(DatasetDO::getType, DataConstants.dataTypePublic));
        List<DatasetRespVO> datasetRespVOS1 = BeanUtils.toBean(datasetDOS1, DatasetRespVO.class);
        List<List<DatasetRespVO>> result = new ArrayList<>();
        result.add(datasetRespVOS0);
        result.add(datasetRespVOS1);*/
        List<DatasetDO> datasetDOS = datasetMapper.selectList(null); // 获取所有数据集

        // 创建两个根节点，分别代表两种 type
        DatasetTreeNode privateRoot = new DatasetTreeNode(DataConstants.dataTypePrivate);
        DatasetTreeNode publicRoot = new DatasetTreeNode(DataConstants.dataTypePublic);

        for (DatasetDO datasetDO : datasetDOS) {
            DatasetRespVO datasetRespVO = BeanUtils.toBean(datasetDO, DatasetRespVO.class);
            // 根据 type 字段决定节点的位置
            if (datasetRespVO.getType()==DataConstants.dataTypePrivate) {
                privateRoot.getChildren().add(datasetRespVO);
            } else if (datasetRespVO.getType()==DataConstants.dataTypePublic) {
                publicRoot.getChildren().add(datasetRespVO);
            }
        }
        List<DatasetTreeNode> root = new ArrayList<>();
        root.add(privateRoot);
        root.add(publicRoot);
        return root;
    }


    public void readJsonFile(List<DatasetFilesDO> jsonFiles){
        jsonFiles.forEach(datasetFilesDO -> {
            HttpURLConnection connection = DataSetReadFileUtils.readFile(datasetFilesDO.getDatasetFileUrl());
            if (connection != null){
                try (BufferedReader in = new BufferedReader(new InputStreamReader(connection.getInputStream(), "UTF-8"))) {
                    StringBuilder content = new StringBuilder();
                    String line;
                    while ((line = in.readLine()) != null) {
                        content.append(line);
                    }
                    // 使用Jackson解析 Json 字符串为List<String>对象
                    // 使用Jackson解析 Json 字符串为List<String>对象
                    ObjectMapper mapper = new ObjectMapper();
                    // 使用 TypeReference 解析 JSON 字符串为 List<String>
                    List<DataJsonTemplate> jsonList = mapper.readValue(content.toString(), new TypeReference<List<DataJsonTemplate>>() {});
                    jsonList.forEach(
                            dataJsonTemplate -> {
                                DatasetQuestionDO datasetQuestionDO = BeanUtils.toBean(dataJsonTemplate, DatasetQuestionDO.class);
                                datasetQuestionDO.setDatasetId(datasetFilesDO.getDatasetId());
                                datasetQuestionDO.setDatasetFilesId(datasetFilesDO.getId());
                                datasetQuestionMapper.insert(datasetQuestionDO);
                                List<String> answers = dataJsonTemplate.getAnswers();
                                if (CollectionUtils.isNotEmpty(answers)){
                                    for (String answer : answers) {
                                        DatasetAnswerDO datasetAnswerDO = new DatasetAnswerDO();
                                        datasetAnswerDO.setDatasetId(datasetFilesDO.getDatasetId());
                                        datasetAnswerDO.setDatasetFilesId(datasetFilesDO.getId());
                                        datasetAnswerDO.setQuestionId(datasetQuestionDO.getId());
                                        datasetAnswerDO.setAnswer(answer);
                                        datasetAnswerMapper.insert(datasetAnswerDO);
                                    }
                                }
                            }
                    );
                }catch (Exception e){
                    throw new RuntimeException("请正确上传json格式得数据！！！");
                }finally {
                    connection.disconnect();
                }
            }
        });
    }
    /**
     * txt文本数据
     * @param txtFiles
     */
    public void readTxtFile(List<DatasetFilesDO> txtFiles){
        txtFiles.forEach(datasetFilesDO -> {
            List<String> newContent = new ArrayList<>();

            HttpURLConnection connection = DataSetReadFileUtils.readFile(datasetFilesDO.getDatasetFileUrl());
            if (connection != null){
                try (BufferedReader in = new BufferedReader(new InputStreamReader(connection.getInputStream(), "UTF-8"))) {
                    String inputLine;
                    while ((inputLine = in.readLine()) != null) {
                        inputLine = inputLine.trim(); // 去除行首和行尾的空白字符
                        if (!inputLine.isEmpty()) {
                            DatasetQuestionDO datasetQuestionDO = new DatasetQuestionDO();// 检查是否为空行
                            datasetQuestionDO.setDatasetId(datasetFilesDO.getDatasetId());
                            datasetQuestionDO.setDatasetFilesId(datasetFilesDO.getId());
                            datasetQuestionDO.setQuestion(inputLine);
                            datasetQuestionMapper.insert(datasetQuestionDO);
                        }
                    }
                }catch (Exception e){
                    e.printStackTrace();
                }finally {
                    connection.disconnect();
                }
            }
        });
    }

    public void parseFile(List<DatasetFilesSaveReqVO> datasetFiles) {
        List<DatasetFilesDO> insertDatasetFiles = BeanUtils.toBean(datasetFiles, DatasetFilesDO.class);
        datasetFilesMapper.insertBatch(insertDatasetFiles, 100);
        // 提取文件
        List<DatasetFilesDO> jsonFiles = insertDatasetFiles.stream()
                .filter(datasetFilesDO -> datasetFilesDO.getDatasetFileUrl().toLowerCase().endsWith(".json"))
                .collect(Collectors.toList());
        if (CollectionUtils.isNotEmpty(jsonFiles)){
            readJsonFile(jsonFiles);
        }

        List<DatasetFilesDO> txtFiles = insertDatasetFiles.stream()
                .filter(datasetFilesDO -> datasetFilesDO.getDatasetFileUrl().toLowerCase().endsWith(".txt"))
                .collect(Collectors.toList());
        if (CollectionUtils.isNotEmpty(txtFiles)){
            readTxtFile(txtFiles);
        }
    }
}
