package com.jillesvangurp.iterables;

import java.io.Closeable;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * Iterable that processes the input concurrently using a {@link Processor} to produce its output.
 * 
 * Please note that this class implements {@link Closeable} and that you are supposed to use a try with resources
 * call. The reason for this is to guarantee the executor used for delegating the work is shut down correctly.
 * 
 * @param <Input> type of the input processed by this iterable
 * @param <Output> type of the output processed by this iterable
 */
public class ConcurrentProcessingIterable<Input, Output> implements Iterable<Output>, Closeable {

    private final int blockSize;
    private final int threadPoolSize; 

    private final Processor<Input, Output> processor;
    private final Iterable<Input> input;
    private final ExecutorService executorService;
    private final LinkedBlockingQueue<List<Input>> scheduledWork;
    private final LinkedBlockingQueue<List<Output>> completedWork;
    private final AtomicBoolean doneProducing = new AtomicBoolean(false);

    /**
     * Create a new iterable.
     * 
     * @param input
     *            iterable with the input
     * @param processor
     *            {@link Processor} that processes each element in the input
     * @param blockSize
     *            size of the list of elements that is processed by the worker threads. Having a large list means there
     *            is less contention on the queues by the threads.
     * @param threadPoolSize
     *            number of threads used. What is sensible very much depends on the work load of the processor.
     *            Generally you don't want to have more threads than CPU cores + one for the producer thread used
     *            internally to queue stuff for the worker threads. 
     * @param queueCapacity
     */
    public ConcurrentProcessingIterable(Iterable<Input> input, Processor<Input, Output> processor, int blockSize, int threadPoolSize, int queueCapacity) {
        this.input = input;
        this.processor = processor;
        this.blockSize = blockSize;
        this.threadPoolSize = threadPoolSize;
        executorService = Executors.newFixedThreadPool(threadPoolSize, new ThreadFactory() {
            @Override
            public Thread newThread(Runnable r) {
                return new Thread(r, "processor");
            }
        });
        scheduledWork = new LinkedBlockingQueue<>(queueCapacity);
        completedWork = new LinkedBlockingQueue<>(queueCapacity);
    }

    @Override
    public Iterator<Output> iterator() {
        // start the producer thread
        executorService.execute(new Runnable() {
            @Override
            public void run() {
                try {
                    ArrayList<Input> block = new ArrayList<>(blockSize);
                    for (Input i : input) {
                        block.add(i);
                        if (block.size() == blockSize) {
                            scheduledWork.put(block);
                            block = new ArrayList<>(blockSize);
                        }
                    }
                    ;
                    if (block.size() > 0) {
                        scheduledWork.put(block);
                    }
                    doneProducing.set(true);
                } catch (InterruptedException e) {
                    throw new IllegalStateException(e);
                } 
            }
        });
        final CountDownLatch latch = new CountDownLatch(threadPoolSize-1);

        // start consumer threads
        for (int i = 0; i < threadPoolSize - 1; i++) {
            executorService.execute(new Runnable() {
                @Override
                public void run() {
                    List<Input> block;
                    try {
                        while ((block = scheduledWork.poll(100, TimeUnit.MILLISECONDS)) != null || !doneProducing.get()) {
                            if(block != null) {
                                ArrayList<Output> outputBlock = new ArrayList<>(blockSize);
                                for (Input input : block) {
                                    outputBlock.add(processor.process(input));
                                }
                                if(outputBlock.size() >0) {
                                    completedWork.put(outputBlock);
                                }
                            }
                        }
                    } catch (InterruptedException e) {
                        throw new IllegalStateException(e);
                    } finally {
                        latch.countDown();
                    }
                }
            });
        }

        return new Iterator<Output>() {
            Output next = null;
            List<Output> currentBlock = null;
            int blockIndex = 0;

            @Override
            public boolean hasNext() {
                try {
                    if (next != null) {
                        return true;
                    } else if (currentBlock != null && blockIndex < currentBlock.size()) {
                        next = currentBlock.get(blockIndex++);
                        return true;
                    } else if ((currentBlock = completedWork.poll(100,TimeUnit.MILLISECONDS)) != null) {
                        blockIndex = 0;
                        return hasNext();                       
                    } else if(doneProducing.get() && scheduledWork.size() == 0 && completedWork.size() ==0 && latch.getCount()==0) {
                        return false;
                    } else {
                        return hasNext();
                    }
                } catch (InterruptedException e) {
                    throw new IllegalStateException(e);
                }
            }
            
            @Override
            public Output next() {
                if (hasNext()) {
                    Output result = next;
                    next = null;
                    return result;
                } else {
                    throw new NoSuchElementException();
                }
            }

            @Override
            public void remove() {
                throw new UnsupportedOperationException();
            }
        };
    }

    @Override
    public void close() throws IOException {
        executorService.shutdown();
        try {
            executorService.awaitTermination(1, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            throw new IllegalStateException("executor failed to shut down cleanly within 1 second");
        }
    }
}

@@ -737,17 +737,17 @@
 rces
-%0A *
  call.
+%0A *
  The
@@ -862,24 +862,38 @@
 aram %3CInput%3E
+%0A *           
  type of the
@@ -944,16 +944,30 @@
 %3COutput%3E
+%0A *           
  type of
@@ -1178,17 +1178,16 @@
 oolSize;
- 
 %0A%0A    pr
@@ -2276,17 +2276,16 @@
 threads.
- 
 %0A     * 
@@ -3732,53 +3732,8 @@
   %7D%0A
-                    doneProducing.set(true);%0A
     
@@ -3853,16 +3853,88 @@
       %7D 
+finally %7B%0A                    doneProducing.set(true);%0A                %7D
 %0A       
@@ -3981,21 +3981,31 @@
 wnLatch 
-latch
+activeConsumers
  = new C
@@ -4032,17 +4032,19 @@
 PoolSize
--
+ - 
 1);%0A%0A   
@@ -4462,16 +4462,17 @@
       if
+ 
 (block !
@@ -4780,16 +4780,17 @@
       if
+ 
 (outputB
@@ -4802,16 +4802,17 @@
 size() %3E
+ 
 0) %7B%0A   
@@ -5137,21 +5137,31 @@
         
-latch
+activeConsumers
 .countDo
@@ -5805,16 +5805,17 @@
 oll(100,
+ 
 TimeUnit
@@ -5922,39 +5922,16 @@
 sNext();
-                       
 %0A       
@@ -5952,16 +5952,17 @@
  else if
+ 
 (donePro
@@ -6033,18 +6033,29 @@
 ) ==
+ 
 0 && 
-latch
+activeConsumers
 .get
@@ -6061,18 +6061,20 @@
 tCount()
+ 
 ==
+ 
 0) %7B%0A   
@@ -6340,28 +6340,16 @@
       %7D%0A
-            
 %0A       

make concurrent processor more robust