package com.lanyuanxiaoyao.service.launcher.compaction.service;

import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
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.entity.compaction.ScheduleJob;
import com.lanyuanxiaoyao.service.common.utils.LogHelper;
import com.lanyuanxiaoyao.service.common.utils.NameHelper;
import com.lanyuanxiaoyao.service.common.utils.TableMetaHelper;
import com.lanyuanxiaoyao.service.configuration.entity.hudi.HudiCompactionPlan;
import com.lanyuanxiaoyao.service.configuration.entity.hudi.HudiInstant;
import com.lanyuanxiaoyao.service.configuration.entity.queue.QueueItem;
import com.lanyuanxiaoyao.service.configuration.utils.QueueUtil;
import com.lanyuanxiaoyao.service.forest.service.HudiService;
import com.lanyuanxiaoyao.service.forest.service.InfoService;
import com.lanyuanxiaoyao.service.launcher.ExecutorService;
import com.lanyuanxiaoyao.service.launcher.configuration.ClusterConfiguration;
import com.lanyuanxiaoyao.service.launcher.configuration.HadoopConfiguration;
import com.lanyuanxiaoyao.service.launcher.configuration.ZookeeperConfiguration;
import dev.failsafe.Failsafe;
import dev.failsafe.RetryPolicy;
import java.time.Duration;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import javax.annotation.PreDestroy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.locks.InterProcessLock;
import org.apache.curator.framework.recipes.locks.InterProcessMutex;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.data.Stat;
import org.eclipse.collections.api.factory.Lists;
import org.eclipse.collections.api.list.ImmutableList;
import org.eclipse.collections.api.list.MutableList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.http.converter.json.Jackson2ObjectMapperBuilder;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;

/**
 * @author ZhangJiacheng
 * @date 2023-05-08
 */
@Service
public class CompactionService {
    private static final Logger logger = LoggerFactory.getLogger(CompactionService.class);
    private static final java.util.concurrent.ExecutorService EXECUTOR = Executors.newWorkStealingPool(20);
    private static final RetryPolicy<String> RETRY_POLICY = RetryPolicy.<String>builder()
            .handle(Exception.class)
            .withDelay(Duration.ofSeconds(1))
            .withMaxAttempts(3)
            .build();
    private static final int FILE_PER_PARALLELISM = 2;
    private static final int MAX_PARALLELISM = 500;
    private static final int MIN_PARALLELISM = 10;
    private final HadoopConfiguration hadoopConfiguration;
    private final ClusterConfiguration clusterConfiguration;
    private final DiscoveryClient discoveryClient;
    private final CuratorFramework zookeeperClient;
    private final InfoService infoService;
    private final HudiService hudiService;
    private final ExecutorService executorService;
    private final ObjectMapper mapper;

    @SuppressWarnings("SpringJavaInjectionPointsAutowiringInspection")
    public CompactionService(HadoopConfiguration hadoopConfiguration, ClusterConfiguration clusterConfiguration, ZookeeperConfiguration zookeeperConfiguration, DiscoveryClient discoveryClient, InfoService infoService, HudiService hudiService, ExecutorService executorService, Jackson2ObjectMapperBuilder builder) {
        this.hadoopConfiguration = hadoopConfiguration;
        this.clusterConfiguration = clusterConfiguration;
        this.discoveryClient = discoveryClient;
        this.infoService = infoService;
        this.hudiService = hudiService;
        this.executorService = executorService;
        mapper = builder.build();
        this.zookeeperClient = CuratorFrameworkFactory.builder()
                .connectString(zookeeperConfiguration.getConnectUrl())
                .retryPolicy(new ExponentialBackoffRetry((int) Constants.SECOND, 3))
                .sessionTimeoutMs((int) Constants.SECOND)
                .connectionTimeoutMs((int) (30 * Constants.SECOND))
                .build();
        this.zookeeperClient.start();
    }

    @PreDestroy
    public void destroy() {
        if (ObjectUtil.isNotNull(this.zookeeperClient)) {
            this.zookeeperClient.close();
        }
    }

    @Scheduled(fixedDelay = 5, initialDelay = 30, timeUnit = TimeUnit.SECONDS)
    public void scheduleCompact() {
        java.util.concurrent.ExecutorService threadPool = Executors.newWorkStealingPool(5);
        if (Boolean.FALSE.equals(QueueUtil.isEmpty(discoveryClient, clusterConfiguration.getCompactionQueueName()))) {
            for (int index = 0; index < 5; index++) {
                threadPool.submit(this::compact);
            }
            threadPool.shutdown();
            while (!threadPool.isTerminated()) {
                ThreadUtil.safeSleep(10 * Constants.SECOND);
            }
            logger.info("Finish compaction schedule");
        }
    }

    @SuppressWarnings("DataFlowIssue")
    private void compact() {
        // 每次把由于各种原因无法运行的任务先放在暂存 list 里面，直到遇到一个能执行的任务，就把这批任务送回队列
        MutableList<QueueItem<ScheduleJob>> holder = Lists.mutable.empty();

        QueueItem<ScheduleJob> item = QueueUtil.poll(discoveryClient, this.mapper, clusterConfiguration.getCompactionQueueName());
        // 检查是否获取到新的任务，如果队列为空，获取到的任务也是空
        while (ObjectUtil.isNotNull(item)) {
            LogHelper.setMdc(Constants.LOG_JOB_ID_LABEL, item.getTraceId());
            try {
                // 检查任务内容是否存在
                if (ObjectUtil.isNotNull(item.getData())) {
                    ScheduleJob job = item.getData();
                    LogHelper.setMdcFlinkJobAndAlias(job.getFlinkJobId(), job.getAlias());
                    logger.info("Receive job[{}]({}): {}", item.getTraceId(), item.getCreateTime(), item.getData());
                    // 构造任务相关的锁
                    String lockPath = NameHelper.compactionLauncherLockPath(job.getFlinkJobId(), job.getAlias());
                    InterProcessLock lock = new InterProcessMutex(zookeeperClient, lockPath);
                    try {
                        if (lock.acquire(2, TimeUnit.SECONDS)) {
                            Stat stat = zookeeperClient.checkExists().forPath(NameHelper.compactionRunningLockPath(job.getFlinkJobId(), job.getAlias()));
                            if (ObjectUtil.isNotNull(stat)) {
                                logger.info("Job {} {} is running", job.getFlinkJobId(), job.getAlias());
                                // 运行中的任务放在持有容器中
                                holder.add(item);
                                // 进入下一轮，由于最外层有一个 finally，所以直接 continue 也会尝试获取新的任务
                                continue;
                            }
                            FlinkJob flinkJob = infoService.flinkJobDetail(job.getFlinkJobId());
                            TableMeta meta = infoService.tableMetaDetail(job.getFlinkJobId(), job.getAlias());
                            if (TableMetaHelper.existsTag(meta, Constants.TAGS_NO_COMPACT)) {
                                logger.warn("[{}] [{}] Table tags no compact", flinkJob.getId(), meta.getAlias());
                                clearHolder(holder);
                                continue;
                            }
                            logger.info("[{}] [{}] Execute job", flinkJob.getId(), meta.getAlias());
                            // 判断是否存在 Hudi 表，提前结束掉
                            if (!hudiService.existsHudiTable(flinkJob.getId(), meta.getAlias())) {
                                logger.info("[{}] [{}] Hudi table not found", flinkJob.getId(), meta.getAlias());
                                clearHolder(holder);
                                continue;
                            }
                            // 获取待压缩的时间点
                            ImmutableList<HudiInstant> selectedInstants = hudiService.timelinePendingCompactionList(flinkJob.getId(), meta.getAlias());
                            if (ObjectUtil.isEmpty(selectedInstants)) {
                                logger.info("[{}] [{}] Table not need to compact", flinkJob.getId(), meta.getAlias());
                                clearHolder(holder);
                                continue;
                            }
                            logger.info("[{}] [{}] Selected Instants: {}", flinkJob.getId(), meta.getAlias(), selectedInstants.makeString(","));
                            // 计算待压缩的文件数
                            long count = predictCompactFileCount(meta, selectedInstants);
                            if (ObjectUtil.isNotNull(count)) {
                                // 根据待压缩的文件数计算并行度
                                long parallelism = predictParallelism(count);
                                logger.info("[{}] [{}] Predict compact files: {} {}", flinkJob.getId(), meta.getAlias(), count, parallelism);
                                meta.getHudi().setCompactionTasks((int) parallelism);
                            }
                            logger.info("[{}] [{}] Execution", flinkJob.getId(), meta.getAlias());
                            String applicationId = Failsafe.with(RETRY_POLICY)
                                    .get(() -> executorService.runCompaction(
                                            job.getBatch(),
                                            flinkJob,
                                            meta,
                                            hadoopConfiguration.getKerberosKeytabPath(),
                                            hadoopConfiguration.getKerberosPrincipal(),
                                            selectedInstants.collect(HudiInstant::getTimestamp).makeString(","),
                                            clusterConfiguration.getCluster()
                                    ).toString());
                            Failsafe.with(RETRY_POLICY)
                                    .run(() -> infoService.saveCompactionId(flinkJob.getId(), meta.getAlias(), applicationId));
                            clearHolder(holder);
                        } else {
                            logger.warn("Un acquire lock for " + item.getId());
                            holder.add(item);
                        }
                    } catch (Exception e) {
                        logger.warn(StrUtil.format("[{}] [{}] Try lock something wrong ", job.getFlinkJobId(), job.getAlias()), e);
                        String failCount = item.getMetadata(Constants.SCHEDULE_JOB_FAIL_COUNT);
                        if (StrUtil.isNotBlank(failCount)) {
                            int fail = Integer.parseInt(failCount);
                            if (fail > 5) {
                                logger.error("Job {} cause unaccepted error", item);
                                continue;
                            } else {
                                item.getMetadata().put(Constants.SCHEDULE_JOB_FAIL_COUNT, String.valueOf(fail + 1));
                            }
                        } else {
                            item.getMetadata().put(Constants.SCHEDULE_JOB_FAIL_COUNT, "1");
                        }
                        QueueUtil.add(discoveryClient, this.mapper, Constants.COMPACTION_QUEUE_PRE, item);
                    } finally {
                        // 无论如何，尝试解锁
                        try {
                            if (lock.isAcquiredInThisProcess()) {
                                lock.release();
                            }
                        } catch (Exception e) {
                            logger.error("Release lock failure " + lockPath, e);
                        }
                    }
                } else {
                    logger.warn("Schedule job is empty. [{}]({}): {}", item.getTraceId(), item.getCreateTime(), item);
                }
            } finally {
                // 无论如何尝试获取下个任务
                item = QueueUtil.poll(discoveryClient, this.mapper, clusterConfiguration.getCompactionQueueName());
                LogHelper.removeMdc(Constants.LOG_JOB_ID_LABEL, Constants.LOG_FLINK_JOB_ID_LABEL, Constants.LOG_ALIAS_LABEL);
            }
        }
        clearHolder(holder);
    }

    private void clearHolder(MutableList<QueueItem<ScheduleJob>> holder) {
        if (holder.isEmpty()) {
            return;
        }
        // 等待半分钟
        ThreadUtil.safeSleep(Constants.HALF_MINUTE);
        for (QueueItem<ScheduleJob> item : holder) {
            QueueUtil.add(discoveryClient, this.mapper, Constants.COMPACTION_QUEUE_PRE, item);
        }
        holder.clear();
    }

    private Long predictCompactFileCount(TableMeta meta, ImmutableList<HudiInstant> selectedInstants) {
        long count = selectedInstants.asParallel(EXECUTOR, 1)
                .sumOfLong(instant -> {
                    try {
                        HudiCompactionPlan plan = hudiService.readCompactionPlan(meta.getJob().getId(), meta.getAlias(), instant.getTimestamp());
                        return plan.getOperations().size();
                    } catch (Exception e) {
                        logger.error("Read compaction plan failure", e);
                    }
                    return 0;
                });
        return count < 1 ? meta.getHudi().getBucketIndexNumber() : count;
    }

    private Integer predictParallelism(Long count) {
        long parallelism = Math.round(count * 1.0 / FILE_PER_PARALLELISM);
        if (parallelism > MAX_PARALLELISM) {
            parallelism = MAX_PARALLELISM;
        }
        if (parallelism < MIN_PARALLELISM) {
            parallelism = MIN_PARALLELISM;
        }
        return Math.toIntExact(parallelism);
    }

    private QueueItem<ScheduleJob> deserialize(String body) {
        if (StrUtil.isBlank(body)) {
            return null;
        }
        try {
            return mapper.readValue(body, new TypeReference<QueueItem<ScheduleJob>>() {
            });
        } catch (Throwable error) {
            logger.error("Schedule job parse error", error);
            return null;
        }
    }
}
