package com.lanyuanxiaoyao.service.executor.task;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.lanyuanxiaoyao.service.executor.core.TaskContext;
import com.lanyuanxiaoyao.service.executor.task.entity.RecordView;
import com.lanyuanxiaoyao.service.executor.task.functions.ReadHudiFile;
import com.lanyuanxiaoyao.service.executor.task.functions.pulsar.ReadPulsarSource;
import com.lanyuanxiaoyao.service.executor.task.helper.ArgumentsHelper;
import com.lanyuanxiaoyao.service.executor.task.helper.FlinkHelper;
import com.lanyuanxiaoyao.service.executor.task.helper.HdfsHelper;
import java.util.Map;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.eclipse.collections.api.list.ImmutableList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * avro 日志扫描
 *
 * @author lanyuanxiaoyao
 * @date 2024-01-08
 */
public class DataScanner {
    private static final Logger logger = LoggerFactory.getLogger(DataScanner.class);

    public static void main(String[] args) throws Exception {
        TaskContext taskContext = ArgumentsHelper.getContext(args);
        logger.info("Context: {}", taskContext);

        Map<String, Object> metadata = taskContext.getMetadata();
        ArgumentsHelper.checkMetadata(taskContext, "key");
        String key = (String) metadata.get("key");
        Boolean scanQueue = (Boolean) metadata.getOrDefault("scan_queue", false);
        Boolean scanLog = (Boolean) metadata.getOrDefault("scan_log", false);
        Boolean scanBase = (Boolean) metadata.getOrDefault("scan_base", false);

        if (!scanQueue && !scanLog && !scanBase) {
            throw new RuntimeException("Must choose mode scan_queue or scan_log or scan_data");
        }

        StreamExecutionEnvironment environment = FlinkHelper.getBatchEnvironment();

        DataStream<RecordView> source = null;
        int totalParallelism = 30;
        if (scanQueue) {
            ArgumentsHelper.checkMetadata(taskContext, "pulsar");
            String pulsarUrl = (String) metadata.get("pulsar");
            ArgumentsHelper.checkMetadata(taskContext, "pulsar_topic");
            String pulsarTopic = (String) metadata.get("pulsar_topic");
            logger.info("Scan queue topic: {} url: {}", pulsarTopic, pulsarUrl);
            DataStream<RecordView> stream = environment
                    .fromSource(new ReadPulsarSource(taskContext, pulsarUrl, pulsarTopic), WatermarkStrategy.noWatermarks(), "Read pulsar")
                    .setParallelism(totalParallelism)
                    .disableChaining();
            if (ObjectUtil.isNull(source)) {
                source = stream;
            } else {
                source = source.union(stream);
            }
        }
        if (scanLog || scanBase) {
            ArgumentsHelper.checkMetadata(taskContext, "hdfs");
            String hdfs = (String) metadata.get("hdfs");
            FileSystem fileSystem = FileSystem.get(new Configuration());
            HdfsHelper.checkHdfsPath(fileSystem, hdfs);

            if (scanLog) {
                logger.info("Scan log hdfs: {}", hdfs);
                ImmutableList<String> logPaths = HdfsHelper.logPaths(fileSystem, hdfs);
                int parallelism = HdfsHelper.logScanParallelismPredict(logPaths);
                totalParallelism = Math.max(totalParallelism, parallelism);
                DataStream<RecordView> stream = environment
                        .fromCollection(logPaths.toList())
                        .name("Read log paths")
                        .flatMap(new ReadHudiFile())
                        .name("Read hudi file")
                        .setParallelism(parallelism);
                if (ObjectUtil.isNull(source)) {
                    source = stream;
                } else {
                    source = source.union(stream);
                }
            }
            if (scanBase) {
                logger.info("Scan base hdfs: {}", hdfs);
                ImmutableList<String> basePaths = HdfsHelper.latestBasePaths(fileSystem, hdfs);
                int parallelism = HdfsHelper.baseScanParallelismPredict(basePaths);
                totalParallelism = Math.max(totalParallelism, parallelism);
                DataStream<RecordView> stream = environment
                        .fromCollection(basePaths.toList())
                        .name("Read base paths")
                        .flatMap(new ReadHudiFile())
                        .name("Read hudi file")
                        .setParallelism(parallelism);
                if (ObjectUtil.isNull(source)) {
                    source = stream;
                } else {
                    source = source.union(stream);
                }
            }
        }
        if (ObjectUtil.isNull(source)) {
            throw new RuntimeException("Source cannot be null");
        }

        environment.setParallelism(Math.max(50, Math.min(totalParallelism / 2, 200)));

        source
                .map(RecordView::toString)
                .name("Covert record to string")
                .filter(line -> StrUtil.contains(line, key))
                .name("Filter target key")
                .sinkTo(FlinkHelper.createFileSink(taskContext))
                .setParallelism(10)
                .name("Output results");
        environment.execute(StrUtil.format("Search {}", key));
    }
}
