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

import cn.iocoder.yudao.module.llm.controller.admin.dataset.vo.DatasetRespVO;
import cn.iocoder.yudao.module.llm.dal.dataobject.basemodel.BaseModelDO;
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.finetuningtask.FineTuningTaskMapper;
import cn.iocoder.yudao.module.llm.enums.FinetuningTaskStatusEnum;
import cn.iocoder.yudao.module.llm.service.dataset.DatasetService;
import cn.iocoder.yudao.module.llm.service.http.TrainHttpService;
import cn.iocoder.yudao.module.llm.service.http.vo.AigcFineTuningCreateReqVO;
import cn.iocoder.yudao.module.llm.service.http.vo.AigcFineTuningCreateRespVO;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Lazy;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.HashMap;

@Service
@Slf4j
public class AsyncFineTuningTaskService {

    @Resource
    private FineTuningTaskMapper fineTuningTaskMapper;

    @Resource
    private TrainHttpService trainHttpService;

    @Lazy
    @Resource
    private DatasetService datasetService;
    @Resource
    private BaseModelMapper baseModelMapper;

    //大模型平台创建调优任务
    @Async
    public void createTuning(FineTuningTaskDO fineTuningTask) {
        try {
            AigcFineTuningCreateReqVO req = getAigcFineTuningCreateReqVO(fineTuningTask);
            BaseModelDO baseModelDO = baseModelMapper.selectById(fineTuningTask.getBaseModelId());
            Long datasetId = fineTuningTask.getDataset();
            DatasetRespVO dataset = datasetService.getDataset(datasetId);
//            req.setFileId(dataset.getJobId());
            if (baseModelDO != null){
                req.setModel(baseModelDO.getAigcModelName());
            }
            req.setDataset(StringUtils.isNotBlank(dataset.getDatasetFileUrl())?dataset.getDatasetFileUrl():"storage/json/2025/01/_xVKpfDH8.json");
            req.setDataset(dataset.getFileUrl());
            req.setSuffix(String.valueOf(fineTuningTask.getId()));
            AigcFineTuningCreateRespVO resp = trainHttpService.finetuningCreate(new HashMap<>(), req);
            FineTuningTaskDO updateObj = new FineTuningTaskDO();
            updateObj.setId(fineTuningTask.getId());
            if (resp != null) {
                updateObj.setJobId(resp.getJobId());
                updateObj.setStatus(FinetuningTaskStatusEnum.WAITING.getStatus());
                updateObj.setJobModelName(resp.getFineTunedModel());
                updateObj.setTrainLog(resp.getTrainLog());
            } else {
                updateObj.setStatus(FinetuningTaskStatusEnum.CANCELLED.getStatus());
            }
            fineTuningTaskMapper.updateById(updateObj);
        } catch (Exception e){
            e.printStackTrace();
        };
    }

    private static AigcFineTuningCreateReqVO getAigcFineTuningCreateReqVO(FineTuningTaskDO fineTuningTask) {
        AigcFineTuningCreateReqVO req = new AigcFineTuningCreateReqVO();
        req.setModel(fineTuningTask.getBaseModel());
        req.setTrainEpoch( (fineTuningTask.getEpoch() == 0) ? 1 :fineTuningTask.getEpoch());
        req.setSuffix(fineTuningTask.getTaskName());
        req.setRemark(fineTuningTask.getTaskIntro());
        req.setTrainBatchSize(fineTuningTask.getBatchSize());
        req.setEvalBatchSize(fineTuningTask.getBatchSize());
        req.setAccumulationSteps(fineTuningTask.getGradientAccumulation());
        req.setProcPerNode(fineTuningTask.getGpuCount());
        req.setLearningRate(fineTuningTask.getLearningRate());
        req.setModelMaxLength(fineTuningTask.getCutoffLen());
        req.setLora(true);
        return req;
    }

    //调优任务部署
    @Async
    public void startFineTuningTask(FineTuningTaskDO fineTuningTask) {
        try {
            createTuning(fineTuningTask);
        }catch(Exception e){
            e.printStackTrace();

        };
    }

    //调优任务停止
    @Async
    public void stopFineTuningTask(FineTuningTaskDO fineTuningTask) {
        try {
            trainHttpService.finetuningCancle(new HashMap<>(),fineTuningTask.getJobId());
            trainHttpService.finetuningDelete(new HashMap<>(), fineTuningTask.getJobId());
            fineTuningTask.setStatus(0);
            fineTuningTaskMapper.updateById(fineTuningTask);
        }catch(Exception e){
            e.printStackTrace();

        };
    }
}
