package com.lanyuanxiaoyao.service.sync;

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.Record;
import com.lanyuanxiaoyao.service.common.entity.RunMeta;
import com.lanyuanxiaoyao.service.common.entity.TableMeta;
import com.lanyuanxiaoyao.service.common.exception.CheckpointRootPathNotFoundException;
import com.lanyuanxiaoyao.service.common.exception.ZookeeperUrlNotFoundException;
import com.lanyuanxiaoyao.service.common.utils.NameHelper;
import com.lanyuanxiaoyao.service.common.utils.TableMetaHelper;
import com.lanyuanxiaoyao.service.sync.configuration.GlobalConfiguration;
import com.lanyuanxiaoyao.service.sync.functions.PulsarMessage2RecordFunction;
import com.lanyuanxiaoyao.service.sync.functions.PulsarMessageSourceReader;
import com.lanyuanxiaoyao.service.sync.functions.ValidateRecordFilter;
import com.lanyuanxiaoyao.service.sync.utils.*;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.serialization.SimpleStringEncoder;
import org.apache.flink.core.fs.Path;
import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.filesystem.StreamingFileSink;
import org.apache.flink.streaming.api.functions.sink.filesystem.rollingpolicies.DefaultRollingPolicy;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static com.lanyuanxiaoyao.service.common.Constants.*;

/**
 * 同步应用
 *
 * @author lanyuanxiaoyao
 * @version 0.0.1
 * @date 2021-11-26
 */
public class Synchronizer {
    private static final Logger logger = LoggerFactory.getLogger(Synchronizer.class);
    private static final ObjectMapper mapper = JacksonUtils.getMapper();

    private static String findConfigFromList(List<TableMeta> metas, Function<TableMeta, String> getter, Supplier<Exception> notFoundException) throws Exception {
        return metas.stream()
                .map(getter)
                .distinct()
                .findFirst()
                .orElseThrow(notFoundException);
    }

    public static void main(String[] args) throws Exception {
        FlinkJob flinkJob = ArgumentsUtils.getFlinkJob(args);
        List<TableMeta> tableMetaList = ArgumentsUtils.getTableMetaList(args);
        String cluster = ArgumentsUtils.getCluster(args);

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

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

        String zkUrl = findConfigFromList(tableMetaList, meta -> meta.getConfig().getZookeeperUrl(), ZookeeperUrlNotFoundException::new);
        for (TableMeta tableMeta : tableMetaList) {
            RunMeta runMeta = new RunMeta(cluster, flinkJob.getId(), tableMeta.getAlias());
            logger.info("Run meta: {}", runMeta);
            ZkUtils.createSynchronizerLock(flinkJob, tableMeta, zkUrl, mapper.writeValueAsString(runMeta));
        }
        RunMeta runMeta = new RunMeta(cluster, flinkJob.getId());
        logger.info("Run meta: {}", runMeta);
        ZkUtils.createSynchronizerLock(flinkJob, zkUrl, mapper.writeValueAsString(runMeta));
        logger.info("Lock for {} success", flinkJob.getId());

        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();

        environment.enableCheckpointing(15 * MINUTE);
        String checkpointRootPath = findConfigFromList(tableMetaList, meta -> meta.getConfig().getCheckpointRootPath(), CheckpointRootPathNotFoundException::new);
        environment.getCheckpointConfig().setCheckpointStorage(new Path(checkpointRootPath + "/" + flinkJob.getId()));
        environment.getCheckpointConfig().setMaxConcurrentCheckpoints(1);
        environment.getCheckpointConfig().setCheckpointTimeout(2 * HOUR);
        environment.getCheckpointConfig().setMinPauseBetweenCheckpoints(15 * MINUTE);
        environment.getCheckpointConfig().setCheckpointingMode(CheckpointingMode.EXACTLY_ONCE);
        environment.getCheckpointConfig().setForceUnalignedCheckpoints(true);
        environment.getCheckpointConfig().enableUnalignedCheckpoints();
        environment.getCheckpointConfig().setTolerableCheckpointFailureNumber(5);

        if (tableMetaList.stream().anyMatch(meta -> TableMetaHelper.existsTag(meta, TAGS_DISABLE_CHAINING))) {
            logger.warn("Disable operator chaining");
            environment.disableOperatorChaining();
        }

        environment.setRestartStrategy(RestartStrategies.fixedDelayRestart(5, MINUTE));

        environment.setStateBackend(new HashMapStateBackend());

        switch (flinkJob.getRunMode()) {
            case ALL_IN_ONE:
                for (TableMeta tableMeta : tableMetaList) {
                    GlobalConfiguration globalConfiguration = new GlobalConfiguration(cluster, applicationId, tableMeta);
                    createFlinkJob(environment, globalConfiguration, flinkJob, tableMeta);
                    publishSyncStart(globalConfiguration, flinkJob, tableMeta);
                }
                environment.execute(NameHelper.syncFlinkName(flinkJob.getId(), flinkJob.getName()));
                break;
            case ONE_IN_ONE:
                for (TableMeta tableMeta : tableMetaList) {
                    GlobalConfiguration globalConfiguration = new GlobalConfiguration(cluster, applicationId, tableMeta);
                    createFlinkJob(environment, globalConfiguration, flinkJob, tableMeta);
                    publishSyncStart(globalConfiguration, flinkJob, tableMeta);
                    environment.execute(NameHelper.syncFlinkName(flinkJob.getId(), flinkJob.getName(), tableMeta.getAlias()));
                }
                break;
            case ALL_IN_ONE_BY_TABLE:
                scheduleOneInOneRegistryByField(environment, cluster, applicationId, flinkJob, tableMetaList, TableMeta::getTable);
                break;
            case ALL_IN_ONE_BY_SCHEMA:
                scheduleOneInOneRegistryByField(environment, cluster, applicationId, flinkJob, tableMetaList, TableMeta::getSchema);
                break;
            default:
                throw new IllegalArgumentException("Unsupported run mode: " + flinkJob.getRunMode());
        }
    }

    private static void scheduleOneInOneRegistryByField(StreamExecutionEnvironment environment, String cluster, String applicationId, FlinkJob flinkJob, List<TableMeta> tableMetaList, Function<TableMeta, String> field) throws Exception {
        Map<String, List<TableMeta>> map = tableMetaList.stream()
                .collect(Collectors.groupingBy(field));
        for (Map.Entry<String, List<TableMeta>> entry : map.entrySet()) {
            for (TableMeta tableMeta : entry.getValue()) {
                GlobalConfiguration globalConfiguration = new GlobalConfiguration(cluster, applicationId, tableMeta);
                createFlinkJob(environment, globalConfiguration, flinkJob, tableMeta);
                publishSyncStart(globalConfiguration, flinkJob, tableMeta);
            }
            environment.execute(NameHelper.syncFlinkName(flinkJob.getId(), flinkJob.getName(), entry.getKey()));
        }
    }

    private static void createFlinkJob(StreamExecutionEnvironment environment, GlobalConfiguration configuration, FlinkJob flinkJob, TableMeta tableMeta) throws IOException {
        logger.info("Table meta: {}", mapper.writeValueAsString(tableMeta));
        logger.info("Config meta: {}", mapper.writeValueAsString(configuration));
        SingleOutputStreamOperator<String> source = environment
                .addSource(new PulsarMessageSourceReader(configuration, flinkJob, tableMeta))
                .setParallelism(tableMeta.getHudi().getSourceTasks());
        if (TableMetaHelper.existsTag(tableMeta, Constants.TAGS_PULSAR_BACKUP)) {
            Path path = new Path(StrUtil.format("hdfs://b2/apps/datalake/hive_test/source/{}/{}", String.join("_", flinkJob.getName().split("\\s")), tableMeta.getAlias()));
            StreamingFileSink<String> fileSink = StreamingFileSink.<String>forRowFormat(path, new SimpleStringEncoder<>("UTF-8"))
                    .withRollingPolicy(DefaultRollingPolicy.builder()
                            .withInactivityInterval(HOUR)
                            .withMaxPartSize(GB)
                            .build())
                    .build();
            source.addSink(fileSink).name("Backup pulsar data");
        }
        SingleOutputStreamOperator<Record> middle = source
                .map(new PulsarMessage2RecordFunction(configuration, flinkJob, tableMeta))
                .name("Json ( " + tableMeta.getSchema() + "-" + tableMeta.getAlias() + " )")
                .filter(new ValidateRecordFilter(configuration, flinkJob, tableMeta))
                .name("Reject json parse failure");
        SyncUtils.sinkToHoodieByTable(configuration, flinkJob, tableMeta, environment, middle);
    }

    private static void publishSyncStart(GlobalConfiguration configuration, FlinkJob flinkJob, TableMeta tableMeta) {
        StatusUtils.syncStart(configuration, flinkJob, tableMeta);
    }
}
