package com.lanyuanxiaoyao.service.web.controller;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.lanyuanxiaoyao.service.configuration.ExecutorProvider;
import com.lanyuanxiaoyao.service.configuration.entity.flink.FlinkVertex;
import com.lanyuanxiaoyao.service.configuration.entity.flink.FlinkVertexOverview;
import com.lanyuanxiaoyao.service.configuration.entity.yarn.YarnApplication;
import com.lanyuanxiaoyao.service.configuration.entity.yarn.YarnQueue;
import com.lanyuanxiaoyao.service.configuration.entity.yarn.YarnRootQueue;
import com.lanyuanxiaoyao.service.configuration.utils.ComparatorUtil;
import com.lanyuanxiaoyao.service.forest.service.FlinkService;
import com.lanyuanxiaoyao.service.forest.service.YarnService;
import com.lanyuanxiaoyao.service.web.controller.base.AmisCrudResponse;
import com.lanyuanxiaoyao.service.web.controller.base.AmisMapResponse;
import com.lanyuanxiaoyao.service.web.controller.base.AmisResponse;
import com.lanyuanxiaoyao.service.web.controller.base.BaseController;
import com.lanyuanxiaoyao.service.web.entity.YarnApplicationVO;
import com.lanyuanxiaoyao.service.web.entity.YarnClusterVO;
import java.util.Comparator;
import java.util.List;
import java.util.Optional;
import java.util.function.Function;
import org.eclipse.collections.api.factory.Lists;
import org.eclipse.collections.api.factory.Maps;
import org.eclipse.collections.api.list.ImmutableList;
import org.eclipse.collections.api.map.ImmutableMap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

/**
 * Sync Yarn
 *
 * @author lanyuanxiaoyao
 * @date 2023-04-21
 */
@RestController
@RequestMapping("yarn")
public class YarnController extends BaseController {
    private static final Logger logger = LoggerFactory.getLogger(YarnController.class);
    private static final ImmutableMap<String, Function<YarnApplication, Long>> SORT_MAP = Maps.immutable.of(
            "startedTime",
            YarnApplication::getStartedTime,
            "finishedTime",
            YarnApplication::getFinishedTime
    );
    private final YarnService yarnService;
    private final FlinkService flinkService;

    @SuppressWarnings("SpringJavaInjectionPointsAutowiringInspection")
    public YarnController(YarnService yarnService, FlinkService flinkService) {
        this.yarnService = yarnService;
        this.flinkService = flinkService;
    }

    @GetMapping("job_list")
    public AmisMapResponse jobList(
            @RequestParam("clusters") List<String> clusters,
            @RequestParam(value = "page", defaultValue = "1") Integer page,
            @RequestParam(value = "count", defaultValue = "10") Integer count,
            @RequestParam(value = "order", defaultValue = "startedTime") String order,
            @RequestParam(value = "direction", defaultValue = "desc") String direction,
            @RequestParam(value = "filter_state", defaultValue = "") String filterState,
            @RequestParam(value = "filter_final_status", defaultValue = "") String filterFinalStatus,
            @RequestParam(value = "search_id", defaultValue = "") String searchId,
            @RequestParam(value = "search_name", defaultValue = "") String searchName,
            @RequestParam(value = "precise", defaultValue = "false") Boolean precise,
            @RequestParam(value = "completion", defaultValue = "false") Boolean completion
    ) {
        boolean isFilterState = StrUtil.isNotBlank(filterState);
        boolean isFilterFinalStatus = StrUtil.isNotBlank(filterFinalStatus);
        boolean isSearchId = StrUtil.isNotBlank(searchId);
        boolean isSearchName = StrUtil.isNotBlank(searchName);
        Comparator<YarnApplication> comparator = ComparatorUtil.longComparator(order, direction, SORT_MAP);
        ImmutableList<YarnApplication> applications = Lists.immutable.ofAll(clusters)
                .asParallel(ExecutorProvider.EXECUTORS, 1)
                .flatCollect(yarnService::jobList)
                .select(app -> !isFilterState || ObjectUtil.contains(filterState, app.getState()))
                .select(app -> !isFilterFinalStatus || ObjectUtil.contains(filterFinalStatus, app.getFinalStatus()))
                .select(app -> !isSearchId || (precise ? StrUtil.equals(app.getId(), searchId) : StrUtil.contains(app.getId(), searchId)))
                .select(app -> !isSearchName || (precise ? StrUtil.equals(app.getName(), searchName) : StrUtil.contains(app.getName(), searchName)))
                .toSortedList(comparator)
                .toImmutable();
        int running = applications.count(app -> ObjectUtil.equal(app.getState(), "RUNNING"));
        int unRunning = applications.count(app -> ObjectUtil.notEqual(app.getState(), "RUNNING"));
        ImmutableList<YarnApplicationVO> result = applications
                .drop(Math.max(page - 1, 0) * count)
                .take(count)
                .asParallel(ExecutorProvider.EXECUTORS, 1)
                .collect(yarnApplication -> {
                    YarnApplicationVO vo = new YarnApplicationVO(yarnApplication);
                    if (!completion) {
                        return vo;
                    }
                    try {
                        if (vo.getCompactionApplication()
                                && StrUtil.equals(yarnApplication.getState(), "RUNNING")
                                && StrUtil.isNotBlank(yarnApplication.getTrackingUrl())) {
                            FlinkVertexOverview vertexOverview = flinkService.vertexOverview(yarnApplication.getTrackingUrl());
                            if (ObjectUtil.isNotNull(vertexOverview) && ObjectUtil.isNotEmpty(vertexOverview.getJobs())) {
                                Optional<FlinkVertex> vertex = vertexOverview.getJobs().getFirstOptional();
                                if (vertex.isPresent() && ObjectUtil.isNotNull(vertex.get().getTasks())) {
                                    FlinkVertex.Tasks tasks = vertex.get().getTasks();
                                    vo.setCompactionCompletionRatio(
                                            ObjectUtil.isNotNull(tasks.getTotal()) && tasks.getTotal() != 0
                                                    ? tasks.getRunning() * 1.0 / tasks.getTotal()
                                                    : 0
                                    );
                                }
                            }
                        }
                    } catch (Throwable e) {
                        logger.error("Get ratio failure", e);
                    }
                    return vo;
                })
                .toList()
                .toImmutable();
        return AmisResponse.responseCrudData(result, applications.size())
                .setData("running", running)
                .setData("unRunning", unRunning);
    }

    @GetMapping("job_current")
    public AmisMapResponse jobCurrent(@RequestParam("clusters") List<String> clusters, @RequestParam("name") String name) {
        Optional<YarnApplication> currentApp = Lists.immutable.ofAll(clusters)
                .asParallel(ExecutorProvider.EXECUTORS, 1)
                .flatCollect(cluster -> yarnService.jobListEquals(cluster, name))
                .select(app -> ObjectUtil.equals(app.getState(), "RUNNING"))
                .toSortedList(ComparatorUtil.longComparator("startedTime", ComparatorUtil.DESC, SORT_MAP))
                .getFirstOptional();
        if (currentApp.isPresent()) {
            return AmisResponse.responseMapData()
                    .setData("hasCurrent", true)
                    .setData("current", currentApp.get());
        } else {
            return AmisResponse.responseMapData("hasCurrent", false);
        }
    }

    @GetMapping("queue_list")
    public AmisCrudResponse queueList(@RequestParam("clusters") List<String> clusters, @RequestParam(value = "names", defaultValue = "") String names) {
        boolean isFilterNames = StrUtil.isNotBlank(names);
        ImmutableList<String> filterNames = Lists.immutable.of(names.split(","));
        ImmutableList<YarnClusterVO> results = Lists.immutable.ofAll(clusters)
                .asParallel(ExecutorProvider.EXECUTORS, 1)
                .collect(cluster -> {
                    YarnRootQueue root = yarnService.cluster(cluster);
                    ImmutableList<YarnQueue> queues = yarnService.queueList(cluster).select(queue -> !isFilterNames || filterNames.anySatisfy(n -> StrUtil.equals(queue.getQueueName(), n)));
                    return new YarnClusterVO(root, queues);
                })
                .toList()
                .toImmutable();
        return AmisResponse.responseCrudData(results, results.size());
    }

    @GetMapping("queue_names")
    public AmisMapResponse queueNames(@RequestParam("clusters") List<String> clusters) {
        ImmutableList<String> names = Lists.immutable.ofAll(clusters)
                .asParallel(ExecutorProvider.EXECUTORS, 1)
                .flatCollect(yarnService::queueList)
                .collect(YarnQueue::getQueueName)
                .toList()
                .toImmutable();
        return AmisResponse.responseMapData("queueNames", names);
    }

    @GetMapping("clusters")
    public AmisMapResponse clusters(@RequestParam("clusters") List<String> clusters) {
        ImmutableList<YarnRootQueue> roots = Lists.immutable.ofAll(clusters)
                .asParallel(ExecutorProvider.EXECUTORS, 1)
                .collect(yarnService::cluster)
                .toList()
                .toImmutable();
        return AmisResponse.responseMapData("cluster", clusters);
    }
}
