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

import cn.iocoder.yudao.framework.common.exception.ErrorCode;
import cn.iocoder.yudao.framework.common.pojo.PageResult;
import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
import cn.iocoder.yudao.module.llm.controller.admin.finetuningtask.vo.FineTuningTaskPageReqVO;
import cn.iocoder.yudao.module.llm.controller.admin.finetuningtask.vo.FineTuningTaskRespVO;
import cn.iocoder.yudao.module.llm.controller.admin.finetuningtask.vo.FineTuningTaskSaveReqVO;
import cn.iocoder.yudao.module.llm.dal.dataobject.basemodel.BaseModelDO;
import cn.iocoder.yudao.module.llm.dal.dataobject.dataset.DatasetDO;
import cn.iocoder.yudao.module.llm.dal.dataobject.finetuningtask.FineTuningTaskDO;
import cn.iocoder.yudao.module.llm.dal.mysql.basemodel.BaseModelMapper;
import cn.iocoder.yudao.module.llm.dal.mysql.dataset.DatasetMapper;
import cn.iocoder.yudao.module.llm.dal.mysql.finetuningtask.FineTuningTaskMapper;
import cn.iocoder.yudao.module.llm.enums.FinetuningTaskStatusEnum;
import cn.iocoder.yudao.module.llm.service.async.AsyncFineTuningTaskService;
import cn.iocoder.yudao.module.llm.service.http.TrainHttpService;
import cn.iocoder.yudao.module.llm.service.http.vo.AigcFineTuningDetailRespVO;
import cn.iocoder.yudao.module.llm.service.modelservice.ModelServiceService;
import com.alibaba.fastjson.JSON;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import jodd.util.StringUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;
import org.springframework.validation.annotation.Validated;

import javax.annotation.Resource;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
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.FINE_TUNING_TASK_NAME_NOT_EXISTS;
import static cn.iocoder.yudao.module.llm.enums.ErrorCodeConstants.FINE_TUNING_TASK_NOT_EXISTS;


/**
 * 模型调优 —— 微调任务 Service 实现类
 *
 * @author 华大大模型
 */
@Service
@Validated
public class FineTuningTaskServiceImpl implements FineTuningTaskService {

    private static final Logger log = LoggerFactory.getLogger(FineTuningTaskServiceImpl.class);
    @Resource
    private FineTuningTaskMapper fineTuningTaskMapper;
    @Resource
    private BaseModelMapper baseModelMapper;
    @Resource
    private DatasetMapper datasetMapper;
    @Resource
    private AsyncFineTuningTaskService asyncFineTuningTaskService;
    @Resource
    TrainHttpService trainHttpService;

    @Resource
    private ModelServiceService modelServiceService;

    @Override
    public Long createFineTuningTask (FineTuningTaskSaveReqVO createReqVO) {
        // 插入
        validateFineTuningNameTaskExists(createReqVO);
        FineTuningTaskDO fineTuningTask = BeanUtils.toBean(createReqVO, FineTuningTaskDO.class);
        fineTuningTask.setStatus(FinetuningTaskStatusEnum.WAITING.getStatus());
        fineTuningTaskMapper.insert(fineTuningTask);
        //todo 调用模型服务，创建调优任务
        asyncFineTuningTaskService.createTuning(fineTuningTask);
        // 返回
        return fineTuningTask.getId();
    }

    @Override
    public void updateFineTuningTask (FineTuningTaskSaveReqVO updateReqVO) {
        // 校验存在
        validateFineTuningTaskExists(updateReqVO.getId());
        validateFineTuningNameTaskExists(updateReqVO);
        // 更新
        FineTuningTaskDO updateObj = BeanUtils.toBean(updateReqVO, FineTuningTaskDO.class);
        fineTuningTaskMapper.updateById(updateObj);
    }

    @Override
    public void deleteFineTuningTask (Long id) {
        // 校验存在
        validateFineTuningTaskExists(id);
        // 校验使用
        validateFineTuningTaskUse(id);
        FineTuningTaskDO fineTuningTaskDO = fineTuningTaskMapper.selectById(id);
        // 删除
        LambdaUpdateWrapper<FineTuningTaskDO> wrapper = new LambdaUpdateWrapper<>();
        wrapper.eq(FineTuningTaskDO::getId, id)
                .set(FineTuningTaskDO::getDeleted, true);
        fineTuningTaskMapper.update(null, wrapper);
        asyncFineTuningTaskService.stopFineTuningTask(fineTuningTaskDO);
        //        fineTuningTaskMapper.deleteById(id);
    }

    /**
     * 校验使用
     * <p>
     * 删除调优任务时，需要判断是否正在使用，如果正在使用，则不允许删除
     * </p>
     *
     * @param id 任务id
     */
    private void validateFineTuningTaskUse (Long id) {
        FineTuningTaskDO fineTuningTask = fineTuningTaskMapper.selectById(id);
        String taskName = fineTuningTask.getTaskName();

        // 校验 微调任务 是否在 模型服务管理 有使用
        validateTaskUsesInModelService(id, taskName);
    }

    /**
     * 校验微调任务是否在模型服务管理中使用
     *
     * @param id       任务id
     * @param taskName 任务名称
     */
    private void validateTaskUsesInModelService (Long id, String taskName) {
        Map<Long, String> modelService = modelServiceService.getModelServiceByTaskId(id);
        if (com.baomidou.mybatisplus.core.toolkit.CollectionUtils.isNotEmpty(modelService)) {
            String msg = String.format("微调任务【%s】在模型服务管理 %s 有使用，请先取消使用", taskName, modelService.values());
            throw exception(new ErrorCode(10005_001, msg));
        }
    }

    private FineTuningTaskDO validateFineTuningTaskExists (Long id) {
        FineTuningTaskDO fineTuningTaskDO = fineTuningTaskMapper.selectById(id);
        if (fineTuningTaskDO == null) {
            throw exception(FINE_TUNING_TASK_NOT_EXISTS);
        }
        return fineTuningTaskDO;
    }

    private void validateFineTuningNameTaskExists (FineTuningTaskSaveReqVO reqVO) {
        LambdaQueryWrapper<FineTuningTaskDO> wrapper = new LambdaQueryWrapper<FineTuningTaskDO>()
                .eq(FineTuningTaskDO::getTaskName, reqVO.getTaskName());

        if (reqVO.getId() != null) {
            wrapper.ne(FineTuningTaskDO::getId, reqVO.getId());
        }
        List<FineTuningTaskDO> dineTuningTaskDos = fineTuningTaskMapper.selectList(wrapper);
        if (CollectionUtils.isNotEmpty(dineTuningTaskDos)) {
            throw exception(FINE_TUNING_TASK_NAME_NOT_EXISTS);
        }
    }

    @Override
    public FineTuningTaskRespVO getFineTuningTask (Long id) {
        validateFineTuningTaskExists(id);
        FineTuningTaskDO fineTuningTaskDO = fineTuningTaskMapper.selectById(id);
        FineTuningTaskRespVO result = BeanUtils.toBean(fineTuningTaskDO, FineTuningTaskRespVO.class);
        Long dataset = fineTuningTaskDO.getDataset();
        if (dataset != null) {
            DatasetDO datasetDO = datasetMapper.selectById(dataset);
            if (datasetDO != null) {
                result.setDatasetName(datasetDO.getDatasetName());
            }
        }
        Long baseModelId = fineTuningTaskDO.getBaseModelId();
        if (baseModelId != null) {
            BaseModelDO baseModelDO = baseModelMapper.selectById(baseModelId);
            if (baseModelDO != null) {
                result.setBaseModelName(baseModelDO.getModelName());
            }
        }
        if (StringUtil.isNotEmpty(fineTuningTaskDO.getJobId())) {
            AigcFineTuningDetailRespVO aigcFineTuningDetailRespVO = trainHttpService.finetuningDetail(new HashMap<>(), fineTuningTaskDO.getJobId());
            result.setDetailRespVO(aigcFineTuningDetailRespVO);
        }
        result.setCheckPointDataJSON(JSON.parseObject(fineTuningTaskDO.getCheckPointData()));
        return result;
    }

    @Override
    public PageResult<FineTuningTaskRespVO> getFineTuningTaskPage1 (FineTuningTaskPageReqVO pageReqVO) {
        PageResult<FineTuningTaskDO> pageResult = fineTuningTaskMapper.selectPage(pageReqVO);
        PageResult<FineTuningTaskRespVO> result = BeanUtils.toBean(pageResult, FineTuningTaskRespVO.class);
        if (CollectionUtils.isNotEmpty(pageResult.getList())) {
            List<Long> modelIds = pageResult.getList().stream().map(FineTuningTaskDO::getBaseModelId).collect(Collectors.toList());
            List<BaseModelDO> baseModelDOS = baseModelMapper.selectList(new LambdaQueryWrapper<BaseModelDO>()
                    .in(BaseModelDO::getId, modelIds));
            Map<Long, BaseModelDO> longModelServiceDOMap = cn.iocoder.yudao.framework.common.util.collection.
                    CollectionUtils.convertMap(baseModelDOS, BaseModelDO::getId);
            result.getList().forEach(item -> {
                BaseModelDO baseModelDO = longModelServiceDOMap.get(item.getBaseModelId());
                if (baseModelDO != null) {
                    item.setBaseModelName(baseModelDO.getModelName());
                }
            });
            List<Long> dataSetlIds = pageResult.getList().stream().map(FineTuningTaskDO::getDataset).collect(Collectors.toList());
            List<DatasetDO> datasetDOS = datasetMapper.selectList(new LambdaQueryWrapper<DatasetDO>().in(DatasetDO::getId, dataSetlIds));
            Map<Long, DatasetDO> datasetDOMap = cn.iocoder.yudao.framework.common.util.collection.
                    CollectionUtils.convertMap(datasetDOS, DatasetDO::getId);
            result.getList().forEach(item -> {
                DatasetDO datasetDO = datasetDOMap.get(item.getDataset());
                if (datasetDO != null) {
                    item.setDatasetName(datasetDO.getDatasetName());
                }
            });

        }
        return result;
    }


    @Override
    public void startFineTuningTask (Long id) {
        FineTuningTaskDO fineTuningTaskDO = validateFineTuningTaskExists(id);
        fineTuningTaskMapper.stopStartTask(id, FinetuningTaskStatusEnum.WAITING.getStatus());
        //调用模型服务，开启调优任务
        asyncFineTuningTaskService.startFineTuningTask(fineTuningTaskDO);
    }

    @Override
    public List<FineTuningTaskRespVO> selectAll () {
        List<FineTuningTaskDO> fineTuningTaskDOS = fineTuningTaskMapper.selectList();
        return BeanUtils.toBean(fineTuningTaskDOS, FineTuningTaskRespVO.class);
    }

    /**
     * 获得模型调优 —— 获取已部署完成的调优任务
     *
     * @return 已部署完成的调优任务
     */
    @Override
    public List<FineTuningTaskRespVO> selectEnable () {
        List<FineTuningTaskDO> fineTuningTaskDOS = fineTuningTaskMapper.selectList();
        fineTuningTaskDOS = fineTuningTaskDOS.stream()
                .filter(item -> item.getStatus().equals(FinetuningTaskStatusEnum.FINISHED.getStatus()))
                .collect(Collectors.toList());
        return BeanUtils.toBean(fineTuningTaskDOS, FineTuningTaskRespVO.class);
    }

    /**
     * 根据 基础模型id 获取 模型调优
     *
     * @param baseModelId 基础模型id
     * @return 模型调优
     */
    @Override
    public Map<Long, String> getModelTuningByModelId (Long baseModelId) {

        if (baseModelId == null) {
            return Collections.emptyMap();
        }

        List<FineTuningTaskDO> fineTuningTaskDOS = fineTuningTaskMapper.selectList(new LambdaQueryWrapper<FineTuningTaskDO>()
                .eq(FineTuningTaskDO::getBaseModelId, baseModelId)
                .in(FineTuningTaskDO::getStatus, FinetuningTaskStatusEnum.WAITING.getStatus(), FinetuningTaskStatusEnum.TRAINING.getStatus()));

        if (CollectionUtils.isNotEmpty(fineTuningTaskDOS)) {
            return fineTuningTaskDOS.stream().collect(Collectors.toMap(FineTuningTaskDO::getId, FineTuningTaskDO::getTaskName));
        }

        return Collections.emptyMap();
    }

    /**
     * 根据 数据集id 获取 模型调优
     *
     * @param datasetId 数据集id
     * @return 模型调优
     */
    @Override
    public Map<Long, String> getModelTuningByDatasetId (Long datasetId) {
        if (datasetId == null) {
            return Collections.emptyMap();
        }

        List<FineTuningTaskDO> fineTuningTaskDOS = fineTuningTaskMapper.selectList(new LambdaQueryWrapper<FineTuningTaskDO>()
                .eq(FineTuningTaskDO::getDataset, datasetId)
                .in(FineTuningTaskDO::getStatus, FinetuningTaskStatusEnum.WAITING.getStatus(), FinetuningTaskStatusEnum.TRAINING.getStatus()));

        if (CollectionUtils.isNotEmpty(fineTuningTaskDOS)) {
            return fineTuningTaskDOS.stream().collect(Collectors.toMap(FineTuningTaskDO::getId, FineTuningTaskDO::getTaskName));
        }

        return Collections.emptyMap();
    }

    @Override
    public String getFineTuningLog (Long id) {
        FineTuningTaskDO fineTuningTaskDO = fineTuningTaskMapper.selectById(id);
        if (fineTuningTaskDO != null && StringUtils.isNotBlank(fineTuningTaskDO.getTrainLog())) {
            String query = "?log_path="+fineTuningTaskDO.getTrainLog();
            Map<String, String> map = new HashMap<>();
            map.put("log_path", fineTuningTaskDO.getTrainDuration());
            return trainHttpService.getTrainLog(query);
        }
        return null;
    }

    @Override
    public void stopFineTuningTask (Long id) {
        FineTuningTaskDO fineTuningTaskDO = validateFineTuningTaskExists(id);
        fineTuningTaskMapper.stopStartTask(id, FinetuningTaskStatusEnum.WAITING.getStatus());
        //todo 调用模型服务，停止调优任务
        asyncFineTuningTaskService.stopFineTuningTask(fineTuningTaskDO);

    }


    @Override
    public PageResult<FineTuningTaskDO> getFineTuningTaskPage (FineTuningTaskPageReqVO pageReqVO) {
        return fineTuningTaskMapper.selectPage(pageReqVO);
    }

}
