package com.lanyuanxiaoyao.service.command.pro.commands;

import cn.hutool.core.util.StrUtil;
import com.lanyuanxiaoyao.service.common.entity.TableMeta;
import com.lanyuanxiaoyao.service.configuration.ExecutorProvider;
import com.lanyuanxiaoyao.service.configuration.entity.hudi.HudiCompactionPlan;
import com.lanyuanxiaoyao.service.configuration.entity.hudi.HudiInstant;
import com.lanyuanxiaoyao.service.forest.service.HudiService;
import com.lanyuanxiaoyao.service.forest.service.InfoService;
import java.io.IOException;
import java.time.Instant;
import java.util.List;
import java.util.concurrent.atomic.LongAdder;
import java.util.stream.Collectors;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.LocatedFileStatus;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.fs.RemoteIterator;
import org.apache.hudi.common.fs.ConsistencyGuardConfig;
import org.apache.hudi.common.fs.FSUtils;
import org.apache.hudi.common.model.BaseFile;
import org.apache.hudi.common.model.HoodieBaseFile;
import org.apache.hudi.common.model.HoodieFileFormat;
import org.apache.hudi.common.model.HoodieFileGroup;
import org.apache.hudi.common.table.HoodieTableMetaClient;
import org.apache.hudi.common.table.timeline.HoodieInstant;
import org.apache.hudi.common.table.timeline.HoodieTimeline;
import org.apache.hudi.common.table.timeline.versioning.TimelineLayoutVersion;
import org.apache.hudi.common.table.view.HoodieTableFileSystemView;
import org.apache.hudi.common.util.BaseFileUtils;
import org.apache.hudi.common.util.ClosableIterator;
import org.apache.hudi.common.util.Option;
import org.apache.hudi.common.util.ParquetUtils;
import org.apache.hudi.io.storage.HoodieParquetReader;
import org.apache.hudi.org.apache.avro.Schema;
import org.apache.hudi.org.apache.avro.generic.GenericRecord;
import org.apache.hudi.org.apache.avro.generic.IndexedRecord;
import org.apache.parquet.avro.AvroParquetReader;
import org.apache.parquet.avro.AvroReadSupport;
import org.apache.parquet.example.data.Group;
import org.apache.parquet.hadoop.ParquetReader;
import org.apache.parquet.hadoop.example.GroupReadSupport;
import org.apache.parquet.hadoop.util.HadoopInputFile;
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.shell.standard.ShellComponent;
import org.springframework.shell.standard.ShellMethod;
import org.springframework.shell.standard.ShellOption;

/**
 * Hudi相关操作
 *
 * @author lanyuanxiaoyao
 * @date 2024-03-19
 */
@SuppressWarnings("SpringJavaInjectionPointsAutowiringInspection")
@ShellComponent("Hudi相关操作")
public class HudiCommand {
    private static final Logger logger = LoggerFactory.getLogger(HudiCommand.class);

    private final InfoService infoService;
    private final HudiService hudiService;

    public HudiCommand(InfoService infoService, HudiService hudiService) {
        this.infoService = infoService;
        this.hudiService = hudiService;
    }

    public static void time(String name, Runnable runnable) {
        logger.info(name);
        long startTime = Instant.now().toEpochMilli();
        LongAdder counter = new LongAdder();
        runnable.run(counter);
        logger.info("Count: {}", counter.sum());
        logger.info("Cost: {}s", (Instant.now().toEpochMilli() - startTime) / 1000.0);
    }

    public static void reader1(LongAdder counter, Configuration configuration, Path root) {
        try (HoodieParquetReader<IndexedRecord> reader = new HoodieParquetReader<>(configuration, root)) {
            try (ClosableIterator<IndexedRecord> recordIterator = reader.getRecordIterator()) {
                while (recordIterator.hasNext()) {
                    recordIterator.next();
                    counter.increment();
                }
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }

    public static void reader2(LongAdder counter, Configuration configuration, Path root) {
        GroupReadSupport readSupport = new GroupReadSupport();
        try (ParquetReader<Group> reader = AvroParquetReader.builder(readSupport, root).withConf(configuration).build()) {
            while (reader.read() != null) {
                counter.increment();
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public static void reader3(LongAdder counter, Configuration configuration, Path root) {
        try (ParquetReader<GenericRecord> reader = AvroParquetReader.genericRecordReader(HadoopInputFile.fromPath(root, configuration))) {
            while (reader.read() != null) {
                counter.increment();
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @ShellMethod("Test filesystem view")
    public void testFilesystemView() throws IOException {
        HoodieTableMetaClient metaClient = HoodieTableMetaClient.builder()
                .setConf(new Configuration())
                .setBasePath("hdfs://b2/apps/datalake/hive/dws_acct_fs/external_table_hudi/dws_bill_b")
                .setLoadActiveTimelineOnLoad(false)
                .setConsistencyGuardConfig(ConsistencyGuardConfig.newBuilder().build())
                .setLayoutVersion(Option.of(TimelineLayoutVersion.CURR_LAYOUT_VERSION))
                .build();
        FileSystem fs = metaClient.getFs();
        String globPath = StrUtil.format("{}/{}/*", metaClient.getBasePathV2().toString(), "*/*/*");
        List<FileStatus> statuses = FSUtils.getGlobStatusExcludingMetaFolder(fs, new Path(globPath));
        HoodieTimeline timeline = metaClient.getActiveTimeline().getWriteTimeline();
        HoodieTableFileSystemView view = new HoodieTableFileSystemView(metaClient, timeline, statuses.toArray(new FileStatus[0]));
        System.out.println(
                view.fetchAllStoredFileGroups()
                        .map(HoodieFileGroup::getLatestDataFile)
                        .filter(Option::isPresent)
                        .map(Option::get)
                        .map(HoodieBaseFile::getBootstrapBaseFile)
                        .map(Option::get)
                        .map(BaseFile::getPath)
                        .collect(Collectors.toList())
        );
    }

    @ShellMethod("Test glob")
    public void testGlob(@ShellOption(help = "parquet file path") String path) throws IOException {
        FileSystem fileSystem = FileSystem.get(new Configuration());
        MutableList<FileStatus> results = Lists.mutable.empty();
        RemoteIterator<LocatedFileStatus> iterator = fileSystem.listFiles(new Path(path), true);
        while (iterator.hasNext()) {
            LocatedFileStatus status = iterator.next();
            if (fileSystem.exists(status.getPath())) {
                results.add(iterator.next());
            }
        }

        List<FileStatus> statuses = FSUtils.getGlobStatusExcludingMetaFolder(FileSystem.get(new Configuration()), new Path(path, "*/*/*/*"));
        logger.info("Sum1: {} Sum2: {}", results.size(), statuses.size());
    }

    @ShellMethod("Test read parquet file by column")
    public void testColumnReadParquet(@ShellOption(help = "parquet file path") String path) throws IOException {
        Configuration configuration = new Configuration();
        Path root = new Path(path);

        time("reader 1", counter -> reader1(counter, configuration, root));
        time("reader 2", counter -> reader2(counter, configuration, root));

        ParquetUtils baseFileUtils = (ParquetUtils) BaseFileUtils.getInstance(HoodieFileFormat.PARQUET);
        Schema schema = baseFileUtils.readAvroSchema(configuration, root);
        MutableList<Schema.Field> fields = Lists.mutable.ofAll(schema.getFields())
                .select(field -> field.name().equals("CODE"))
                .collect(field -> new Schema.Field(field.name(), field.schema(), field.doc(), field.defaultVal()))
                .toList();
        Schema readSchema = Schema.createRecord(schema.getName(), null, null, false, fields);
        AvroReadSupport.setRequestedProjection(configuration, readSchema);
        time("reader 3", counter -> reader3(counter, configuration, root));
    }

    @ShellMethod("Count meta files")
    public void countMetaFiles(@ShellOption(help = "alias", defaultValue = "") String alias) throws IOException {
        LongAdder count = new LongAdder();
        FileSystem fileSystem = FileSystem.get(new Configuration());
        infoService
                .tableMetaList()
                .select(meta -> StrUtil.isBlank(alias) || StrUtil.equals(meta.getAlias(), alias))
                .collect(TableMeta::getHudi)
                .collect(TableMeta.HudiMeta::getTargetHdfsPath)
                .asParallel(ExecutorProvider.EXECUTORS_5, 1)
                .forEach(hdfs -> {
                    Path root = new Path(hdfs, ".hoodie");
                    try {
                        FileStatus[] statuses = fileSystem.listStatus(root);
                        for (FileStatus status : statuses) {
                            if (status.isFile()) {
                                count.increment();
                            }
                            if (StrUtil.containsIgnoreCase(status.getPath().toString(), "INVALID")) {
                                logger.info("{}", status.getPath().toString());
                            }
                        }
                    } catch (IOException e) {
                        logger.warn("List file error", e);
                    }
                });
        logger.info("Count: {}", count.longValue());
        fileSystem.close();
    }

    @ShellMethod("Max meta files")
    public void maxMetaFiles() throws IOException {
        MutableList<P> list = Lists.mutable.<P>empty().asSynchronized();
        FileSystem fileSystem = FileSystem.get(new Configuration());
        infoService
                .tableMetaList()
                .collect(TableMeta::getHudi)
                .collect(TableMeta.HudiMeta::getTargetHdfsPath)
                .asParallel(ExecutorProvider.EXECUTORS_20, 1)
                .forEach(hdfs -> {
                    Path root = new Path(hdfs, ".hoodie");
                    try {
                        FileStatus[] statuses = fileSystem.listStatus(root);
                        long num = 0;
                        for (FileStatus status : statuses) {
                            if (status.isFile()) {
                                num++;
                            }
                            if (StrUtil.containsIgnoreCase(status.getPath().toString(), "INVALID")) {
                                logger.info("{}", status.getPath().toString());
                            }
                        }
                        list.add(new P(num, hdfs));
                        logger.info("Count: {} Hdfs: {}", num, hdfs);
                    } catch (IOException e) {
                        logger.warn("List file error", e);
                    }
                });
        MutableList<P> listP = list.select(p -> p.count > 1000);
        for (P maxP : listP) {
            logger.info("Max: {} Hdfs: {}", maxP.count, maxP.hdfs);
        }
        fileSystem.close();
    }

    @ShellMethod("Get timeline instants")
    public void timelineInstant(@ShellOption(help = "root hdfs path") String hdfs) {
        hudiService.timelineHdfsAllActive(hdfs).forEach(instant -> logger.info(instant.toString()));
    }

    @ShellMethod("Check compaction")
    public void checkCompaction(@ShellOption(help = "root hdfs path") String hdfs) {
        logger.info("{}", hudiService.timelineHdfsAllActive(hdfs));
        ImmutableList<HudiInstant> finishCompactionInstants = hudiService.timelineHdfsAllActive(hdfs)
                .select(instant -> StrUtil.equals(instant.getAction(), HoodieTimeline.COMPACTION_ACTION))
                .select(instant -> StrUtil.equals(instant.getState(), HoodieInstant.State.COMPLETED.name()));
        for (HudiInstant instant : finishCompactionInstants) {
            logger.info("Compaction: {}", instant);
            HudiCompactionPlan compactionPlan = hudiService.readCompactionPlanHdfs(hdfs, instant.getTimestamp());
            for (HudiCompactionPlan.Operation operation : compactionPlan.getOperations()) {
                logger.info("Base: {}", operation.getDataFilePath());
                operation.getDeltaFilePaths().forEach(path -> logger.info("  Delta: {}", path));
            }
        }
    }

    public interface Runnable {
        void run(LongAdder counter);
    }

    private static final class P {
        String hdfs;
        long count;

        public P(long count, String hdfs) {
            this.count = count;
            this.hdfs = hdfs;
        }
    }
}
