package com.lanyuanxiaoyao.service.hudi.utils;

import cn.hutool.core.util.ReUtil;
import cn.hutool.core.util.StrUtil;
import com.lanyuanxiaoyao.service.configuration.entity.hudi.HudiInstant;
import java.io.IOException;
import java.util.Arrays;
import java.util.Comparator;
import java.util.List;
import java.util.Set;
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.Path;
import org.apache.hudi.common.table.HoodieTableMetaClient;
import org.apache.hudi.common.table.TableSchemaResolver;
import org.apache.hudi.common.table.timeline.HoodieActiveTimeline;
import org.apache.hudi.common.table.timeline.HoodieArchivedTimeline;
import org.apache.hudi.common.table.timeline.HoodieInstant;
import org.apache.hudi.common.table.timeline.HoodieTimeline;
import org.apache.hudi.org.apache.avro.Schema;
import org.eclipse.collections.api.factory.Lists;
import org.eclipse.collections.api.list.ImmutableList;
import org.eclipse.collections.api.list.MutableList;
import org.eclipse.collections.api.map.MutableMap;
import org.eclipse.collections.api.multimap.list.MutableListMultimap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Hudi 工具
 *
 * @author lanyuanxiaoyao
 * @date 2023-05-01
 */
public class HoodieUtils {
    private static final Logger logger = LoggerFactory.getLogger(HoodieUtils.class);

    public static Schema getSchema(String hdfs) throws Exception {
        HoodieTableMetaClient client = HoodieTableMetaClient.builder()
                .setConf(new Configuration())
                .setBasePath(hdfs)
                .build();
        return getSchema(client);
    }

    public static Schema getSchema(HoodieTableMetaClient client) throws Exception {
        TableSchemaResolver schemaUtil = new TableSchemaResolver(client);
        return schemaUtil.getTableAvroSchema(true);
    }

    public static ImmutableList<HudiInstant> getAllActiveInstants(HoodieTableMetaClient client) throws IOException {
        return getAllInstants(client, true, false);
    }

    public static ImmutableList<HudiInstant> getAllArchiveInstants(HoodieTableMetaClient client) throws IOException {
        return getAllInstants(client, false, true);
    }

    public static ImmutableList<HudiInstant> getAllInstants(HoodieTableMetaClient client) throws IOException {
        return getAllInstants(client, true, true);
    }

    public static ImmutableList<HudiInstant> getAllInstants(HoodieTableMetaClient client, Boolean active, Boolean archive) throws IOException {
        FileSystem fileSystem = client.getRawFs();
        Path metadataPath = new Path(client.getBasePathV2().toString(), HoodieTableMetaClient.METAFOLDER_NAME);
        MutableList<HudiInstant> instants = Lists.mutable.empty();

        if (active) {
            instants.addAllIterable(activeInstants(client));
        }
        if (archive) {
            instants.addAllIterable(archiveInstants(client));
        }

        MutableMap<String, Long> fileModifiedTimeMap = Lists.immutable.of(fileSystem.listStatus(metadataPath))
                .toMap(status -> status.getPath().toString(), FileStatus::getModificationTime);
        return instants
                .collect(instant -> {
                    String instantPath = StrUtil.format("{}/{}", client.getMetaPath(), instant.getFileName());
                    instant.setFileTime(fileModifiedTimeMap.getOrDefault(instantPath, 0L));
                    return instant;
                })
                .sortThis(Comparator
                        .<HudiInstant>comparingLong(instant -> Long.parseLong(instant.getTimestamp()))
                        .thenComparingLong(HudiInstant::getFileTime))
                .toImmutable();
    }

    public static HudiInstant convert(HoodieInstant instant, String type) {
        return new HudiInstant(
                instant.getAction(),
                instant.getState().name(),
                instant.getTimestamp(),
                instant.getFileName(),
                0L,
                type
        );
    }

    public static ImmutableList<HudiInstant> activeInstants(HoodieTableMetaClient client) throws IOException {
        HoodieActiveTimeline timeline = client.getActiveTimeline();
        Set<String> committedTimestamps = timeline.getCommitsTimeline()
                .filterCompletedInstants()
                .getInstants()
                .map(HoodieInstant::getTimestamp)
                .collect(Collectors.toSet());
        Set<String> compactionCommittedTimestamps = Arrays.stream(client.getFs().listStatus(new Path(client.getMetaPath())))
                .filter(status -> status.getPath().toString().endsWith(HoodieTimeline.REQUESTED_COMPACTION_EXTENSION))
                .map(status -> status.getPath().getName())
                .map(name -> ReUtil.get("^(\\d+)\\..+", name, 1))
                .filter(committedTimestamps::contains)
                .collect(Collectors.toSet());
        List<HudiInstant> instants = timeline.getInstants()
                .map(instant -> convert(instant, "active"))
                .map(instant -> {
                    if (compactionCommittedTimestamps.contains(instant.getTimestamp())) {
                        return new HudiInstant(
                                HoodieTimeline.COMPACTION_ACTION,
                                HoodieInstant.State.COMPLETED.name(),
                                instant.getTimestamp(),
                                instant.getFileName(),
                                instant.getFileTime(),
                                instant.getType()
                        );
                    }
                    return instant;
                })
                .collect(Collectors.toList());
        return Lists.immutable.ofAll(instants);
    }

    public static ImmutableList<HudiInstant> archiveInstants(HoodieTableMetaClient client) {
        HoodieArchivedTimeline timeline = client.getArchivedTimeline();
        MutableList<HoodieInstant> instants = Lists.mutable.ofAll(timeline.getInstants().collect(Collectors.toList()));
        instants.forEach(instant -> logger.info(instant.toString()));
        MutableListMultimap<HoodieInstant.State, String> stateMap = instants.groupBy(HoodieInstant::getState).collectValues(HoodieInstant::getTimestamp);
        return instants
                .select(instant -> HoodieInstant.State.REQUESTED.equals(instant.getState()))
                .collect(instant -> convert(instant, "archive"))
                .collect(instant -> {
                    if (stateMap.containsKeyAndValue(HoodieInstant.State.INVALID, instant.getTimestamp())) {
                        instant.setState(HoodieInstant.State.INVALID.name());
                    } else if (stateMap.containsKeyAndValue(HoodieInstant.State.COMPLETED, instant.getTimestamp())) {
                        instant.setState(HoodieInstant.State.COMPLETED.name());
                    } else if (stateMap.containsKeyAndValue(HoodieInstant.State.INFLIGHT, instant.getTimestamp())) {
                        instant.setState(HoodieInstant.State.INFLIGHT.name());
                    }
                    return instant;
                })
                .toImmutable();
    }
}

