package com.lanyuanxiaoyao.service.launcher;

import cn.hutool.core.lang.Tuple;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.ReUtil;
import cn.hutool.core.util.StrUtil;
import com.lanyuanxiaoyao.service.common.Constants;
import com.lanyuanxiaoyao.service.common.entity.FlinkJob;
import com.lanyuanxiaoyao.service.common.entity.TableMeta;
import com.lanyuanxiaoyao.service.common.exception.CheckpointRootPathNotFoundException;
import com.lanyuanxiaoyao.service.common.utils.NameHelper;
import com.lanyuanxiaoyao.service.executor.Runner;
import com.lanyuanxiaoyao.service.forest.service.InfoService;
import com.lanyuanxiaoyao.service.launcher.configuration.HadoopConfiguration;
import com.lanyuanxiaoyao.service.launcher.configuration.HudiConfiguration;
import com.lanyuanxiaoyao.service.launcher.utils.HadoopUtil;
import com.lanyuanxiaoyao.service.launcher.utils.JacksonUtil;
import java.io.IOException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.regex.Pattern;
import javax.annotation.Resource;
import org.apache.flink.client.cli.ClientOptions;
import org.apache.flink.configuration.*;
import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.core.JsonProcessingException;
import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.yarn.configuration.YarnConfigOptions;
import org.apache.flink.yarn.configuration.YarnDeploymentTarget;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hudi.common.util.collection.Pair;
import org.eclipse.collections.api.factory.Lists;
import org.eclipse.collections.api.list.ImmutableList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;

import static com.lanyuanxiaoyao.service.common.Constants.HALF_HOUR;
import static com.lanyuanxiaoyao.service.common.Constants.MINUTE;

/**
 * 执行器
 *
 * @author ZhangJiacheng
 * @date 2022-06-06
 */
@Service
public class ExecutorService {
    private static final Logger logger = LoggerFactory.getLogger(ExecutorService.class);
    private static final ObjectMapper MAPPER = JacksonUtil.getMapper();
    private static final Pattern EXECUTOR_JAR_NAME = Pattern.compile(".+sync-(\\d+)\\.jar");
    @Resource
    private InfoService infoService;
    @Resource
    private HadoopConfiguration hadoopConfiguration;
    @Resource
    private HudiConfiguration hudiConfiguration;

    public static String[] syncArgs(FlinkJob flinkJob, ImmutableList<TableMeta> tableMetaList, String cluster) throws JsonProcessingException {
        List<String> argsList = Lists.mutable.empty();
        argsList.add(Constants.FLINK_JOB_OPTION);
        argsList.add(MAPPER.writeValueAsString(flinkJob));
        argsList.add(Constants.TABLE_META_LIST_OPTION);
        argsList.add(MAPPER.writeValueAsString(tableMetaList));
        argsList.add(Constants.CLUSTER_OPTION);
        argsList.add(cluster);
        return argsList.toArray(new String[]{});
    }

    public static String[] compactionArgs(FlinkJob flinkJob, TableMeta tableMeta, String instants, String cluster) throws JsonProcessingException {
        List<String> argsList = Lists.mutable.empty();
        argsList.add(Constants.FLINK_JOB_OPTION);
        argsList.add(MAPPER.writeValueAsString(flinkJob));
        argsList.add(Constants.TABLE_META_OPTION);
        argsList.add(MAPPER.writeValueAsString(tableMeta));
        if (ObjectUtil.isNotEmpty(instants)) {
            argsList.add(Constants.INSTANTS_OPTION);
            argsList.add(instants);
        }
        argsList.add(Constants.CLUSTER_OPTION);
        argsList.add(cluster);
        return argsList.toArray(new String[]{});
    }

    private Memory jobMemory(FlinkJob flinkJob, ImmutableList<TableMeta> tableMetas) {
        int jobMemory = 1024;
        int jobMetaspaceMemory = 256;
        int taskMemory = 512;
        int taskMetaspaceMemory = 256;

        switch (flinkJob.getRunMode()) {
            case ALL_IN_ONE:
                jobMemory = tableMetas.collectInt(meta -> meta.getSyncYarn().getJobManagerMemory()).maxIfEmpty(jobMemory);
                jobMetaspaceMemory = tableMetas.size() * 25;
                taskMemory = (int) tableMetas.sumOfInt(meta -> meta.getSyncYarn().getTaskManagerMemory());
                break;
            case ALL_IN_ONE_BY_SCHEMA:
                jobMemory = tableMetas.collectInt(meta -> meta.getSyncYarn().getJobManagerMemory()).maxIfEmpty(jobMemory);
                jobMetaspaceMemory = tableMetas.distinctBy(TableMeta::getSchema).size() * 25;
                taskMemory = (int) tableMetas.sumByInt(TableMeta::getSchema, meta -> meta.getSyncYarn().getTaskManagerMemory()).max();
                break;
            case ALL_IN_ONE_BY_TABLE:
                jobMemory = tableMetas.collectInt(meta -> meta.getSyncYarn().getJobManagerMemory()).maxIfEmpty(jobMemory);
                jobMetaspaceMemory = tableMetas.distinctBy(TableMeta::getTable).size() * 25;
                taskMemory = (int) tableMetas.sumByInt(TableMeta::getTable, meta -> meta.getSyncYarn().getTaskManagerMemory()).max();
                break;
            case ONE_IN_ONE:
                // jobMemory = tableMetas.collectInt(meta -> meta.getSyncYarn().getJobManagerMemory()).maxIfEmpty(jobMemory);
                jobMemory = flinkJob.getOneInOneSyncYarn().getJobManagerMemory();
                jobMetaspaceMemory = tableMetas.size() * 25;
                taskMemory = flinkJob.getOneInOneSyncYarn().getTaskManagerMemory();
                break;
            default:
        }

        taskMetaspaceMemory = Math.min(taskMemory / 3, tableMetas.size() * 25);

        return new Memory(
                Math.max(512, jobMemory),
                Math.max(256, jobMetaspaceMemory),
                Math.max(512, taskMemory),
                Math.max(256, taskMetaspaceMemory)
        );
    }

    private String getLatestExecutorJarPath() throws IOException {
        try (FileSystem fileSystem = HadoopUtil.createFileSystem(HadoopUtil.createConfiguration(hadoopConfiguration))) {
            Path root = new Path(hudiConfiguration.getAppHdfsPath());
            return Lists.immutable.of(fileSystem.listStatus(root))
                    .select(FileStatus::isFile)
                    .collect(FileStatus::getPath)
                    .collect(Path::toString)
                    .select(path -> ReUtil.isMatch(EXECUTOR_JAR_NAME, path))
                    .collect(path -> new Tuple(path, getLatestExecutorJarVersion(path)))
                    .reject(tuple -> tuple.<Long>get(1) < 0)
                    .maxByOptional(tuple -> tuple.<Long>get(1))
                    .orElseThrow(() -> new RuntimeException("Latest jar path not found"))
                    .get(0);
        }
    }

    private Long getLatestExecutorJarVersion(String path) {
        return Optional.ofNullable(ReUtil.get(EXECUTOR_JAR_NAME, path, 1)).map(Long::valueOf).orElse(-1L);
    }

    private void setEnvironment(Configuration configuration, String key, String value) {
        configuration.setString(ResourceManagerOptions.CONTAINERIZED_MASTER_ENV_PREFIX + key, value);
        configuration.setString(ResourceManagerOptions.CONTAINERIZED_TASK_MANAGER_ENV_PREFIX + key, value);
    }

    private Configuration commonConfiguration(String keytabPath, String principal) {
        Configuration configuration = new Configuration();

        configuration.setBoolean(JobManagerOptions.JVM_DIRECT_MEMORY_LIMIT_ENABLED, true);
        // configuration.setString(JobManagerOptions.ARCHIVE_DIR, hudiConfiguration.getArchiveHdfsPath());

        configuration.setString(AkkaOptions.ASK_TIMEOUT, "10 min");
        configuration.setString(AkkaOptions.TCP_TIMEOUT, "15 min");
        configuration.setString(AkkaOptions.LOOKUP_TIMEOUT, "10 min");
        configuration.set(ClientOptions.CLIENT_TIMEOUT, Duration.ofMinutes(30));

        // Kerberos认证
        configuration.setBoolean(SecurityOptions.KERBEROS_LOGIN_USETICKETCACHE, true);
        if (StrUtil.isNotBlank(keytabPath)) {
            configuration.setString(SecurityOptions.KERBEROS_LOGIN_KEYTAB, keytabPath);
        }
        if (StrUtil.isNotBlank(principal)) {
            configuration.setString(SecurityOptions.KERBEROS_LOGIN_PRINCIPAL, principal);
        }

        // configuration.setString(HistoryServerOptions.HISTORY_SERVER_WEB_ADDRESS, "0.0.0.0");
        // configuration.setInteger(HistoryServerOptions.HISTORY_SERVER_WEB_PORT, 8083);

        configuration.setLong(HeartbeatManagerOptions.HEARTBEAT_INTERVAL, MINUTE);
        configuration.setLong(HeartbeatManagerOptions.HEARTBEAT_TIMEOUT, HALF_HOUR);

        configuration.setString(AkkaOptions.ASK_TIMEOUT, "1 min");
        configuration.setString(AkkaOptions.TCP_TIMEOUT, "2 min");

        configuration.setBoolean(CoreOptions.CHECK_LEAKED_CLASSLOADER, false);

        configuration.setString(YarnConfigOptions.APPLICATION_ATTEMPTS, "4");
        // configuration.setString(YarnConfigOptions.STAGING_DIRECTORY, "hdfs://b2/apps/datalake/yarn");

        configuration.setString(ResourceManagerOptions.CONTAINERIZED_MASTER_ENV_PREFIX + "MALLOC_ARENA_MAX", "1");
        configuration.setString(ResourceManagerOptions.CONTAINERIZED_TASK_MANAGER_ENV_PREFIX + "MALLOC_ARENA_MAX", "1");

        configuration.setString(RestOptions.BIND_PORT, "8084-9400");

        configuration.setString(DeploymentOptions.TARGET, YarnDeploymentTarget.APPLICATION.getName());

        if (StrUtil.isNotBlank(hudiConfiguration.getVictoriaPushUrl())) {
            configuration.setString("metrics.reporter.victoria.class", "com.eshore.odcp.hudi.connector.utils.executor.metrics.VictoriaMetricsReporter");
            configuration.setString("metrics.reporter.victoria.endpoint", hudiConfiguration.getVictoriaPushUrl());
            configuration.setBoolean("metrics.reporter.victoria.enable.auth", true);
            configuration.setString("metrics.reporter.victoria.auth.username", "EsCFVuNkiDWv7PKmcF");
            configuration.setString("metrics.reporter.victoria.auth.password", "Abf%x9ocS^iKr3tgrd");
        }

        // 设置 JVM 参数
        // configuration.setString(CoreOptions.FLINK_TM_JVM_OPTIONS, "-XX:CMSInitiatingOccupancyFraction=40");
        // configuration.setString(CoreOptions.FLINK_TM_JVM_OPTIONS, "-XX:+UseG1GC -XX:MaxGCPauseMillis=1500 -XX:InitiatingHeapOccupancyPercent=30 -XX:SurvivorRatio=4 -XX:ParallelGCThreads=60");

        return configuration;
    }

    public ApplicationId runSync(Long flinkJobId, String keytabPath, String principal, String cluster) throws Exception {
        FlinkJob flinkJob = infoService.flinkJobDetail(flinkJobId);
        ImmutableList<TableMeta> tableMetas = infoService.tableMetaList(flinkJobId);
        if (ObjectUtil.isEmpty(tableMetas)) {
            throw new RuntimeException("Table metas not found");
        }
        return runSync(flinkJob, tableMetas, keytabPath, principal, cluster);
    }

    public ApplicationId runSync(FlinkJob flinkJob, ImmutableList<TableMeta> tableMetaList, String keytabPath, String principal, String cluster) throws Exception {
        logger.info("Running Sync: {} {} {}", flinkJob.getId(), keytabPath, principal);

        try (FileSystem fileSystem = HadoopUtil.createFileSystem(HadoopUtil.createConfiguration(hadoopConfiguration))) {
            String checkpointRootPath = tableMetaList.collect(meta -> meta.getConfig().getCheckpointRootPath())
                    .distinct()
                    .getFirstOptional()
                    .orElseThrow(CheckpointRootPathNotFoundException::new);
            String checkpointPath = checkpointRootPath + "/" + flinkJob.getId();
            logger.info("Delete checkpoint path {}", checkpointPath);
            fileSystem.delete(new Path(checkpointPath), true);
        }

        Configuration configuration = commonConfiguration(keytabPath, principal);

        Memory memory = jobMemory(flinkJob, tableMetaList);
        // JobManager的内存
        MemorySize jobMemorySize = MemorySize.parse(memory.jobMemory + "m");
        configuration.set(JobManagerOptions.TOTAL_PROCESS_MEMORY, jobMemorySize);
        configuration.set(JobManagerOptions.JVM_METASPACE, MemorySize.parse(memory.jobMetaspaceMemory + "m"));

        // TaskManager的内存
        MemorySize taskMemorySize = MemorySize.parse(memory.taskMemory + "m");
        configuration.set(TaskManagerOptions.TOTAL_FLINK_MEMORY, taskMemorySize);
        configuration.set(TaskManagerOptions.MANAGED_MEMORY_SIZE, MemorySize.parse("0m"));
        configuration.set(TaskManagerOptions.JVM_METASPACE, MemorySize.parse(memory.taskMetaspaceMemory + "m"));

        // 并行度
        configuration.setInteger(TaskManagerOptions.NUM_TASK_SLOTS, 10);

        configuration.setString(YarnConfigOptions.APPLICATION_NAME, NameHelper.syncJobName(flinkJob.getId(), flinkJob.getName()));

        // configuration.setString(HistoryServerOptions.HISTORY_SERVER_ARCHIVE_DIRS, "hdfs://b2/apps/flink/completed-jobs/");
        // configuration.setLong(HistoryServerOptions.HISTORY_SERVER_ARCHIVE_REFRESH_INTERVAL, 10000);

        // 业务jar包
        String executorJarPath = getLatestExecutorJarPath();
        logger.info("Executor jar path: {}", executorJarPath);
        Long executorJarVersion = getLatestExecutorJarVersion(executorJarPath);
        configuration.set(PipelineOptions.JARS, new ArrayList<String>() {{
            add(executorJarPath);
        }});

        String flinkJobName = flinkJob.getName().replaceAll("\\s", "_");

        setEnvironment(configuration, Constants.METRICS_LABEL_RUN_TYPE, Constants.METRICS_RUN_TYPE_SYNC);
        setEnvironment(configuration, Constants.METRICS_LABEL_FLINK_JOB_ID, String.valueOf(flinkJob.getId()));
        setEnvironment(configuration, Constants.METRICS_LABEL_FLINK_JOB_NAME, flinkJobName);
        setEnvironment(configuration, Constants.METRICS_LABEL_EXECUTOR_VERSION, String.valueOf(executorJarVersion));
        setEnvironment(configuration, Constants.METRICS_LABEL_CLUSTER, cluster);

        setEnvironment(configuration, Constants.LOKI_PUSH_URL, hudiConfiguration.getLokiPushUrl());

        configuration.setString("metrics.reporter.victoria.tags", Lists.immutable.of(
                Pair.of(Constants.METRICS_LABEL_RUN_TYPE, Constants.METRICS_RUN_TYPE_SYNC),
                Pair.of(Constants.METRICS_LABEL_FLINK_JOB_ID, flinkJob.getId()),
                Pair.of(Constants.METRICS_LABEL_FLINK_JOB_NAME, flinkJobName),
                Pair.of(Constants.METRICS_LABEL_EXECUTOR_VERSION, executorJarVersion)
        ).collect(pair -> StrUtil.format("{}={}", pair.getLeft(), pair.getRight())).makeString(";"));

        return Runner.run(
                configuration,
                "com.eshore.odcp.hudi.connector.sync.Synchronizer",
                syncArgs(flinkJob, tableMetaList, cluster)
        );
    }

    public ApplicationId runCompaction(String batchCode, FlinkJob flinkJob, TableMeta tableMeta, String keytabPath, String principal, String instants, String cluster) throws Exception {
        logger.info("Running Compaction: {} {} {} {} {}", batchCode, flinkJob.getId(), tableMeta.getAlias(), keytabPath, principal);

        Configuration configuration = commonConfiguration(keytabPath, principal);

        MemorySize jobMemorySize = MemorySize.parse(tableMeta.getCompactionYarn().getJobManagerMemory() + "m");
        configuration.set(JobManagerOptions.TOTAL_PROCESS_MEMORY, jobMemorySize);

        MemorySize taskMemorySize = MemorySize.parse(tableMeta.getCompactionYarn().getTaskManagerMemory() + "m");
        configuration.set(TaskManagerOptions.TOTAL_FLINK_MEMORY, taskMemorySize);
        configuration.set(TaskManagerOptions.MANAGED_MEMORY_SIZE, MemorySize.parse("0m"));

        configuration.setInteger(TaskManagerOptions.NUM_TASK_SLOTS, 10);

        configuration.setString(YarnConfigOptions.APPLICATION_NAME, NameHelper.compactionJobName(flinkJob.getId(), tableMeta.getAlias()));

        String executorJarPath = getLatestExecutorJarPath();
        logger.info("Executor jar path: {}", executorJarPath);
        Long executorJarVersion = getLatestExecutorJarVersion(executorJarPath);
        configuration.set(PipelineOptions.JARS, new ArrayList<String>() {{
            add(executorJarPath);
        }});

        String flinkJobName = flinkJob.getName().replaceAll("\\s", "_");

        setEnvironment(configuration, Constants.METRICS_LABEL_RUN_TYPE, Constants.METRICS_RUN_TYPE_COMPACTION);
        setEnvironment(configuration, Constants.METRICS_LABEL_FLINK_JOB_ID, String.valueOf(flinkJob.getId()));
        setEnvironment(configuration, Constants.METRICS_LABEL_FLINK_JOB_NAME, flinkJobName);
        setEnvironment(configuration, Constants.METRICS_LABEL_SCHEMA, tableMeta.getSchema());
        setEnvironment(configuration, Constants.METRICS_LABEL_TABLE, tableMeta.getTable());
        setEnvironment(configuration, Constants.METRICS_LABEL_ALIAS, tableMeta.getAlias());
        setEnvironment(configuration, Constants.METRICS_LABEL_BATCH_ID, batchCode);
        setEnvironment(configuration, Constants.METRICS_LABEL_EXECUTOR_VERSION, String.valueOf(executorJarVersion));
        setEnvironment(configuration, Constants.METRICS_LABEL_CLUSTER, cluster);

        setEnvironment(configuration, Constants.LOKI_PUSH_URL, hudiConfiguration.getLokiPushUrl());

        configuration.setString("metrics.reporter.victoria.tags", Lists.immutable.of(
                Pair.of(Constants.METRICS_LABEL_RUN_TYPE, Constants.METRICS_RUN_TYPE_COMPACTION),
                Pair.of(Constants.METRICS_LABEL_FLINK_JOB_ID, flinkJob.getId()),
                Pair.of(Constants.METRICS_LABEL_FLINK_JOB_NAME, flinkJob.getName().replaceAll("\\s", "_")),
                Pair.of(Constants.METRICS_LABEL_SCHEMA, tableMeta.getSchema()),
                Pair.of(Constants.METRICS_LABEL_TABLE, tableMeta.getTable()),
                Pair.of(Constants.METRICS_LABEL_ALIAS, tableMeta.getAlias()),
                Pair.of(Constants.METRICS_LABEL_BATCH_ID, batchCode),
                Pair.of(Constants.METRICS_LABEL_EXECUTOR_VERSION, executorJarVersion)
        ).collect(pair -> StrUtil.format("{}={}", pair.getLeft(), pair.getRight())).makeString(";"));

        return Runner.run(
                configuration,
                "com.eshore.odcp.hudi.connector.sync.Compactor",
                compactionArgs(flinkJob, tableMeta, instants, cluster)
        );
    }

    private static final class Memory {
        private final int jobMemory;
        private final int jobMetaspaceMemory;
        private final int taskMemory;
        private final int taskMetaspaceMemory;

        public Memory(int jobMemory, int jobMetaspaceMemory, int taskMemory, int taskMetaspaceMemory) {
            this.jobMemory = jobMemory;
            this.jobMetaspaceMemory = jobMetaspaceMemory;
            this.taskMemory = taskMemory;
            this.taskMetaspaceMemory = taskMetaspaceMemory;
        }
    }
}
