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

import cn.hutool.core.util.ObjectUtil;
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 javax.annotation.Nullable;
import org.apache.flink.api.connector.source.SplitEnumerator;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;
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;

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

    @Override
    public void start() {
    }

    @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);
        } else {
            logger.info("t{} No more split for {}", subtaskId, subtaskId);
            context.signalNoMoreSplits(subtaskId);
        }
    }

    @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 {
    }
}
