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

import cn.iocoder.yudao.module.llm.dal.dataobject.finetuningtask.FineTuningTaskDO;
import cn.iocoder.yudao.module.llm.dal.mysql.finetuningtask.FineTuningTaskMapper;
import cn.iocoder.yudao.module.llm.enums.FineTuningTaskStatusConstants;
import cn.iocoder.yudao.module.llm.enums.FinetuningTaskStatusEnum;
import cn.iocoder.yudao.module.llm.service.basemodel.vo.ModelListRes;
import cn.iocoder.yudao.module.llm.service.http.TrainHttpService;
import cn.iocoder.yudao.module.llm.service.http.vo.AigcFineTuningDetailRespVO;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.esotericsoftware.minlog.Log;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLConnection;
import java.time.Duration;
import java.time.LocalDateTime;
import java.util.*;

@Component
@Slf4j
public class FineTuningTaskSyncService {

    @Resource
    TrainHttpService trainHttpService;

    @Resource
    FineTuningTaskMapper fineTuningTaskMapper;

    @Scheduled(cron ="10 * * * * ?")
    public void updateFineTuningTaskStatus() {
        Log.info("FineTuningTaskSync 定时任务启动");
        List<FineTuningTaskDO> fineTuningTaskDOList = fineTuningTaskMapper.selectList();
        for (FineTuningTaskDO fineTuningTaskDO : fineTuningTaskDOList) {

            if(Objects.equals(fineTuningTaskDO.getStatus(), FinetuningTaskStatusEnum.TRAINING.getStatus())
                || Objects.equals(fineTuningTaskDO.getStatus(), FinetuningTaskStatusEnum.WAITING.getStatus())){
                AigcFineTuningDetailRespVO resp = trainHttpService.finetuningDetail(new HashMap<>(), fineTuningTaskDO.getJobId());
                if (resp == null){
                    continue;
                }
                FineTuningTaskDO updateObj = new FineTuningTaskDO();
                Integer status = FineTuningTaskStatusConstants.getStatus(resp.getTrainStatus());
                if(status != null){
                    updateObj.setId(fineTuningTaskDO.getId());
                }
                //如果微调任务由训练中变为已完成，则取模型列表中获取该任务生成的模型id
                if((fineTuningTaskDO.getStatus() == 1 && status == 2) || (fineTuningTaskDO.getStatus() == 3 && status == 2)){
                    String jobModelName = fineTuningTaskDO.getJobModelName();
                /*    String resStr = trainHttpService.modelsList("");
                    Log.info("获取aicg模型列表返回数据内容{}",resStr);
                    JSONObject jsonObject = JSONObject.parseObject(resStr);
                    JSONObject data = jsonObject.getJSONObject("data");
                    JSONArray list = data.getJSONArray("list");
                    String s = list.toJSONString();
                    List<ModelListRes> modelListRess = JSONObject.parseArray(s, ModelListRes.class);
                    if(modelListRess != null){
                        ModelListRes modelListRes = modelListRess.get(0);
                        updateObj.setJobModelListId(modelListRes.getId());

                    }else{
                        //todo 已训练完成的微调任务，没有在模型列表找对对应模型的处理方式
                    }*/

                    try {
                        // 获取模型调优时长
                        String query = "?filter={\"job_id\":\""+fineTuningTaskDO.getJobId()+"\"}";
                        String res = trainHttpService.modelTableQuery(new HashMap<>(), "fine_tuning_train_job",query);
                        log.info("获取 aigc fine_tuning_train_job 表数据 info {}",res);
                        JSONArray jsonArray = JSONArray.parseArray(res);
                        if(jsonArray.size() > 0){
                            JSONObject jsonObject1 = jsonArray.getJSONObject(0);
                            updateObj.setTrainDuration(jsonObject1.getString("train_duration"));
                        }
                        updateObj.setStatus(2);
                        // 获取模型id
                        String querModels = "?filter={\"model_name\":\""+resp.getFineTunedModel()+"\"}";
                        String resModels = trainHttpService.modelTableQuery(new HashMap<>(), "models",querModels);
                        log.info("获取 aigc models 表数据 info {}",resModels);
                        JSONArray jsonArrayModels = JSONArray.parseArray(resModels);

                        // 时长获取
                        Duration duration = Duration.between(fineTuningTaskDO.getCreateTime(),LocalDateTime.now());
                        long minutes = duration.toMinutes();
                        updateObj.setTrainDuration(String.valueOf(minutes));

                        if(jsonArrayModels.size() > 0){
                            JSONObject jsonObjectModels = jsonArrayModels.getJSONObject(0);
                            updateObj.setJobModelListId(jsonObjectModels.getLong("id"));
                        }

                    }catch (Exception e){
                        log.error(" error {}",e.getMessage());
                    }

                    try {
                        //获取检查点信息
                        //todo 模型工厂的功能有问题，暂时写死
                        jobModelName = "Qwen2.5-0.5B-Instruct-147";
                        String checkFileList = trainHttpService.getCheckFileList(jobModelName);
                        List<String> checkpoints = new ArrayList<>();
                        List<String> fileUrls = new ArrayList<>();
                        List<String> fileList = JSONArray.parseArray(checkFileList,String.class);
                        Map<String,JSONObject> map = new HashMap<>();
                        for (String s : fileList) {
                            if(s.contains("checkpoint")){
                                checkpoints.add(s);
                            }
                        }
                        if(checkpoints.size() > 0){
                            for (String checkpoint : checkpoints) {
                                String filePath = "/" + checkpoint + "/trainer_state.json";
                                fileUrls.add(filePath);
                                String fileUrl = trainHttpService.getCheckFile(jobModelName, filePath);
                                try {
                                    URL url = new URL(fileUrl);
                                    URLConnection urlConnection = url.openConnection();
                                    InputStream inputStream = urlConnection.getInputStream();
                                    BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream));
                                    String line;
                                    StringBuilder content = new StringBuilder();
                                    while ((line = reader.readLine()) != null) {
                                        content.append(line + "\n");
                                    }
                                    reader.close();
                                    map.put(checkpoint,JSONObject.parseObject(content.toString()));
                                } catch (MalformedURLException e) {
                                    e.printStackTrace();
                                } catch (IOException e) {
                                    e.printStackTrace();
                                }
                            }
                        }

                        updateObj.setCheckPointFilePath(JSONObject.toJSONString(fileUrls));
                        updateObj.setCheckPointData(JSONObject.toJSONString(map));

                    }catch (Exception e){
                        log.error(" error {}",e.getMessage());
                    }

                    fineTuningTaskMapper.updateById(updateObj);
                }
                else {
                    updateObj.setStatus(status);
                }
                fineTuningTaskMapper.updateById(updateObj);

            }

        }
    }


}
