package com.ruoyi.platform.scheduling;

import com.ruoyi.system.api.constant.Constant;
import com.ruoyi.platform.domain.Experiment;
import com.ruoyi.platform.domain.ExperimentIns;
import com.ruoyi.platform.mapper.ExperimentDao;
import com.ruoyi.platform.mapper.ExperimentInsDao;
import com.ruoyi.platform.service.AimService;
import com.ruoyi.platform.service.ExperimentInsService;
import com.ruoyi.platform.service.ResourceOccupyService;
import com.ruoyi.platform.utils.JacksonUtil;
import com.ruoyi.platform.utils.JsonUtils;
import com.ruoyi.platform.vo.InsMetricInfoVo;
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.time.Instant;
import java.util.*;

@Component()
public class ExperimentInstanceStatusTask {
    @Autowired
    private ExperimentInsService experimentInsService;
    @Resource
    private ExperimentDao experimentDao;
    @Resource
    private ExperimentInsDao experimentInsDao;
    @Resource
    private AimService aimService;
    @Resource
    private ResourceOccupyService resourceOccupyService;

    private HashSet<Integer> experimentIds = new HashSet<>();

    @Scheduled(cron = "0/10 * * * * ?") // 每10S执行一次
    public void executeExperimentInsStatus() throws Exception {
        // 首先查到所有非终止态的实验实例
        List<ExperimentIns> experimentInsList = experimentInsService.queryByExperimentIsNotTerminated();
        // 去argo查询状态
        List<ExperimentIns> updateList = new ArrayList<>();
        if (experimentInsList != null && experimentInsList.size() > 0) {
            for (ExperimentIns experimentIns : experimentInsList) {
                //当原本状态为null或非终止态时才调用argo接口
                try {
                    experimentIns = experimentInsService.queryStatusFromArgo(experimentIns);
                } catch (Exception e) {
                    experimentIns.setStatus(Constant.Failed);
                    // 结束扣除积分
                    resourceOccupyService.endDeduce(Constant.TaskType_Workflow, null, Long.valueOf(experimentIns.getId()), null, null);
                }

                // 扣除积分
                if (StringUtils.isNotEmpty(experimentIns.getNodesStatus())) {
                    Map<String, Object> nodesStatusMap = JsonUtils.jsonToMap(experimentIns.getNodesStatus());
                    for (String key : nodesStatusMap.keySet()) {
                        Map<String, Object> value = (Map<String, Object>) nodesStatusMap.get(key);
                        String startedAt = (String) value.get("startedAt");
                        Instant instant = Instant.parse(startedAt);
                        Date startTime = Date.from(instant);
                        String phase = (String) value.get("phase");
                        String finishedAt = (String) value.get("finishedAt");

                        if (StringUtils.isEmpty(finishedAt) && Constant.Running.equals(phase)) {
                            resourceOccupyService.deducing(Constant.TaskType_Workflow, null, Long.valueOf(experimentIns.getId()), key, startTime);
                        } else if (StringUtils.isNotEmpty(finishedAt)) {
                            resourceOccupyService.endDeduce(Constant.TaskType_Workflow, null, Long.valueOf(experimentIns.getId()), key, startTime);
                        }
                    }
                }
                //运行成功的实验实例记录指标数值
                Map<String, Object> metricRecord = JacksonUtil.parseJSONStr2Map(experimentIns.getMetricRecord());
                List<Map<String, Object>> trainMetricRecords = (List<Map<String, Object>>) metricRecord.get("train");
                List<Map<String, Object>> evaluateMetricRecords = (List<Map<String, Object>>) metricRecord.get("evaluate");

                HashMap<String, Object> metricValue = new HashMap<>();
                HashMap<String, Object> trainMetricValues = new HashMap<>();
                HashMap<String, Object> evaluateMetricValues = new HashMap<>();

                if (trainMetricRecords != null && !trainMetricRecords.isEmpty()) {
                    for (Map<String, Object> trainMetricRecord : trainMetricRecords) {
                        HashMap<String, Object> trainMetricValue = new HashMap<>();
                        String taskId = (String) trainMetricRecord.get("task_id");
                        if (taskId.startsWith("model-train")) {
                            String runId = (String) trainMetricRecord.get("run_id");
                            List<InsMetricInfoVo> expTrainInfos = aimService.getExpInfos1(true, experimentIns.getExperimentId(), runId);
                            for (InsMetricInfoVo expTrainInfo : expTrainInfos) {
                                Map metrics = expTrainInfo.getMetrics();
                                trainMetricValue.putAll(metrics);
                                trainMetricValue.put("run_hash", expTrainInfo.getRunId());
                            }
                            if (trainMetricValue.size() > 0) {
                                trainMetricValues.put(taskId, trainMetricValue);
                            }
                        }
                    }
                }

                if (evaluateMetricRecords != null && !evaluateMetricRecords.isEmpty()) {
                    for (Map<String, Object> evaluateMetricRecord : evaluateMetricRecords) {
                        HashMap<String, Object> evaluateMetricValue = new HashMap<>();
                        String taskId = (String) evaluateMetricRecord.get("task_id");
                        if (taskId.startsWith("model-evaluate")) {
                            String runId = (String) evaluateMetricRecord.get("run_id");
                            List<InsMetricInfoVo> expTrainInfos = aimService.getExpInfos1(false, experimentIns.getExperimentId(), runId);
                            for (InsMetricInfoVo expTrainInfo : expTrainInfos) {
                                Map metrics = expTrainInfo.getMetrics();
                                evaluateMetricValue.putAll(metrics);
                                evaluateMetricValue.put("run_hash", expTrainInfo.getRunId());
                            }
                            if (evaluateMetricValue.size() > 0) {
                                evaluateMetricValues.put(taskId, evaluateMetricValue);
                            }
                        }
                    }
                }
                if (trainMetricValues.size() > 0) {
                    metricValue.put("train", trainMetricValues);
                } else {
                    metricValue.put("train", null);
                }
                if (evaluateMetricValues.size() > 0) {
                    metricValue.put("evaluate", evaluateMetricValues);
                } else {
                    metricValue.put("evaluate", null);
                }
                experimentIns.setMetricValue(JsonUtils.mapToJson(metricValue));
                experimentIns.setUpdateTime(new Date());
                // 线程安全的添加操作
                synchronized (experimentIds) {
                    experimentIds.add(experimentIns.getExperimentId());
                }
                updateList.add(experimentIns);
            }
            if (updateList.size() > 0) {
                experimentInsDao.insertOrUpdateBatch(updateList);
            }
        }
    }

    @Scheduled(cron = "0/10 * * * * ?") // / 每30S执行一次
    public void executeExperimentStatus() throws IOException {
        if (experimentIds.isEmpty()) {
            return;
        }
        // 存储需要更新的实验对象列表
        List<Experiment> updateExperiments = new ArrayList<>();
        Iterator<Integer> iterator1 = experimentIds.iterator();

        while (iterator1.hasNext()) {
            Integer experimentId = iterator1.next();
            // 获取当前实验的所有实例列表
            List<ExperimentIns> insList = experimentInsService.getByExperimentId(experimentId);
            List<String> statusList = new ArrayList<>();
            // 更新实验状态列表
            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 (experiment == null) {
                iterator1.remove();
            } else {
                // 如果实验状态列表发生变化，则更新实验对象，并加入到需要更新的列表中
                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();
                    }
                }
            }
        }
    }
}
