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

import cn.hutool.core.collection.ListUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.lanyuanxiaoyao.service.executor.task.entity.RecordView;
import com.lanyuanxiaoyao.service.executor.task.functions.pulsar.event.AddEvent;
import com.lanyuanxiaoyao.service.executor.task.functions.pulsar.event.EndEvent;
import com.lanyuanxiaoyao.service.executor.task.functions.pulsar.event.StartEvent;
import com.lanyuanxiaoyao.service.executor.task.helper.TimeRangeHelper;
import java.io.Serializable;
import java.time.Instant;
import java.time.ZoneId;
import java.time.format.DateTimeFormatter;
import java.util.ArrayDeque;
import java.util.List;
import java.util.Locale;
import java.util.Queue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import org.apache.flink.api.connector.source.ReaderOutput;
import org.apache.flink.api.connector.source.SourceReader;
import org.apache.flink.api.connector.source.SourceReaderContext;
import org.apache.flink.core.io.InputStatus;
import org.apache.pulsar.client.api.Message;
import org.apache.pulsar.client.api.MessageId;
import org.apache.pulsar.client.api.PulsarClient;
import org.apache.pulsar.client.api.Reader;
import org.apache.pulsar.client.impl.schema.StringSchema;
import org.eclipse.collections.api.list.ImmutableList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author lanyuanxiaoyao
 * @date 2024-01-18
 */
public class ReadPulsarSourceReader implements SourceReader<RecordView, ReadPulsarSplit>, Serializable {
    private static final Logger logger = LoggerFactory.getLogger(ReadPulsarSourceReader.class);
    private static final DateTimeFormatter FORMATTER = DateTimeFormatter.ofPattern("yyyyMMddHHmmssSSS")
            .withLocale(Locale.CHINA)
            .withZone(ZoneId.systemDefault());
    private final Queue<ReadPulsarSplit> readQueue = new ArrayDeque<>();
    private final SourceReaderContext readerContext;
    private CompletableFuture<Void> availability = new CompletableFuture<>();
    private ReadPulsarSplit currentSplit;
    private boolean noMoreSplits = false;

    public ReadPulsarSourceReader(SourceReaderContext readerContext) {
        this.readerContext = readerContext;
    }

    @Override
    public void start() {
        if (readQueue.isEmpty()) {
            readerContext.sendSplitRequest();
        }
    }

    private RecordView parsePulsarMessage(Message<String> message) {
        return new RecordView(
                RecordView.Operation.QUEUE,
                message.getValue(),
                FORMATTER.format(Instant.ofEpochMilli(message.getPublishTime())),
                message.getMessageId().toString()
        );
    }

    @Override
    public InputStatus pollNext(ReaderOutput<RecordView> output) throws Exception {
        logger.info("t{} Poll Next", readerContext.getIndexOfSubtask());
        readerContext.sendSourceEventToCoordinator(new StartEvent());
        if (ObjectUtil.isNotNull(currentSplit)) {
            logger.info("t{} Read split: {}", readerContext.getIndexOfSubtask(), currentSplit.getSplitId());
            long startTimestamp = currentSplit.getStartTimestamp();
            long endTimestamp = currentSplit.getEndTimestamp();
            try (PulsarClient client = PulsarClient.builder()
                    .serviceUrl(currentSplit.getPulsarUrl())
                    .build()) {
                try (Reader<String> reader = client.newReader(new StringSchema())
                        .topic(currentSplit.getPulsarTopic())
                        .receiverQueueSize(50000)
                        .subscriptionName(StrUtil.format(
                                "Task_Reader_{}_{}_{}_{}",
                                currentSplit.getTaskId(),
                                readerContext.getIndexOfSubtask(),
                                startTimestamp,
                                endTimestamp
                        ))
                        .startMessageIdInclusive()
                        .startMessageId(MessageId.earliest)
                        .create()) {
                    reader.seek(startTimestamp);
                    long count = 0;
                    Message<String> message = reader.readNext(10, TimeUnit.SECONDS);
                    while (ObjectUtil.isNotNull(message)) {
                        long publishTime = message.getPublishTime();
                        if (publishTime > endTimestamp) {
                            logger.info("t{} Break for {} -> {}, Queue rest: {}", readerContext.getIndexOfSubtask(), message.getPublishTime(), currentSplit.getEndTimestamp(), readQueue.size());
                            break;
                        }
                        if (++count > 500000) {
                            ImmutableList<TimeRangeHelper.TimeRange> range = TimeRangeHelper.range(publishTime, endTimestamp, currentSplit.getGap() / 5);
                            if (ObjectUtil.isNotEmpty(range)) {
                                readerContext.sendSourceEventToCoordinator(new AddEvent(
                                        range.collect(r -> new ReadPulsarSplit(
                                                currentSplit.getTaskId(),
                                                IdUtil.nanoId(10),
                                                currentSplit.getPulsarUrl(),
                                                currentSplit.getPulsarTopic(),
                                                currentSplit.getLatestMessageId(),
                                                r.getStart(),
                                                r.getEnd(),
                                                r.getGap()
                                        )).toList()
                                ));
                                break;
                            } else {
                                count = 0;
                            }
                        }
                        output.collect(parsePulsarMessage(message));
                        message = reader.readNext(10, TimeUnit.SECONDS);
                    }
                }
            }
            readerContext.sendSourceEventToCoordinator(new EndEvent(currentSplit));
        }

        currentSplit = null;
        if (ObjectUtil.isEmpty(readQueue) && !noMoreSplits) {
            readerContext.sendSplitRequest();
            logger.info("t{} Request new split", readerContext.getIndexOfSubtask());
        }

        currentSplit = readQueue.poll();
        logger.info("t{} Queue rest: {}, Current split: {}", readerContext.getIndexOfSubtask(), readQueue.size(), currentSplit);
        if (ObjectUtil.isNotNull(currentSplit)) {
            return InputStatus.MORE_AVAILABLE;
        } else if (noMoreSplits) {
            return InputStatus.END_OF_INPUT;
        } else {
            if (availability.isDone()) {
                availability = new CompletableFuture<>();
            }
            return InputStatus.NOTHING_AVAILABLE;
        }
    }

    @Override
    public List<ReadPulsarSplit> snapshotState(long checkpointId) {
        return ListUtil.empty();
    }

    @Override
    public CompletableFuture<Void> isAvailable() {
        return availability;
    }

    @Override
    public void addSplits(List<ReadPulsarSplit> splits) {
        logger.info("t{} Receive add splits: {}", readerContext.getIndexOfSubtask(), splits.stream().map(ReadPulsarSplit::getSplitId).collect(Collectors.toList()));
        readQueue.addAll(splits);
        availability.complete(null);
    }

    @Override
    public void notifyNoMoreSplits() {
        logger.info("t{} Receive no more splits", readerContext.getIndexOfSubtask());
        noMoreSplits = true;
        availability.complete(null);
    }

    @Override
    public void close() {
    }
}
