package com.ruoyi.platform.service.impl;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.platform.domain.ExperimentIns;
import com.ruoyi.platform.mapper.ExperimentInsDao;
import com.ruoyi.platform.service.AimService;
import com.ruoyi.platform.service.ExperimentInsService;
import com.ruoyi.platform.utils.AIM64EncoderUtil;
import com.ruoyi.platform.utils.HttpUtils;
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.Value;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageImpl;
import org.springframework.data.domain.PageRequest;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.util.*;
import java.util.stream.Collectors;

@Service
public class AimServiceImpl implements AimService {
    @Resource
    private ExperimentInsService experimentInsService;
    @Resource
    private ExperimentInsDao experimentInsDao;

    @Value("${aim.url}")
    private String aimUrl;
    @Value("${aim.proxyUrl}")
    private String aimProxyUrl;

    @Override
    public List<InsMetricInfoVo> getExpTrainInfos(Integer experimentId, String offset, int limit) throws Exception {
        return getAimRunInfos(true, experimentId, offset, limit);
    }

    @Override
    public List<InsMetricInfoVo> getExpEvaluateInfos(Integer experimentId, String offset, int limit) throws Exception {
        return getAimRunInfos(false, experimentId, offset, limit);
    }

    @Override
    public String getExpMetrics(List<String> runIds) throws Exception {
        String decode = AIM64EncoderUtil.decode(runIds);
        return aimUrl + "/metrics?select=" + decode;
    }

    private List<InsMetricInfoVo> getAimRunInfos(boolean isTrain, Integer experimentId, String offset, int limit) throws Exception {
        String experimentName = "experiment-" + experimentId + "-train";
        if (!isTrain) {
            experimentName = "experiment-" + experimentId + "-evaluate";
        }

        StringBuffer query = new StringBuffer("run.experiment==\"" + experimentName + "\"");

        String encodedUrlString = URLEncoder.encode(String.valueOf(query), "UTF-8");
        encodedUrlString = encodedUrlString + "&limit=" + limit;
        if (StringUtils.isNotEmpty(offset)) {
            encodedUrlString = encodedUrlString + "&offset=" + offset;
        }

        String url = aimProxyUrl + "/api/runs/search/run?query=" + encodedUrlString;
        String s = HttpUtils.sendGet(url, null);
        List<Map<String, Object>> response = JacksonUtil.parseJSONStr2MapList(s);
        System.out.println("response: " + JacksonUtil.toJSONString(response));
        if (response == null || response.size() == 0) {
            return new ArrayList<>();
        }
        //查询实例数据
        List<ExperimentIns> byExperimentId = experimentInsService.queryByExperimentId(experimentId);

        if (byExperimentId == null || byExperimentId.size() == 0) {
            return new ArrayList<>();
        }
        List<InsMetricInfoVo> aimRunInfoList = new ArrayList<>();
        for (Map<String, Object> run : response) {
            InsMetricInfoVo aimRunInfo = new InsMetricInfoVo();
            String runHash = (String) run.get("run_hash");

            aimRunInfo.setRunId(runHash);

            Map params = (Map) run.get("params");
            Map<String, Object> paramMap = JsonUtils.flattenJson("", params);
            aimRunInfo.setParams(paramMap);
            String aimrunId = (String) paramMap.get("id");
            Map<String, Object> tracesMap = (Map<String, Object>) run.get("traces");
            List<Map<String, Object>> metricList = (List<Map<String, Object>>) tracesMap.get("metric");
            //过滤name为__system__开头的对象
            aimRunInfo.setMetrics(new HashMap<>());
            if (metricList != null && metricList.size() > 0) {
                List<Map<String, Object>> metricRelList = metricList.stream()
                        .filter(map -> !StringUtils.startsWith((String) map.get("name"), "__system__"))
                        .collect(Collectors.toList());
                if (metricRelList != null && metricRelList.size() > 0) {
                    Map<String, Object> relMetricMap = new HashMap<>();
                    for (Map<String, Object> metricMap : metricRelList) {
                        relMetricMap.put((String) metricMap.get("name"), metricMap.get("last_value"));
                    }
                    aimRunInfo.setMetrics(relMetricMap);
                }
            }
            //找到ins
            for (ExperimentIns ins : byExperimentId) {
                String metricRecordString = ins.getMetricRecord();
                if (StringUtils.isEmpty(metricRecordString)) {
                    continue;
                }
                if (metricRecordString.contains(aimrunId)) {
                    aimRunInfo.setExperimentInsId(ins.getId());
                    aimRunInfo.setStatus(ins.getStatus());
                    aimRunInfo.setStartTime(ins.getCreateTime());
                    Map<String, Object> metricRecordMap = JacksonUtil.parseJSONStr2Map(metricRecordString);
                    if (isTrain) {
                        List<Map<String, Object>> records = (List<Map<String, Object>>) metricRecordMap.get("train");
                        List<String> datasetList = getTrainDateSet(records, aimrunId);
                        aimRunInfo.setDataset(datasetList);
                    } else {
                        List<Map<String, Object>> records = (List<Map<String, Object>>) metricRecordMap.get("evaluate");
                        List<String> datasetList = getTrainDateSet(records, aimrunId);
                        aimRunInfo.setDataset(datasetList);
                    }
                    aimRunInfoList.add(aimRunInfo);
                }
            }

        }
        //判断哪个最长

        // 获取所有 metrics 的 key 的并集
        Set<String> metricsKeys = (Set<String>) aimRunInfoList.stream()
                .map(InsMetricInfoVo::getMetrics)
                .flatMap(metrics -> metrics.keySet().stream())
                .collect(Collectors.toSet());
        // 将并集赋值给每个 InsMetricInfoVo 的 metricsNames 属性
        aimRunInfoList.forEach(vo -> vo.setMetricsNames(new ArrayList<>(metricsKeys)));

        // 获取所有 params 的 key 的并集
        Set<String> paramKeys = (Set<String>) aimRunInfoList.stream()
                .map(InsMetricInfoVo::getParams)
                .flatMap(params -> params.keySet().stream())
                .collect(Collectors.toSet());
        // 将并集赋值给每个 InsMetricInfoVo 的 paramsNames 属性
        aimRunInfoList.forEach(vo -> vo.setParamsNames(new ArrayList<>(paramKeys)));

        return aimRunInfoList;
    }

    public Page<InsMetricInfoVo> getExpInfos(boolean isTrain, Integer experimentId, int page, int size) {
        PageRequest pageRequest = PageRequest.of(page, size);
        long count = experimentInsDao.countTorE(experimentId, isTrain);
        List<ExperimentIns> experimentInsList = experimentInsDao.queryTorE(experimentId, isTrain, pageRequest);

        List<InsMetricInfoVo> aimRunInfoList = new ArrayList<>();
        for (ExperimentIns experimentIns : experimentInsList) {
            InsMetricInfoVo aimRunInfo = new InsMetricInfoVo();
            aimRunInfo.setExperimentInsId(experimentIns.getId());
            aimRunInfo.setStartTime(experimentIns.getCreateTime());
            aimRunInfo.setStatus(experimentIns.getStatus());

            //解析参数
            JSONArray params = JSON.parseArray(experimentIns.getGlobalParam());
            HashMap<String, Object> paramsMap = new HashMap<>();
            List<String> paramsNames = new ArrayList<>();

            for (int i = 0; i < params.size(); i++) {
                JSONObject jsonObject = params.getJSONObject(i);
                String paramName = jsonObject.getString("param_name");
                String paramValue = jsonObject.getString("param_value");
                paramsMap.put(paramName, paramValue);
                paramsNames.add(paramName);
            }
            aimRunInfo.setParams(paramsMap);
            aimRunInfo.setParamsNames(paramsNames);

            //解析数据集
            Map<String, Object> metricRecord = JacksonUtil.parseJSONStr2Map(experimentIns.getMetricRecord());
            if (isTrain) {
                aimRunInfo.setDataset(getDataset("train", metricRecord));
            } else {
                aimRunInfo.setDataset(getDataset("evaluate", metricRecord));
            }

            //解析指标
            if (StringUtils.isNotEmpty(experimentIns.getMetricValue())) {
                Map<String, Object> metricValue = JacksonUtil.parseJSONStr2Map(experimentIns.getMetricValue());
                if (isTrain) {
                    setMetricValue("train", metricValue, aimRunInfo);
                } else {
                    setMetricValue("evaluate", metricValue, aimRunInfo);
                }
            }
            aimRunInfoList.add(aimRunInfo);
        }
        return new PageImpl<>(aimRunInfoList, pageRequest, count);
    }


    @Override
    public List<InsMetricInfoVo> getExpInfos1(boolean isTrain, Integer experimentId, String runId) throws Exception {
        String encodedUrlString = URLEncoder.encode("run.id==\"" + runId + "\"", "UTF-8");
        String url = aimProxyUrl + "/api/runs/search/run?query=" + encodedUrlString;
        String s = HttpUtils.sendGet(url, null);
        List<Map<String, Object>> response = JacksonUtil.parseJSONStr2MapList(s);
        System.out.println("response: " + JacksonUtil.toJSONString(response));
        if (response == null || response.size() == 0) {
            return new ArrayList<>();
        }
        //查询实例数据
        List<ExperimentIns> byExperimentId = experimentInsService.queryByExperimentId(experimentId);

        if (byExperimentId == null || byExperimentId.size() == 0) {
            return new ArrayList<>();
        }
        List<InsMetricInfoVo> aimRunInfoList = new ArrayList<>();
        for (Map<String, Object> run : response) {
            InsMetricInfoVo aimRunInfo = new InsMetricInfoVo();
            String runHash = (String) run.get("run_hash");

            aimRunInfo.setRunId(runHash);

            Map params = (Map) run.get("params");
            Map<String, Object> paramMap = JsonUtils.flattenJson("", params);
            aimRunInfo.setParams(paramMap);
            String aimrunId = (String) paramMap.get("id");
            Map<String, Object> tracesMap = (Map<String, Object>) run.get("traces");
            List<Map<String, Object>> metricList = (List<Map<String, Object>>) tracesMap.get("metric");
            //过滤name为__system__开头的对象
            aimRunInfo.setMetrics(new HashMap<>());
            if (metricList != null && metricList.size() > 0) {
                List<Map<String, Object>> metricRelList = metricList.stream()
                        .filter(map -> !StringUtils.startsWith((String) map.get("name"), "__system__"))
                        .collect(Collectors.toList());
                if (metricRelList != null && metricRelList.size() > 0) {
                    Map<String, Object> relMetricMap = new HashMap<>();
                    for (Map<String, Object> metricMap : metricRelList) {
                        relMetricMap.put((String) metricMap.get("name"), metricMap.get("last_value"));
                    }
                    aimRunInfo.setMetrics(relMetricMap);
                }
            }
            //找到ins
            for (ExperimentIns ins : byExperimentId) {
                String metricRecordString = ins.getMetricRecord();
                if (StringUtils.isEmpty(metricRecordString)) {
                    continue;
                }
                if (metricRecordString.contains(aimrunId)) {
                    aimRunInfo.setExperimentInsId(ins.getId());
                    aimRunInfo.setStatus(ins.getStatus());
                    aimRunInfo.setStartTime(ins.getCreateTime());
                    Map<String, Object> metricRecordMap = JacksonUtil.parseJSONStr2Map(metricRecordString);
                    if (isTrain) {
                        List<Map<String, Object>> records = (List<Map<String, Object>>) metricRecordMap.get("train");
                        List<String> datasetList = getTrainDateSet(records, aimrunId);
                        aimRunInfo.setDataset(datasetList);
                    } else {
                        List<Map<String, Object>> records = (List<Map<String, Object>>) metricRecordMap.get("evaluate");
                        List<String> datasetList = getTrainDateSet(records, aimrunId);
                        aimRunInfo.setDataset(datasetList);
                    }
                    aimRunInfoList.add(aimRunInfo);
                }
            }
        }

        //判断哪个最长

        // 获取所有 metrics 的 key 的并集
        Set<String> metricsKeys = (Set<String>) aimRunInfoList.stream()
                .map(InsMetricInfoVo::getMetrics)
                .flatMap(metrics -> metrics.keySet().stream())
                .collect(Collectors.toSet());
        // 将并集赋值给每个 InsMetricInfoVo 的 metricsNames 属性
        aimRunInfoList.forEach(vo -> vo.setMetricsNames(new ArrayList<>(metricsKeys)));

        // 获取所有 params 的 key 的并集
        Set<String> paramKeys = (Set<String>) aimRunInfoList.stream()
                .map(InsMetricInfoVo::getParams)
                .flatMap(params -> params.keySet().stream())
                .collect(Collectors.toSet());
        // 将并集赋值给每个 InsMetricInfoVo 的 paramsNames 属性
        aimRunInfoList.forEach(vo -> vo.setParamsNames(new ArrayList<>(paramKeys)));

        return aimRunInfoList;
    }

    private List<String> getDataset(String isTrain, Map<String, Object> metricRecord) {
        List<String> datasetList = new ArrayList<>();
        List<Map<String, Object>> trainMetricRecords = (List<Map<String, Object>>) metricRecord.get(isTrain);
        if (trainMetricRecords != null) {
            for (Map<String, Object> trainMetricRecord : trainMetricRecords) {
                String taskId = (String) trainMetricRecord.get("task_id");
                if (taskId.startsWith("model-" + isTrain)) {
                    List<Map<String, Object>> datasets = (List<Map<String, Object>>) trainMetricRecord.get("datasets");
                    for (Map<String, Object> dataset : datasets) {
                        String datasetName = dataset.get("dataset_name") + ":" + dataset.get("dataset_version");
                        datasetList.add(datasetName);
                    }
                }
            }
        }
        return datasetList;
    }

    private void setMetricValue(String isTrain, Map<String, Object> metricValue, InsMetricInfoVo aimRunInfo) {
        Map<String, Object> metricValues = (Map<String, Object>) metricValue.get(isTrain);

        HashMap<String, Object> metrics = new HashMap<>();
        List<String> metricsNames = new ArrayList<>();

        if (metricValues != null) {
            for (String key : metricValues.keySet()) {
                Map<String, Object> valueMap = (Map<String, Object>) metricValues.get(key);
                aimRunInfo.setRunId((String) valueMap.get("run_hash"));

                metrics.putAll(valueMap.entrySet().stream().filter(entry -> !entry.getKey().equals("run_hash")).collect(Collectors.toMap(
                        Map.Entry::getKey,
                        Map.Entry::getValue
                )));
                metricsNames.addAll(valueMap.keySet().stream().filter(entry -> !entry.equals("run_hash")).collect(Collectors.toList()));
            }
        }
        aimRunInfo.setMetrics(metrics);
        aimRunInfo.setMetricsNames(metricsNames);
    }

    private List<String> getTrainDateSet(List<Map<String, Object>> records, String aimrunId) {
        List<String> datasetList = new ArrayList<>();
        for (Map<String, Object> record : records) {
            if (StringUtils.equals(aimrunId, (String) record.get("run_id"))) {
                List<Map<String, Object>> datasets = (List<Map<String, Object>>) record.get("datasets");
                if (datasets == null || datasets.size() == 0) {
                    continue;
                }
                for (Map<String, Object> dataset : datasets) {
                    String datasetName = (String) dataset.get("dataset_name") + ":" + (String) dataset.get("dataset_version");
                    datasetList.add(datasetName);
                }
                break;
            }
        }
        return datasetList;
    }

    @Override
    public HashMap<String, Object> queryMetricsParams(String runId) throws UnsupportedEncodingException {
        String encodedUrlString = URLEncoder.encode("run.id==\"" + runId + "\"", "UTF-8");
        String url = aimProxyUrl + "/api/runs/search/run?query=" + encodedUrlString;
        String s = HttpUtils.sendGet(url, null);
        List<Map<String, Object>> response = JacksonUtil.parseJSONStr2MapList(s);
        if (response == null || response.size() == 0) {
            return new HashMap<>();
        }

        HashMap<String, Object> resultMap = new HashMap<>();
        List<Map<String, Object>> paramList = new ArrayList<>();

        Map<String, Object> map = response.get(0);
        LinkedHashMap<String, ArrayList> traces = (LinkedHashMap<String, ArrayList>) map.get("traces");
        if (traces != null) {
            List<Map<String, Object>> metrics = traces.get("metric");
            for (Map<String, Object> metric : metrics) {
                Map<String, Object> metricParam = new HashMap<>();
                metricParam.put("context", metric.get("context"));
                metricParam.put("name", metric.get("name"));
                paramList.add(metricParam);
            }
            resultMap.put("params", JSON.toJSONString(paramList));
        }
        resultMap.put("run_hash", map.get("run_hash"));
        return resultMap;
    }

    @Override
    public List<Map<String, Object>> getBatchMetric(String runHash, String body) {
        String url = aimUrl + "/api/runs/" + runHash + "/metric/get-batch";
        String response = HttpUtils.sendPost(url, null, body);
        if (StringUtils.isNotEmpty(response)) {
            return JacksonUtil.parseJSONStr2MapList(response);
        }
        return null;
    }

}
