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

import cn.hutool.core.util.ObjectUtil;
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 java.io.IOException;
import java.io.Serializable;
import java.util.ArrayDeque;
import java.util.Collection;
import java.util.List;
import java.util.Queue;
import java.util.Set;
import javax.annotation.Nullable;
import org.apache.flink.api.connector.source.SourceEvent;
import org.apache.flink.api.connector.source.SplitEnumerator;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;
import org.eclipse.collections.api.factory.Lists;
import org.eclipse.collections.api.factory.Sets;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author lanyuanxiaoyao
 * @date 2024-01-18
 */
public class ReadPulsarSourceEnumerator implements SplitEnumerator<ReadPulsarSplit, Collection<ReadPulsarSplit>>, Serializable {
    private static final Logger logger = LoggerFactory.getLogger(ReadPulsarSourceEnumerator.class);
    private final SplitEnumeratorContext<ReadPulsarSplit> context;
    private final Queue<ReadPulsarSplit> readQueue;
    private final Set<String> runningTasks = Sets.mutable.<String>empty().asSynchronized();

    public ReadPulsarSourceEnumerator(SplitEnumeratorContext<ReadPulsarSplit> context, Collection<ReadPulsarSplit> splits) {
        this.context = context;
        this.readQueue = new ArrayDeque<>(splits);
    }

    @Override
    public void start() {
    }

    @Override
    public void handleSourceEvent(int subtaskId, SourceEvent sourceEvent) {
        if (sourceEvent instanceof AddEvent) {
            AddEvent event = (AddEvent) sourceEvent;
            logger.info("t{} Add: {}", subtaskId, event.getSplits());
            readQueue.addAll(event.getSplits());

            logger.info("t{} ReSend task, running tasks: {}", subtaskId, runningTasks);
            context.registeredReaders()
                    .values()
                    .forEach(info -> handleSourceEvent(info.getSubtaskId(), null));
        } else if (sourceEvent instanceof EndEvent) {
            EndEvent event = (EndEvent) sourceEvent;
            logger.info("t{} End: {}", subtaskId, event.getSplit());
            if (ObjectUtil.isNotNull(event.getSplit())) {
                runningTasks.remove(event.getSplit().getSplitId());
            }
            logger.info("t{} Queue: {} running tasks: {}", subtaskId, readQueue.size(), runningTasks);
            if (ObjectUtil.isEmpty(readQueue) && ObjectUtil.isEmpty(runningTasks)) {
                logger.info("t{} No more", subtaskId);
                context.registeredReaders()
                        .values()
                        .forEach(info -> context.signalNoMoreSplits(info.getSubtaskId()));
            }
        }
    }

    @Override
    public void handleSplitRequest(int subtaskId, @Nullable String requesterHostname) {
        final ReadPulsarSplit split = readQueue.poll();
        if (ObjectUtil.isNotNull(split)) {
            logger.info("t{} Assign split for {}, Queue rest: {}", subtaskId, subtaskId, readQueue.size());
            context.assignSplit(split, subtaskId);
            runningTasks.add(split.getSplitId());
        }
    }

    @Override
    public void addSplitsBack(List<ReadPulsarSplit> splits, int subtaskId) {
        readQueue.addAll(splits);
    }

    @Override
    public void addReader(int subtaskId) {
    }

    @Override
    public Collection<ReadPulsarSplit> snapshotState(long checkpointId) throws Exception {
        return readQueue;
    }

    @Override
    public void close() throws IOException {
    }
}
