/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.paimon.flink.source;

import org.apache.paimon.annotation.VisibleForTesting;
import org.apache.paimon.flink.source.assigners.FIFOSplitAssigner;
import org.apache.paimon.flink.source.assigners.PreAssignSplitAssigner;
import org.apache.paimon.flink.source.assigners.SplitAssigner;
import org.apache.paimon.io.DataFileMeta;
import org.apache.paimon.postpone.PostponeBucketFileStoreWrite;
import org.apache.paimon.table.BucketMode;
import org.apache.paimon.table.sink.ChannelComputer;
import org.apache.paimon.table.source.ChainSplit;
import org.apache.paimon.table.source.DataSplit;
import org.apache.paimon.table.source.EndOfScanException;
import org.apache.paimon.table.source.IncrementalSplit;
import org.apache.paimon.table.source.SnapshotNotExistPlan;
import org.apache.paimon.table.source.StreamTableScan;
import org.apache.paimon.table.source.TableScan;

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.apache.flink.api.connector.source.SplitsAssignment;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;

import static org.apache.paimon.utils.Preconditions.checkArgument;
import static org.apache.paimon.utils.Preconditions.checkNotNull;

/** A continuously monitoring enumerator. */
public class ContinuousFileSplitEnumerator
        implements SplitEnumerator<FileStoreSourceSplit, PendingSplitsCheckpoint> {

    private static final Logger LOG = LoggerFactory.getLogger(ContinuousFileSplitEnumerator.class);

    protected final SplitEnumeratorContext<FileStoreSourceSplit> context;

    protected final long discoveryInterval;

    protected final Set<Integer> readersAwaitingSplit;

    protected final FileStoreSourceSplitGenerator splitGenerator;

    protected final StreamTableScan scan;

    protected final int splitMaxPerTask;

    protected final SplitAssigner splitAssigner;

    protected final ConsumerProgressCalculator consumerProgressCalculator;

    private final int splitMaxNum;

    private final boolean shuffleBucketWithPartition;

    @Nullable protected Long nextSnapshotId;

    protected boolean finished = false;

    private boolean stopTriggerScan = false;

    private long handledSnapshotCount = 0;

    private final int maxSnapshotCount;

    private final int sourceParallelismUpperBound;

    /**
     * Metric name for source scaling max parallelism. This metric provides a recommended upper
     * bound of parallelism for auto-scaling systems.
     */
    public static final String SOURCE_PARALLELISM_UPPER_BOUND = "sourceParallelismUpperBound";

    public ContinuousFileSplitEnumerator(
            SplitEnumeratorContext<FileStoreSourceSplit> context,
            Collection<FileStoreSourceSplit> remainSplits,
            @Nullable Long nextSnapshotId,
            long discoveryInterval,
            StreamTableScan scan,
            boolean unordered,
            int splitMaxPerTask,
            boolean shuffleBucketWithPartition,
            int maxSnapshotCount,
            int sourceParallelismUpperBound) {
        checkArgument(discoveryInterval > 0L);
        this.context = checkNotNull(context);
        this.nextSnapshotId = nextSnapshotId;
        this.discoveryInterval = discoveryInterval;
        this.readersAwaitingSplit = new LinkedHashSet<>();
        this.splitGenerator = new FileStoreSourceSplitGenerator();
        this.scan = scan;
        this.splitMaxPerTask = splitMaxPerTask;
        this.splitMaxNum = context.currentParallelism() * splitMaxPerTask;
        this.splitAssigner = createSplitAssigner(unordered);
        this.shuffleBucketWithPartition = shuffleBucketWithPartition;
        addSplits(remainSplits);

        this.consumerProgressCalculator =
                new ConsumerProgressCalculator(context.currentParallelism());
        this.maxSnapshotCount = maxSnapshotCount;
        this.sourceParallelismUpperBound = sourceParallelismUpperBound;
    }

    @VisibleForTesting
    void enableTriggerScan() {
        this.stopTriggerScan = false;
    }

    protected void addSplits(Collection<FileStoreSourceSplit> splits) {
        splits.forEach(this::addSplit);
    }

    private void addSplit(FileStoreSourceSplit split) {
        splitAssigner.addSplit(assignSuggestedTask(split), split);
    }

    @Override
    public void start() {
        registerMetrics();
        context.callAsync(
                this::scanNextSnapshot, this::processDiscoveredSplits, 0, discoveryInterval);
    }

    private void registerMetrics() {
        try {
            context.metricGroup()
                    .gauge(SOURCE_PARALLELISM_UPPER_BOUND, () -> sourceParallelismUpperBound);
        } catch (Exception e) {
            LOG.warn("Failed to register enumerator metrics.", e);
        }
    }

    @Override
    public void close() throws IOException {
        // no resources to close
    }

    @Override
    public void addReader(int subtaskId) {
        // this source is purely lazy-pull-based, nothing to do upon registration
    }

    @Override
    public void handleSplitRequest(int subtaskId, @Nullable String requesterHostname) {
        readersAwaitingSplit.add(subtaskId);
        assignSplits();
        // if current task assigned no split, we check conditions to scan one more time
        if (readersAwaitingSplit.contains(subtaskId)) {
            if (stopTriggerScan) {
                return;
            }
            stopTriggerScan = true;
            context.callAsync(this::scanNextSnapshot, this::processDiscoveredSplits);
        }
    }

    @Override
    public void handleSourceEvent(int subtaskId, SourceEvent sourceEvent) {
        if (sourceEvent instanceof ReaderConsumeProgressEvent) {
            consumerProgressCalculator.updateConsumeProgress(
                    subtaskId, (ReaderConsumeProgressEvent) sourceEvent);
        } else {
            LOG.error("Received unrecognized event: {}", sourceEvent);
        }
    }

    @Override
    public void addSplitsBack(List<FileStoreSourceSplit> splits, int subtaskId) {
        LOG.debug("File Source Enumerator adds splits back: {}", splits);
        splitAssigner.addSplitsBack(subtaskId, splits);
    }

    @Override
    public PendingSplitsCheckpoint snapshotState(long checkpointId) throws Exception {
        List<FileStoreSourceSplit> splits = new ArrayList<>(splitAssigner.remainingSplits());
        final PendingSplitsCheckpoint checkpoint =
                new PendingSplitsCheckpoint(splits, nextSnapshotId);

        consumerProgressCalculator.notifySnapshotState(
                checkpointId,
                readersAwaitingSplit,
                subtask -> splitAssigner.getNextSnapshotId(subtask).orElse(nextSnapshotId),
                context.currentParallelism());

        LOG.debug("Source Checkpoint is {}", checkpoint);
        return checkpoint;
    }

    @Override
    public void notifyCheckpointComplete(long checkpointId) throws Exception {
        consumerProgressCalculator
                .notifyCheckpointComplete(checkpointId)
                .ifPresent(scan::notifyCheckpointComplete);
        handledSnapshotCount = 0;
    }

    // ------------------------------------------------------------------------

    // this need to be synchronized because scan object is not thread safe. handleSplitRequest and
    // context.callAsync will invoke this. This method runs in workerExecutorThreadPool in
    // parallelism.
    protected synchronized Optional<PlanWithNextSnapshotId> scanNextSnapshot() {
        if (splitAssigner.numberOfRemainingSplits() >= splitMaxNum) {
            return Optional.empty();
        }
        if (maxSnapshotCount > 0 && handledSnapshotCount >= maxSnapshotCount) {
            LOG.debug(
                    "There is {} in-flight snapshot, pending to scan next snapshot.",
                    handledSnapshotCount);
            return Optional.empty();
        }

        TableScan.Plan plan = scan.plan();
        Long nextSnapshotId = scan.checkpoint();
        if (nextSnapshotId != null && !plan.splits().isEmpty()) {
            if (this.nextSnapshotId == null) {
                handledSnapshotCount++;
            } else if (!nextSnapshotId.equals(this.nextSnapshotId)) {
                handledSnapshotCount++;
            }
        }
        return Optional.of(new PlanWithNextSnapshotId(plan, nextSnapshotId));
    }

    // this mothod could not be synchronized, because it runs in coordinatorThread, which will make
    // it serialize.
    protected void processDiscoveredSplits(
            Optional<PlanWithNextSnapshotId> planWithNextSnapshotIdOptional, Throwable error) {
        if (error != null) {
            if (error instanceof EndOfScanException) {
                // finished
                LOG.debug("Catching EndOfStreamException, the stream is finished.");
                finished = true;
                assignSplits();
            } else {
                LOG.error("Failed to enumerate files", error);
                throw new RuntimeException(error);
            }
            return;
        }

        if (!planWithNextSnapshotIdOptional.isPresent()) {
            return;
        }
        PlanWithNextSnapshotId planWithNextSnapshotId = planWithNextSnapshotIdOptional.get();
        nextSnapshotId = planWithNextSnapshotId.nextSnapshotId;
        TableScan.Plan plan = planWithNextSnapshotId.plan;
        if (plan.equals(SnapshotNotExistPlan.INSTANCE)) {
            stopTriggerScan = true;
            return;
        }

        stopTriggerScan = false;
        if (plan.splits().isEmpty()) {
            return;
        }

        addSplits(splitGenerator.createSplits(plan));
        assignSplits();
    }

    /**
     * Method should be synchronized because {@link #handleSplitRequest} and {@link
     * #processDiscoveredSplits} have thread conflicts.
     */
    protected synchronized void assignSplits() {
        // create assignment
        Map<Integer, List<FileStoreSourceSplit>> assignment = new HashMap<>();
        Iterator<Integer> readersAwait = readersAwaitingSplit.iterator();
        Set<Integer> subtaskIds = context.registeredReaders().keySet();
        while (readersAwait.hasNext()) {
            Integer task = readersAwait.next();
            if (!subtaskIds.contains(task)) {
                readersAwait.remove();
                continue;
            }
            List<FileStoreSourceSplit> splits = splitAssigner.getNext(task, null);
            if (!splits.isEmpty()) {
                assignment.put(task, splits);
                consumerProgressCalculator.updateAssignInformation(task, splits.get(0));
            }
        }

        // remove readers who fetched splits
        if (noMoreSplits()) {
            Iterator<Integer> iterator = readersAwaitingSplit.iterator();
            while (iterator.hasNext()) {
                Integer reader = iterator.next();
                if (!assignment.containsKey(reader)) {
                    context.signalNoMoreSplits(reader);
                    iterator.remove();
                }
            }
        }
        assignment.keySet().forEach(readersAwaitingSplit::remove);
        context.assignSplits(new SplitsAssignment<>(assignment));
    }

    protected int assignSuggestedTask(FileStoreSourceSplit split) {
        int task;
        if (split.split() instanceof DataSplit) {
            task = assignSuggestedTask((DataSplit) split.split());
        } else if (split.split() instanceof ChainSplit) {
            task = assignSuggestedTask((ChainSplit) split.split());
        } else {
            task = assignSuggestedTask((IncrementalSplit) split.split());
        }

        // Split assigners keep splits in a map keyed by task, but only ever hand out splits for
        // tasks within the parallelism, so a task outside of it silently strands the split instead
        // of failing.
        int parallelism = context.currentParallelism();
        checkArgument(
                task >= 0 && task < parallelism,
                "Split %s is suggested to task %s, which is out of the parallelism [0, %s). This is unexpected.",
                split.splitId(),
                task,
                parallelism);
        return task;
    }

    protected int assignSuggestedTask(DataSplit split) {
        int parallelism = context.currentParallelism();

        int bucketId;
        if (split.bucket() == BucketMode.POSTPONE_BUCKET) {
            bucketId =
                    PostponeBucketFileStoreWrite.getWriteId(split.dataFiles().get(0).fileName())
                            % parallelism;
        } else {
            bucketId = split.bucket();
        }

        if (shuffleBucketWithPartition) {
            return ChannelComputer.select(split.partition(), bucketId, parallelism);
        } else {
            return ChannelComputer.select(bucketId, parallelism);
        }
    }

    protected int assignSuggestedTask(IncrementalSplit split) {
        int parallelism = context.currentParallelism();

        int bucketId;
        if (split.bucket() == BucketMode.POSTPONE_BUCKET) {
            // A diff which only removes files has no after files, so fall back to the before files.
            List<DataFileMeta> files =
                    split.afterFiles().isEmpty() ? split.beforeFiles() : split.afterFiles();
            bucketId =
                    PostponeBucketFileStoreWrite.getWriteId(files.get(0).fileName()) % parallelism;
        } else {
            bucketId = split.bucket();
        }

        if (shuffleBucketWithPartition) {
            return ChannelComputer.select(split.partition(), bucketId, parallelism);
        } else {
            return ChannelComputer.select(bucketId, parallelism);
        }
    }

    protected int assignSuggestedTask(ChainSplit split) {
        int parallelism = context.currentParallelism();
        // Extract bucket id from the bucket path stored in fileBucketPathMapping.
        // The bucket path ends with "bucket-{id}".
        int bucketId = 0;
        if (!split.fileBucketPathMapping().isEmpty()) {
            String bucketPath = split.fileBucketPathMapping().values().iterator().next();
            int lastSlash = bucketPath.lastIndexOf('/');
            if (lastSlash >= 0) {
                String bucketDir = bucketPath.substring(lastSlash + 1);
                if (bucketDir.startsWith("bucket-")) {
                    try {
                        bucketId = Integer.parseInt(bucketDir.substring("bucket-".length()));
                    } catch (NumberFormatException e) {
                        LOG.warn(
                                "Failed to parse bucket id from path '{}', falling back to 0.",
                                bucketPath,
                                e);
                    }
                }
            }
        }
        if (shuffleBucketWithPartition) {
            return ChannelComputer.select(split.logicalPartition(), bucketId, parallelism);
        } else {
            return ChannelComputer.select(bucketId, parallelism);
        }
    }

    protected SplitAssigner createSplitAssigner(boolean unordered) {
        return unordered
                ? new FIFOSplitAssigner(Collections.emptyList())
                : new PreAssignSplitAssigner(
                        this.splitMaxPerTask, context, Collections.emptyList());
    }

    protected boolean noMoreSplits() {
        return finished;
    }

    /** The result of scan. */
    protected static class PlanWithNextSnapshotId {
        private final TableScan.Plan plan;
        private final Long nextSnapshotId;

        public PlanWithNextSnapshotId(TableScan.Plan plan, Long nextSnapshotId) {
            this.plan = plan;
            this.nextSnapshotId = nextSnapshotId;
        }

        public TableScan.Plan plan() {
            return plan;
        }

        public Long nextSnapshotId() {
            return nextSnapshotId;
        }
    }
}
