package com.lanyuanxiaoyao.service.command.commands;

import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.lanyuanxiaoyao.service.command.entity.TableMetaJob;
import com.lanyuanxiaoyao.service.command.utils.CommandLineUtils;
import com.lanyuanxiaoyao.service.command.utils.ShellTableUtils;
import com.lanyuanxiaoyao.service.common.Constants;
import com.lanyuanxiaoyao.service.common.entity.RunMeta;
import com.lanyuanxiaoyao.service.common.utils.NameHelper;
import com.lanyuanxiaoyao.service.configuration.ExecutorProvider;
import com.lanyuanxiaoyao.service.configuration.entity.yarn.YarnApplication;
import com.lanyuanxiaoyao.service.configuration.entity.yarn.YarnClusters;
import com.lanyuanxiaoyao.service.configuration.entity.zookeeper.ZookeeperNode;
import com.lanyuanxiaoyao.service.forest.service.InfoService;
import com.lanyuanxiaoyao.service.forest.service.YarnService;
import com.lanyuanxiaoyao.service.forest.service.ZookeeperService;
import com.lanyuanxiaoyao.service.forest.service.launcher.LauncherService;
import com.lanyuanxiaoyao.service.forest.service.launcher.LaunchersService;
import java.time.Instant;
import java.time.LocalDateTime;
import java.time.ZoneId;
import me.tongfei.progressbar.ProgressBar;
import org.eclipse.collections.api.factory.Lists;
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.shell.standard.ShellComponent;
import org.springframework.shell.standard.ShellMethod;
import org.springframework.shell.standard.ShellOption;
import org.springframework.shell.table.TableModelBuilder;

/**
 * 表同步任务 description
 *
 * @author ZhangJiacheng
 * @date 2022-04-22
 */
@SuppressWarnings("SpringJavaInjectionPointsAutowiringInspection")
@ShellComponent("同步任务相关操作")
public class YarnCommand extends AbstractUtilShellComponent {
    private static final Logger logger = LoggerFactory.getLogger(YarnCommand.class);
    private final static String KILL_CONFIRMATION_MESSAGE = "操作将会停止上表所显示的表同步任务, 请再次确认操作";
    private final static String RUN_CONFIRMATION_MESSAGE = "操作将会启动上表所显示的表同步任务, 请再次确认操作";
    private final InfoService infoService;
    private final LauncherService launcherService;
    private final ZookeeperService zookeeperService;
    private final YarnService yarnService;
    private final YarnClusters yarnClusters;

    public YarnCommand(InfoService infoService, LaunchersService launchersService, ZookeeperService zookeeperService, YarnService yarnService, YarnClusters yarnClusters) {
        this.infoService = infoService;
        this.zookeeperService = zookeeperService;
        this.yarnService = yarnService;
        this.yarnClusters = yarnClusters;
        logger.info("Default sync cluster: {}", yarnClusters.getDefaultSyncCluster());
        this.launcherService = launchersService.getService(yarnClusters.getDefaultSyncCluster());
    }

    private String generateFlinkJobsTable(ImmutableList<TableMetaJob> jobs, Boolean showFinishTime) {
        TableModelBuilder<Object> builder = new TableModelBuilder<>();
        builder.addRow();
        builder.addValue("#");
        builder.addValue("Id");
        builder.addValue("Name");
        builder.addValue("State");
        builder.addValue("Application Id");
        builder.addValue("Launch time");
        if (showFinishTime) {
            builder.addValue("Finish time");
        }
        builder.addValue("Tracking url");
        for (int index = 0; index < jobs.size(); index++) {
            TableMetaJob job = jobs.get(index);
            builder.addRow();
            builder.addValue(index + 1);
            builder.addValue(job.getFlinkJobId());
            builder.addValue(job.getAlias());
            if (job.isRunning()) {
                builder.addValue(job.getState());
                builder.addValue(job.getId());
                LocalDateTime launchTime = LocalDateTime.ofInstant(Instant.ofEpochMilli(job.getStartedTime()), ZoneId.systemDefault());
                builder.addValue(launchTime.format(Constants.FORMATTER));
                if (showFinishTime) {
                    if (ObjectUtil.isNotNull(job.getFinishedTime()) || job.getFinishedTime() != 0) {
                        LocalDateTime finishTime = LocalDateTime.ofInstant(Instant.ofEpochMilli(job.getFinishedTime()), ZoneId.systemDefault());
                        builder.addValue(finishTime.format(Constants.FORMATTER));
                    } else {
                        builder.addValue("");
                    }
                }
                builder.addValue(job.getTrackingUrl() + " ");
            } else {
                builder.addValue("");
                builder.addValue("");
                if (showFinishTime) {
                    builder.addValue("");
                }
                builder.addValue("");
                builder.addValue("");
            }
        }
        return ShellTableUtils.renderTable(builder);
    }

    @ShellMethod("启动表同步任务")
    public String yarnRun(
            @ShellOption(help = "Flink job id") Long flinkJobId,
            @ShellOption(
                    help = "Ignore double check",
                    defaultValue = "false"
            ) Boolean ignoreCheck
    ) {
        return yarnRunBatch(new Long[]{flinkJobId}, ignoreCheck);
    }

    @ShellMethod("停止所有正在运行的表同步任务")
    public String yarnKillRunning(
            @ShellOption(
                    help = "Ignore double check",
                    defaultValue = "false"
            ) Boolean ignoreCheck
    ) {
        return yarnKillBatch(getAllJobId(true).toArray(new Long[]{}), ignoreCheck);
    }

    @ShellMethod("启动所有未运行的表同步任务")
    public String yarnRunUnRunning(
            @ShellOption(
                    help = "Ignore double check",
                    defaultValue = "false"
            ) Boolean ignoreCheck
    ) {
        return yarnRunBatch(getAllJobId(false).toArray(new Long[]{}), ignoreCheck);
    }

    private ImmutableList<Long> getAllJobId(Boolean isRunning) {
        return infoService.allFlinkJobId()
                .select(id -> {
                    String lockPath = NameHelper.syncRunningLockPath(id);
                    if (isRunning) {
                        return zookeeperService.existsPath(lockPath);
                    } else {
                        return !zookeeperService.existsPath(lockPath);
                    }
                });
    }

    private ImmutableList<TableMetaJob> getRunningJobs(Long[] flinkJobIds) {
        ImmutableList<String> runningPaths = zookeeperService.getChildren(NameHelper.ZK_SYNC_RUNNING_LOCK_PATH).collect(ZookeeperNode::getPath);
        return Lists.immutable.of(flinkJobIds)
                .distinct()
                .collect(infoService::tableMetaList)
                .flatCollect(meta -> meta)
                .asParallel(ExecutorProvider.EXECUTORS, 5)
                .collect(meta -> {
                    if (runningPaths.contains(NameHelper.syncRunningLockPath(meta.getJob().getId()))) {
                        RunMeta runMeta = zookeeperService.getSyncRunMeta(meta.getJob().getId());
                        YarnApplication application = yarnService.jobDetail(yarnClusters.getDefaultSyncCluster(), runMeta.getApplicationId());
                        return new TableMetaJob(meta, application);
                    } else {
                        return new TableMetaJob(meta);
                    }
                })
                .toList()
                .toImmutable();
    }

    @ShellMethod("批量停止表同步任务")
    public String yarnKillBatch(
            @ShellOption(help = "Flink job id") Long[] flinkJobIds,
            @ShellOption(
                    help = "Ignore double check",
                    defaultValue = "false"
            ) Boolean ignoreCheck
    ) {
        ImmutableList<TableMetaJob> targetJobs = getRunningJobs(flinkJobIds);
        return CommandLineUtils.generateResultLines(
                () -> {
                    if (ObjectUtil.isEmpty(targetJobs)) {
                        return "没有找到包含指定 flink job id 对应的任务";
                    } else {
                        System.out.println(generateFlinkJobsTable(targetJobs, false));
                        if (doubleCheck(KILL_CONFIRMATION_MESSAGE, ignoreCheck)) {
                            ImmutableList<Long> ids = targetJobs.select(TableMetaJob::isRunning)
                                    .collect(TableMetaJob::getFlinkJobId)
                                    .distinct();
                            try (ProgressBar pb = CommandLineUtils.progressbarBuilder("Kill jobs", ids.size()).build()) {
                                ids.forEach(id -> {
                                    launcherService.syncStop(id);
                                    pb.setExtraMessage(id.toString());
                                    pb.step();
                                });
                            }
                            return Constants.OPERATION_DONE;
                        } else {
                            return Constants.OPERATION_CANCEL;
                        }
                    }
                }
        );
    }


    @ShellMethod("停止表同步任务")
    public String yarnKill(
            @ShellOption(help = "Flink job id") Long flinkJobId,
            @ShellOption(
                    help = "Ignore double check",
                    defaultValue = "false"
            ) Boolean ignoreCheck
    ) {
        return yarnKillBatch(new Long[]{flinkJobId}, ignoreCheck);
    }

    @ShellMethod("批量启动表同步任务")
    public String yarnRunBatch(
            @ShellOption(help = "Flink job id") Long[] flinkJobIds,
            @ShellOption(
                    help = "Ignore double check",
                    defaultValue = "false"
            ) Boolean ignoreCheck
    ) {
        ImmutableList<TableMetaJob> targetJobs = getRunningJobs(flinkJobIds);
        return CommandLineUtils.generateResultLines(
                () -> {
                    if (ObjectUtil.isEmpty(targetJobs)) {
                        return "没有找到包含指定 flink job id 对应的任务";
                    } else {
                        System.out.println(generateFlinkJobsTable(targetJobs, false));
                        if (doubleCheck(RUN_CONFIRMATION_MESSAGE, ignoreCheck)) {
                            ImmutableList<Long> ids = targetJobs.select(TableMetaJob::isNotRunning)
                                    .collect(TableMetaJob::getFlinkJobId)
                                    .distinct();
                            try (ProgressBar pb = CommandLineUtils.progressbarBuilder("Run jobs", ids.size()).build()) {
                                ids.forEach(id -> {
                                    ThreadUtil.safeSleep(1000);
                                    launcherService.syncStart(id);
                                    pb.setExtraMessage(id.toString());
                                    pb.step();
                                });
                            }
                            return Constants.OPERATION_DONE;
                        } else {
                            return Constants.OPERATION_CANCEL;
                        }
                    }
                }
        );
    }

    @ShellMethod("查询表同步任务状态")
    public String yarnAll(
            @ShellOption(help = "搜索值 (包含查询)",
                    defaultValue = "") String pattern,
            @ShellOption(help = "展示结果列表的 Flink job id",
                    defaultValue = "false") Boolean showFlinkJobId,
            @ShellOption(help = "展示运行的任务",
                    defaultValue = "false") Boolean running,
            @ShellOption(help = "展示非运行的任务",
                    defaultValue = "false") Boolean unRunning
    ) {
        if (running && unRunning) {
            throw new RuntimeException("--running 和 --un-running 参数不能同时使用");
        }
        ImmutableMap<String, YarnApplication> applicationMap = yarnService.jobListRunning(yarnClusters.getDefaultSyncCluster())
                .toMap(YarnApplication::getName, app -> app)
                .toImmutable();
        ImmutableList<String> locks = zookeeperService.getChildren(NameHelper.ZK_SYNC_RUNNING_LOCK_PATH).collect(ZookeeperNode::getPath);
        ImmutableList<TableMetaJob> jobs = infoService.tableMetaList()
                .select(meta -> StrUtil.contains(meta.getAlias(), pattern))
                .select(meta -> {
                    String lockPath = NameHelper.syncRunningLockPath(meta.getJob().getId(), meta.getAlias());
                    if (running && locks.contains(lockPath)) {
                        return true;
                    } else if (unRunning && !locks.contains(lockPath)) {
                        return true;
                    } else {
                        return !running && !unRunning;
                    }
                })
                .collect(meta -> {
                    String jobName = NameHelper.syncJobName(meta.getJob().getId(), meta.getJob().getName());
                    return new TableMetaJob(meta, applicationMap.get(jobName));
                });
        return CommandLineUtils.generateResultLines(
                () -> generateFlinkJobsTable(jobs, false),
                () -> showFlinkJobId ? CommandLineUtils.generateResultLines(
                        () -> "Flink ids",
                        () -> jobs
                                .collect(TableMetaJob::getFlinkJobId)
                                .collect(String::valueOf)
                                .distinct()
                                .makeString(",")
                ) : ""
        );
    }

    @ShellMethod("test")
    public void test() {
    }
}
