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

import cn.hutool.core.util.StrUtil;
import com.lanyuanxiaoyao.service.executor.core.TaskContext;
import com.lanyuanxiaoyao.service.executor.task.entity.RecordView;
import java.io.Serializable;
import java.time.Instant;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Locale;
import java.util.concurrent.TimeUnit;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.connector.source.*;
import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
import org.apache.flink.core.io.SimpleVersionedSerializer;
import org.apache.pulsar.client.api.*;
import org.eclipse.collections.api.factory.Lists;
import org.eclipse.collections.api.list.MutableList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author lanyuanxiaoyao
 * @date 2024-01-18
 */
public class ReadPulsarSource implements Source<RecordView, ReadPulsarSplit, Collection<ReadPulsarSplit>>, ResultTypeQueryable<RecordView>, Serializable {
    private static final Logger logger = LoggerFactory.getLogger(ReadPulsarSource.class);
    private static final Long TASK_GAP = TimeUnit.MINUTES.toMillis(60);
    private static final DateTimeFormatter FORMATTER = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS")
            .withLocale(Locale.CHINA)
            .withZone(ZoneId.systemDefault());
    private final Collection<ReadPulsarSplit> splits;

    public ReadPulsarSource(TaskContext taskContext, String pulsarUrl, String pulsarTopic) throws PulsarClientException {
        try (PulsarClient client = PulsarClient.builder()
                .serviceUrl(pulsarUrl)
                .build()) {
            try (Consumer<byte[]> consumer = client.newConsumer()
                    .topic(pulsarTopic)
                    .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
                    .subscriptionMode(SubscriptionMode.NonDurable)
                    .subscriptionType(SubscriptionType.Exclusive)
                    .subscriptionName(StrUtil.format("Task_Reader_Detect_{}", taskContext.getTaskId()))
                    .startMessageIdInclusive()
                    .subscribe()) {
                MessageId latestMessageId = consumer.getLastMessageId();
                Message<byte[]> message = consumer.receive();
                long startTimestamp = message.getPublishTime();
                long endTimestamp = Instant.now().toEpochMilli();
                MutableList<ReadPulsarSplit> tasks = Lists.mutable.empty();
                while (startTimestamp < endTimestamp) {
                    tasks.add(new ReadPulsarSplit(
                            taskContext.getTaskId(),
                            pulsarUrl,
                            pulsarTopic,
                            latestMessageId.toString(),
                            startTimestamp,
                            Math.min(endTimestamp, startTimestamp + TASK_GAP)
                    ));
                    startTimestamp += TASK_GAP;
                }
                logger.info("Gap: {}, Splits: {}", TASK_GAP, tasks.size());
                for (ReadPulsarSplit split : tasks) {
                    logger.info("Read split: {} -> {}", covertTimestamp(split.getStartTimestamp()), covertTimestamp(split.getEndTimestamp()));
                }
                splits = new ArrayList<>(tasks.shuffleThis());
            }
        }
    }

    private static String covertTimestamp(Long timestamp) {
        return FORMATTER.format(Instant.ofEpochMilli(timestamp));
    }

    @Override
    public Boundedness getBoundedness() {
        return Boundedness.BOUNDED;
    }

    @Override
    public SourceReader<RecordView, ReadPulsarSplit> createReader(SourceReaderContext readerContext) throws PulsarClientException {
        return new ReadPulsarSourceReader(readerContext);
    }

    @Override
    public SplitEnumerator<ReadPulsarSplit, Collection<ReadPulsarSplit>> createEnumerator(SplitEnumeratorContext<ReadPulsarSplit> enumContext) throws Exception {
        return new ReadPulsarSourceEnumerator(enumContext, splits);
    }

    @Override
    public SplitEnumerator<ReadPulsarSplit, Collection<ReadPulsarSplit>> restoreEnumerator(SplitEnumeratorContext<ReadPulsarSplit> enumContext, Collection<ReadPulsarSplit> checkpoint) throws Exception {
        return new ReadPulsarSourceEnumerator(enumContext, checkpoint);
    }

    @Override
    public SimpleVersionedSerializer<ReadPulsarSplit> getSplitSerializer() {
        return new ReadPulsarVersionedSplitSerializer();
    }

    @Override
    public SimpleVersionedSerializer<Collection<ReadPulsarSplit>> getEnumeratorCheckpointSerializer() {
        return new ReadPulsarVersionedCheckpointSerializer();
    }

    @Override
    public TypeInformation<RecordView> getProducedType() {
        return TypeInformation.of(RecordView.class);
    }
}
