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

import cn.hutool.core.collection.ListUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
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.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.*;
import org.apache.pulsar.client.impl.schema.StringSchema;
import org.apache.pulsar.client.internal.DefaultImplementation;
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) throws PulsarClientException {
        this.readerContext = readerContext;
    }

    private static MessageId parseMessageId(String messageIdText) {
        String[] items = messageIdText.split(":");
        return DefaultImplementation.newMessageId(Long.parseLong(items[0]), Long.parseLong(items[1]), -1);
    }

    @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());
        if (ObjectUtil.isNotNull(currentSplit)) {
            logger.info("t{} Read split: {}", readerContext.getIndexOfSubtask(), currentSplit.getStartTimestamp());
            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(),
                                currentSplit.getStartTimestamp(),
                                currentSplit.getEndTimestamp()
                        ))
                        .startMessageIdInclusive()
                        .startMessageId(MessageId.earliest)
                        .create()) {
                    reader.seek(currentSplit.getStartTimestamp());
                    Message<String> message = reader.readNext(10, TimeUnit.SECONDS);
                    while (ObjectUtil.isNotNull(message)) {
                        if (message.getPublishTime() > currentSplit.getEndTimestamp()) {
                            logger.info("t{} Break for {} -> {}, Queue rest: {}", readerContext.getIndexOfSubtask(), message.getPublishTime(), currentSplit.getEndTimestamp(), readQueue.size());
                            break;
                        }
                        output.collect(parsePulsarMessage(message));
                        message = reader.readNext(10, TimeUnit.SECONDS);
                    }
                }
            }
        }

        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{} Add splits: {}", readerContext.getIndexOfSubtask(), splits.stream().map(ReadPulsarSplit::getStartTimestamp).collect(Collectors.toList()));
        readQueue.addAll(splits);
        availability.complete(null);
    }

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

    @Override
    public void close() {
    }
}
