/*
 * Copyright (C) 2014 Indeed Inc.
 *
 * Licensed 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 com.indeed.flamdex.utils;

import com.google.common.base.Throwables;
import com.google.common.collect.AbstractIterator;

import java.util.Iterator;
import java.util.concurrent.BrokenBarrierException;
import java.util.concurrent.CyclicBarrier;

/**
 * @author jsadun
 */
public class BlockingCopyableIterator<E> implements Iterable<E> {
    private final CyclicBarrier barrier;
    private final E[] buffer;
    private int finished;

    public BlockingCopyableIterator(final Iterator<E> it, int numConsumers, int bufferSize) {
        barrier = new CyclicBarrier(numConsumers, new Runnable() {
            @Override
            public void run() {
                for (int i = 0; i < buffer.length; i++) {
                    if (it.hasNext()) {
                        buffer[i] = it.next();
                    } else {
                        finished = i;
                        break;
                    }
                }
            }
        });
        // noinspection unchecked
        buffer = (E[]) new Object[bufferSize];
        finished = -1;
    }

    @Override
    public Iterator<E> iterator() {
        return new AbstractIterator<E>() {
            private int myLoc = buffer.length;

            @Override
            protected E computeNext() {
                try {
                    if (myLoc == buffer.length) {
                        barrier.await();
                        myLoc = 0;
                    }
                    if (myLoc == finished) {
                        return endOfData();
                    }
                    final E ret = buffer[myLoc];
                    myLoc++;
                    return ret;
                } catch (InterruptedException e) {
                    throw Throwables.propagate(e);
                } catch (BrokenBarrierException e) {
                    throw Throwables.propagate(e);
                }
            }
        };
    }
}
