package com.lanyuanxiaoyao.service.executor.task.functions;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.lanyuanxiaoyao.service.common.Constants;
import com.lanyuanxiaoyao.service.executor.task.entity.RecordView;
import java.io.IOException;
import java.util.Arrays;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.util.Collector;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hudi.common.fs.FSUtils;
import org.apache.hudi.common.model.HoodieLogFile;
import org.apache.hudi.common.model.HoodieRecord;
import org.apache.hudi.common.table.TableSchemaResolver;
import org.apache.hudi.common.table.log.HoodieLogFormat;
import org.apache.hudi.common.table.log.block.*;
import org.apache.hudi.common.util.ClosableIterator;
import org.apache.hudi.io.storage.HoodieParquetReader;
import org.apache.hudi.org.apache.avro.Schema;
import org.apache.hudi.org.apache.avro.generic.IndexedRecord;
import org.apache.hudi.org.apache.avro.util.Utf8;
import org.apache.parquet.avro.AvroSchemaConverter;
import org.apache.parquet.schema.MessageType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 读取log/data文件
 *
 * @author lanyuanxiaoyao
 * @date 2024-01-09
 */
public class ReadHudiFile implements FlatMapFunction<String, RecordView> {
    private static final Logger logger = LoggerFactory.getLogger(ReadHudiFile.class);

    private RecordView parseData(String source, RecordView.Operation operation, IndexedRecord record) {
        Schema schema = record.getSchema();
        StringBuilder builder = new StringBuilder();
        for (Schema.Field field : schema.getFields()) {
            builder.append(field.name())
                    .append("=")
                    .append(record.get(field.pos()))
                    .append(" ");
        }
        String timestamp = null;
        Schema.Field commitTimeField = schema.getField(HoodieRecord.COMMIT_TIME_METADATA_FIELD);
        if (ObjectUtil.isNotNull(commitTimeField)) {
            timestamp = ((Utf8) record.get(commitTimeField.pos())).toString();
        }
        String latestOpTs = null;
        Schema.Field latestOpTsField = schema.getField(Constants.LATEST_OPERATION_TIMESTAMP_KEY_NAME);
        if (ObjectUtil.isNotNull(latestOpTsField)) {
            latestOpTs = ((Utf8) record.get(latestOpTsField.pos())).toString();
        }

        String data = builder.toString();
        RecordView recordView = new RecordView(operation, data, timestamp, source);
        if (StrUtil.isNotBlank(latestOpTs)) {
            recordView.getAttributes().put(Constants.LATEST_OPERATION_TIMESTAMP_KEY_NAME, latestOpTs);
        }

        return recordView;
    }

    @Override
    public void flatMap(String value, Collector<RecordView> out) throws IOException {
        Configuration readerConfiguration = new Configuration();
        FileSystem readerFilesystem = FileSystem.get(readerConfiguration);
        Path filePath = new Path(value);
        if (FSUtils.isLogFile(filePath)) {
            readLogFile(readerFilesystem, filePath, out);
        } else if (FSUtils.isDataFile(filePath)) {
            readDataFile(readerFilesystem, filePath, out);
        } else {
            logger.warn("Cannot read file format: {}", filePath);
        }
    }

    private void readDataFile(FileSystem readerFilesystem, Path dataFilePath, Collector<RecordView> out) throws IOException {
        HoodieParquetReader<IndexedRecord> reader = new HoodieParquetReader<>(readerFilesystem.getConf(), dataFilePath);
        ClosableIterator<IndexedRecord> recordIterator = reader.getRecordIterator();
        while (recordIterator.hasNext()) {
            RecordView recordView = parseData(dataFilePath.toString(), RecordView.Operation.RESULT, recordIterator.next());
            out.collect(recordView);
        }
        recordIterator.close();
        reader.close();
    }

    private void readLogFile(FileSystem readerFilesystem, Path logFilePath, Collector<RecordView> out) throws IOException {
        try (HoodieLogFormat.Reader reader = HoodieLogFormat.newReader(readerFilesystem, new HoodieLogFile(logFilePath), null)) {
            while (reader.hasNext()) {
                HoodieLogBlock block = reader.next();
                Map<HoodieLogBlock.HeaderMetadataType, String> logBlockHeader = block.getLogBlockHeader();
                String instant = logBlockHeader.getOrDefault(HoodieLogBlock.HeaderMetadataType.INSTANT_TIME, null);
                switch (block.getBlockType()) {
                    case AVRO_DATA_BLOCK:
                        HoodieAvroDataBlock avroDataBlock = (HoodieAvroDataBlock) block;
                        try (ClosableIterator<IndexedRecord> avroDataBlockRecordIterator = avroDataBlock.getRecordIterator()) {
                            while (avroDataBlockRecordIterator.hasNext()) {
                                RecordView recordView = parseData(logFilePath.toString(), RecordView.Operation.UPSERT, avroDataBlockRecordIterator.next());
                                out.collect(recordView);
                            }
                        }
                        break;
                    case PARQUET_DATA_BLOCK:
                        HoodieParquetDataBlock parquetDataBlock = (HoodieParquetDataBlock) block;
                        try (ClosableIterator<IndexedRecord> parquetDataBlockRecordIterator = parquetDataBlock.getRecordIterator()) {
                            while (parquetDataBlockRecordIterator.hasNext()) {
                                RecordView recordView = parseData(logFilePath.toString(), RecordView.Operation.UPSERT, parquetDataBlockRecordIterator.next());
                                out.collect(recordView);
                            }
                        }
                        break;
                    case CORRUPT_BLOCK:
                        break;
                    case DELETE_BLOCK:
                        HoodieDeleteBlock deleteBlock = (HoodieDeleteBlock) block;
                        String keys = Arrays.stream(deleteBlock.getRecordsToDelete())
                                .map(deleteRecord -> deleteRecord.getHoodieKey().toString())
                                .collect(Collectors.joining(" "));
                        out.collect(new RecordView(RecordView.Operation.DELETE, keys, instant, logFilePath.toString()));
                        break;
                    case COMMAND_BLOCK:
                        HoodieCommandBlock commandBlock = (HoodieCommandBlock) block;
                        Map<HoodieLogBlock.HeaderMetadataType, String> header = commandBlock.getLogBlockHeader();
                        out.collect(new RecordView(RecordView.Operation.ROLLBACK, header.get(HoodieLogBlock.HeaderMetadataType.TARGET_INSTANT_TIME), instant, logFilePath.toString()));
                        break;
                    default:
                        break;
                }
            }
        }
    }
}
