package com.ruoyi.platform.scheduling;

import com.ruoyi.platform.domain.Experiment;
import com.ruoyi.platform.domain.ExperimentIns;
import com.ruoyi.platform.domain.ModelDependency;
import com.ruoyi.platform.mapper.ExperimentDao;
import com.ruoyi.platform.mapper.ExperimentInsDao;
import com.ruoyi.platform.mapper.ModelDependencyDao;
import com.ruoyi.platform.service.ExperimentInsService;
import com.ruoyi.platform.utils.JacksonUtil;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.*;

@Component()
public class ExperimentInstanceStatusTask {
    @Autowired
    private ExperimentInsService experimentInsService;
    @Resource
    private ExperimentDao experimentDao;
    @Resource
    private ExperimentInsDao experimentInsDao;
    @Resource
    private ModelDependencyDao modelDependencyDao;
    private List<Integer> experimentIds = new ArrayList<>();

    @Scheduled(cron = "0/30 * * * * ?") // 每30S执行一次
    public void executeExperimentInsStatus() throws IOException {
        // 首先查到所有非终止态的实验实例
        List<ExperimentIns> experimentInsList = experimentInsService.queryByExperimentIsNotTerminated();
        // 去argo查询状态
        List<ExperimentIns> updateList = new ArrayList<>();
        if (experimentInsList != null && experimentInsList.size() > 0) {
            for (ExperimentIns experimentIns : experimentInsList) {
                //当原本状态为null或非终止态时才调用argo接口
                String oldStatus = experimentIns.getStatus();
                try {
                    experimentIns = experimentInsService.queryStatusFromArgo(experimentIns);
                } catch (Exception e) {
                    experimentIns.setStatus("Failed");
                }
                experimentIns.setUpdateTime(new Date());
                // 线程安全的添加操作
                synchronized (experimentIds) {
                    experimentIds.add(experimentIns.getExperimentId());
                }
                updateList.add(experimentIns);
            }
        }
        if (updateList.size() > 0) {
            experimentInsDao.insertOrUpdateBatch(updateList);

            //遍历模型关系表，找到
            List<ModelDependency> modelDependencyList = new ArrayList<ModelDependency>();
            for (ExperimentIns experimentIns : updateList) {
                ModelDependency modelDependencyquery = new ModelDependency();
                modelDependencyquery.setExpInsId(experimentIns.getId());
                modelDependencyquery.setState(2);

                List<ModelDependency> modelDependencyListquery = modelDependencyDao.queryByModelDependency(modelDependencyquery);
                if (modelDependencyListquery == null || modelDependencyListquery.size() == 0) {
                    continue;
                }
                ModelDependency modelDependency = modelDependencyListquery.get(0);
                //查看状态，
                if (StringUtils.equals("Failed", experimentIns.getStatus())) {
                    //取出节点状态
                    String trainTask = modelDependency.getTrainTask();
                    Map<String, Object> trainMap = JacksonUtil.parseJSONStr2Map(trainTask);
                    String task_id = (String) trainMap.get("task_id");
                    if (StringUtils.isEmpty(task_id)) {
                        continue;
                    }
                    String nodesStatus = experimentIns.getNodesStatus();
                    Map<String, Object> nodeMaps = JacksonUtil.parseJSONStr2Map(nodesStatus);
                    Map<String, Object> nodeMap = JacksonUtil.parseJSONStr2Map(JacksonUtil.toJSONString(nodeMaps.get(task_id)));

                    if (nodeMap == null) {
                        continue;
                    }
                    if (!StringUtils.equals("Succeeded", (String) nodeMap.get("phase"))) {
                        modelDependency.setState(0);
                        modelDependencyList.add(modelDependency);
                    }
                }
            }
            if (modelDependencyList.size() > 0) {
                modelDependencyDao.insertOrUpdateBatch(modelDependencyList);
            }
        }
    }

    @Scheduled(cron = "0/30 * * * * ?") // / 每30S执行一次
    public void executeExperimentStatus() throws IOException {
        if (experimentIds.size() == 0) {
            return;
        }
        // 存储需要更新的实验对象列表
        List<Experiment> updateExperiments = new ArrayList<>();
        for (Integer experimentId : experimentIds) {
            // 获取当前实验的所有实例列表
            List<ExperimentIns> insList = experimentInsService.getByExperimentId(experimentId);
            List<String> statusList = new ArrayList<String>();
            // 更新实验状态列表
            for (int i = 0; i < insList.size(); i++) {
                statusList.add(insList.get(i).getStatus());
            }
            String subStatus = statusList.toString().substring(1, statusList.toString().length() - 1);
            Experiment experiment = experimentDao.queryById(experimentId);
            // 如果实验状态列表发生变化，则更新实验对象，并加入到需要更新的列表中
            if (!StringUtils.equals(subStatus, experiment.getStatusList())) {
                experiment.setStatusList(subStatus);
                updateExperiments.add(experiment);
            }
        }

        if (!updateExperiments.isEmpty()) {
            experimentDao.insertOrUpdateBatch(updateExperiments);
            // 使用Iterator进行安全的删除操作
            Iterator<Integer> iterator = experimentIds.iterator();
            while (iterator.hasNext()) {
                Integer experimentId = iterator.next();
                for (Experiment experiment : updateExperiments) {
                    if (experiment.getId().equals(experimentId)) {
                        iterator.remove();
                    }
                }
            }
        }
    }
}
