/*
 * 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.flink.api.connector.source.mocks;

import org.apache.flink.api.connector.source.SourceSplit;

import java.io.Serializable;
import java.util.Arrays;
import java.util.Objects;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * Simple Mock SourceSplit for unit test. The implementation of this class is just an in-memory
 * queue. The values are integers and each value has an associated index which is its position in
 * the queue over the entire life cycle of the split. For example, a value with index K means it is
 * the K-th element that was polled out of the queue since the creation of this split.
 */
public class MockSourceSplit implements SourceSplit, Serializable {

    private static final long serialVersionUID = 1L;

    private final int id;
    private final BlockingQueue<Integer> records;
    private final int endIndex;
    private int index;

    public MockSourceSplit(int id) {
        this(id, 0);
    }

    public MockSourceSplit(int id, int startingIndex) {
        this(id, startingIndex, Integer.MAX_VALUE);
    }

    public MockSourceSplit(int id, int startingIndex, int endIndex) {
        this.id = id;
        this.endIndex = endIndex;
        this.index = startingIndex;
        this.records = new LinkedBlockingQueue<>();
    }

    @Override
    public String splitId() {
        return Integer.toString(id);
    }

    public int index() {
        return index;
    }

    public int endIndex() {
        return endIndex;
    }

    public boolean isAvailable() {
        return !isFinished() && !records.isEmpty();
    }

    public boolean isFinished() {
        return index == endIndex;
    }

    /** Get the next element. Block if asked. */
    public int[] getNext(boolean blocking) throws InterruptedException {
        Integer value = blocking ? records.take() : records.poll();
        return value == null ? null : new int[] {value, index++};
    }

    /** Add a record to this split. */
    public MockSourceSplit addRecord(int record) {
        if (!records.offer(record)) {
            throw new IllegalStateException("Failed to add record to split.");
        }
        return this;
    }

    @Override
    public int hashCode() {
        return Objects.hash(id, Arrays.hashCode(records.toArray(new Integer[0])), endIndex, index);
    }

    @Override
    public boolean equals(Object obj) {
        if (!(obj instanceof MockSourceSplit)) {
            return false;
        }
        MockSourceSplit that = (MockSourceSplit) obj;
        return id == that.id
                && index == that.index
                && Arrays.equals(
                        records.toArray(new Integer[0]), that.records.toArray(new Integer[0]))
                && endIndex == that.endIndex;
    }

    @Override
    public String toString() {
        return String.format(
                "MockSourceSplit(id=%d, num_records=%d, endIndex=%d, currentIndex=%d)",
                id, records.size(), endIndex, index);
    }
}
