package com.lanyuanxiaoyao.service.sync;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.lanyuanxiaoyao.service.common.entity.FlinkJob;
import com.lanyuanxiaoyao.service.common.entity.RunMeta;
import com.lanyuanxiaoyao.service.common.entity.TableMeta;
import com.lanyuanxiaoyao.service.common.utils.NameHelper;
import com.lanyuanxiaoyao.service.sync.configuration.GlobalConfiguration;
import com.lanyuanxiaoyao.service.sync.functions.CompactionEventHandler;
import com.lanyuanxiaoyao.service.sync.utils.*;
import java.io.IOException;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.operators.ProcessOperator;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hudi.avro.model.HoodieCompactionPlan;
import org.apache.hudi.client.HoodieFlinkWriteClient;
import org.apache.hudi.common.model.HoodieLogFile;
import org.apache.hudi.common.table.HoodieTableMetaClient;
import org.apache.hudi.common.table.TableSchemaResolver;
import org.apache.hudi.common.table.log.HoodieLogFormat;
import org.apache.hudi.common.table.log.block.HoodieAvroDataBlock;
import org.apache.hudi.common.table.log.block.HoodieDeleteBlock;
import org.apache.hudi.common.table.log.block.HoodieLogBlock;
import org.apache.hudi.common.table.log.block.HoodieParquetDataBlock;
import org.apache.hudi.common.table.timeline.HoodieInstant;
import org.apache.hudi.common.table.timeline.HoodieTimeline;
import org.apache.hudi.common.util.ClosableIterator;
import org.apache.hudi.common.util.CompactionUtils;
import org.apache.hudi.common.util.collection.Pair;
import org.apache.hudi.configuration.FlinkOptions;
import org.apache.hudi.exception.HoodieException;
import org.apache.hudi.org.apache.avro.Schema;
import org.apache.hudi.org.apache.avro.generic.IndexedRecord;
import org.apache.hudi.sink.compact.*;
import org.apache.hudi.sink.compact.strategy.CompactionPlanStrategies;
import org.apache.hudi.sink.compact.strategy.CompactionPlanStrategy;
import org.apache.hudi.table.HoodieFlinkTable;
import org.apache.hudi.util.CompactionUtil;
import org.apache.hudi.util.StreamerUtil;
import org.apache.parquet.avro.AvroSchemaConverter;
import org.apache.parquet.schema.MessageType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 单表离线压缩
 *
 * @author ZhangJiacheng
 * @date 2022-06-21
 */
public class Compactor {
    private static final Logger logger = LoggerFactory.getLogger(Compactor.class);
    private static final ObjectMapper mapper = JacksonUtils.getMapper();

    public static void main(String[] args) throws Exception {
        FlinkJob flinkJob = ArgumentsUtils.getFlinkJob(args);
        TableMeta tableMeta = ArgumentsUtils.getTableMeta(args);
        String selectedInstants = ArgumentsUtils.getInstants(args);
        String cluster = ArgumentsUtils.getCluster(args);

        logger.info("Bootstrap flink job: {}", mapper.writeValueAsString(flinkJob));
        logger.info("Bootstrap table meta: {}", mapper.writeValueAsString(tableMeta));
        logger.info("Bootstrap instants: {}", selectedInstants);
        logger.info("Bootstrap cluster: {}", cluster);

        String applicationId = System.getenv("_APP_ID");

        RunMeta runMeta = new RunMeta(cluster, flinkJob.getId(), tableMeta.getAlias());
        logger.info("Run meta: {}", runMeta);
        ZkUtils.createCompactionLock(flinkJob, tableMeta, tableMeta.getConfig().getZookeeperUrl(), mapper.writeValueAsString(runMeta));
        logger.info("Lock for {} {} success", flinkJob.getId(), tableMeta.getAlias());

        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();

        FlinkCompactionConfig config = new FlinkCompactionConfig();
        if (ObjectUtil.isEmpty(selectedInstants)) {
            config.compactionPlanSelectStrategy = CompactionPlanStrategy.ALL;
        } else {
            config.compactionPlanSelectStrategy = CompactionPlanStrategy.INSTANTS;
            config.compactionPlanInstant = selectedInstants;
        }

        GlobalConfiguration globalConfiguration = new GlobalConfiguration(cluster, applicationId, tableMeta);
        Configuration configuration = SyncUtils.getCompactionFlinkConfiguration(
                globalConfiguration,
                new Configuration(),
                flinkJob,
                tableMeta,
                SyncUtils.avroSchemaWithExtraFields(tableMeta),
                1
        );

        CompactionEventHandler eventHandler = new CompactionEventHandler(globalConfiguration, flinkJob, tableMeta);

        HoodieFlinkWriteClient<?> writeClient = StreamerUtil.createWriteClient(configuration);
        HoodieFlinkTable<?> table = writeClient.getHoodieTable();

        table.getMetaClient().reloadActiveTimeline();

        StatusUtils.compactionStart(globalConfiguration, flinkJob, tableMeta);

        // 检查时间线当前状态
        logger.info("{} timeline detail ({})", tableMeta.getAlias(), LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")));
        table.getActiveTimeline()
                .getInstants()
                .forEach(instant -> logger.info("{} {} {}", instant.getTimestamp(), instant.getAction(), instant.getState()));

        HoodieTimeline pendingCompactionTimeline = table.getActiveTimeline().filterPendingCompactionTimeline();
        List<HoodieInstant> requested = CompactionPlanStrategies.getStrategy(config).select(pendingCompactionTimeline);
        if (requested.isEmpty()) {
            logger.info("No compaction plan scheduled");
            eventHandler.closed("No compaction plan scheduled", null);
            return;
        }

        List<String> compactionInstantTimes = requested.stream().map(HoodieInstant::getTimestamp).collect(Collectors.toList());
        compactionInstantTimes.forEach(timestamp -> {
            HoodieInstant inflightInstant = HoodieTimeline.getCompactionInflightInstant(timestamp);
            if (pendingCompactionTimeline.containsInstant(inflightInstant)) {
                logger.info("Rollback inflight compaction instant: [" + timestamp + "]");
                table.rollbackInflightCompaction(inflightInstant);
                table.getMetaClient().reloadActiveTimeline();
            }
        });

        List<Pair<String, HoodieCompactionPlan>> compactionPlans = compactionInstantTimes.stream()
                .map(timestamp -> {
                    try {
                        return Pair.of(timestamp, CompactionUtils.getCompactionPlan(table.getMetaClient(), timestamp));
                    } catch (IOException e) {
                        throw new HoodieException("Get compaction plan at instant " + timestamp + " error", e);
                    }
                })
                .filter(pair -> {
                    HoodieCompactionPlan plan = pair.getRight();
                    return plan != null && plan.getOperations() != null && plan.getOperations().size() > 0;
                })
                .collect(Collectors.toList());

        if (compactionPlans.isEmpty()) {
            logger.info("No compaction plan for instant " + String.join(",", compactionInstantTimes));
            eventHandler.closed("No compaction plan for instant " + String.join(",", compactionInstantTimes), null);
            return;
        }

        List<HoodieInstant> instants = compactionInstantTimes.stream().map(HoodieTimeline::getCompactionRequestedInstant).collect(Collectors.toList());

        logger.info("Start to compaction for instant " + compactionInstantTimes);

        for (HoodieInstant instant : instants) {
            if (!pendingCompactionTimeline.containsInstant(instant)) {
                CompactionUtil.cleanInstant(table.getMetaClient(), instant);
            }
            table.getActiveTimeline().transitionCompactionRequestedToInflight(instant);
        }
        table.getMetaClient().reloadActiveTimeline();

        try {
            compactionPlans.forEach(pair -> preCommit(globalConfiguration, flinkJob, tableMeta, table, pair.getLeft(), pair.getRight()));
        } catch (Throwable e) {
            logger.warn("Cannot submit pre-commit log");
        }

        environment.addSource(new CompactionPlanSourceFunction(compactionPlans))
                .name("compaction_source")
                .uid("uid_compaction_source")
                .rebalance()
                .transform("compact_task",
                        TypeInformation.of(CompactionCommitEvent.class),
                        new ProcessOperator<>(new CompactFunction(configuration)))
                .setParallelism(configuration.getInteger(FlinkOptions.COMPACTION_TASKS))
                .addSink(new CompactionCommitSink(configuration, eventHandler))
                .name("compaction_commit")
                .uid("uid_compaction_commit")
                .setParallelism(1);

        environment.execute(NameHelper.compactionFlinkName(flinkJob.getId(), tableMeta.getSchema(), tableMeta.getAlias()));
    }

    private static void preCommit(GlobalConfiguration configuration, FlinkJob flinkJob, TableMeta tableMeta, HoodieFlinkTable<?> table, String instant, HoodieCompactionPlan compactionPlan) {
        HoodieTableMetaClient client = table.getMetaClient();
        FileSystem fileSystem = client.getRawFs();
        List<Path> deltaFilePaths = compactionPlan.getOperations()
                .stream()
                .flatMap(operation -> operation.getDeltaFilePaths().stream().map(path -> StrUtil.format("{}/{}", operation.getPartitionPath(), path)))
                .map(path -> new Path(StrUtil.format("{}/{}", client.getBasePathV2(), path)))
                .collect(Collectors.toList());
        AtomicLong totalLogFilesCompacted = new AtomicLong(0);
        AtomicLong totalLogFilesSize = new AtomicLong(0);
        AtomicLong totalRecordsDeleted = new AtomicLong(0);
        AtomicLong totalCompactedRecordsUpdated = new AtomicLong(0);
        deltaFilePaths.parallelStream().forEach(path -> {
            try {
                FileStatus fileStatus = fileSystem.getFileStatus(path);
                totalLogFilesCompacted.incrementAndGet();
                totalLogFilesSize.addAndGet(fileStatus.getLen());
                MessageType messageType = TableSchemaResolver.readSchemaFromLogFile(fileSystem, path);
                if (ObjectUtil.isNull(messageType)) {
                    return;
                }
                Schema writerSchema = new AvroSchemaConverter().convert(messageType);
                try (HoodieLogFormat.Reader reader = HoodieLogFormat.newReader(fileSystem, new HoodieLogFile(path), writerSchema)) {
                    while (reader.hasNext()) {
                        HoodieLogBlock block = reader.next();
                        switch (block.getBlockType()) {
                            case AVRO_DATA_BLOCK:
                                HoodieAvroDataBlock avroDataBlock = (HoodieAvroDataBlock) block;
                                try (ClosableIterator<IndexedRecord> avroDataBlockRecordIterator = avroDataBlock.getRecordIterator()) {
                                    while (avroDataBlockRecordIterator.hasNext()) {
                                        // logger.info("totalCompactedRecordsUpdated {}", totalCompactedRecordsUpdated.incrementAndGet());
                                        totalCompactedRecordsUpdated.incrementAndGet();
                                        avroDataBlockRecordIterator.next();
                                    }
                                }
                                break;
                            case PARQUET_DATA_BLOCK:
                                HoodieParquetDataBlock parquetDataBlock = (HoodieParquetDataBlock) block;
                                try (ClosableIterator<IndexedRecord> parquetDataBlockRecordIterator = parquetDataBlock.getRecordIterator()) {
                                    while (parquetDataBlockRecordIterator.hasNext()) {
                                        // logger.info("totalCompactedRecordsUpdated {}", totalCompactedRecordsUpdated.incrementAndGet());
                                        totalCompactedRecordsUpdated.incrementAndGet();
                                        parquetDataBlockRecordIterator.next();
                                    }
                                }
                                break;
                            case DELETE_BLOCK:
                                HoodieDeleteBlock deleteBlock = (HoodieDeleteBlock) block;
                                // logger.info("totalRecordsDeleted {}", totalRecordsDeleted.addAndGet(deleteBlock.getRecordsToDelete().length));
                                totalRecordsDeleted.addAndGet(deleteBlock.getRecordsToDelete().length);
                                break;
                            default:
                                break;
                        }
                    }
                }
            } catch (Exception e) {
                logger.warn("Parse log file failure for " + path, e);
            }
        });

        Map<String, Long> metadata = new ConcurrentHashMap<>(5);
        metadata.put("totalLogFilesCompacted", totalLogFilesCompacted.get());
        metadata.put("totalLogFilesSize", totalLogFilesSize.get());
        metadata.put("totalRecordsDeleted", totalRecordsDeleted.get());
        metadata.put("totalCompactedRecordsUpdated", totalCompactedRecordsUpdated.get());
        metadata.put("totalLogRecordsCompacted", totalRecordsDeleted.get() + totalCompactedRecordsUpdated.get());

        StatusUtils.compactionPreCommit(configuration, flinkJob, tableMeta, instant, metadata);
    }
}
