/*
 * Copyright (C) 2021 Grakn Labs
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 *
 */

package grakn.core.common.producer;

import grakn.common.collection.Either;
import grakn.core.common.concurrent.ExecutorService;
import grakn.core.common.concurrent.ManagedBlockingQueue;
import grakn.core.common.exception.GraknException;
import grakn.core.common.iterator.ResourceIterator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;
import javax.annotation.concurrent.ThreadSafe;
import java.util.LinkedList;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;

public class ProducerIterator<T> implements ResourceIterator<T> {

    private static final Logger LOG = LoggerFactory.getLogger(ProducerIterator.class);
    private static final int BUFFER_MIN_SIZE = 32;
    private static final int BUFFER_MAX_SIZE = 64;

    private final LinkedList<Producer<T>> producers;
    private final Queue queue;

    private T next;
    private State state;

    public ProducerIterator(List<Producer<T>> producers) {
        this(producers, BUFFER_MIN_SIZE, BUFFER_MAX_SIZE);
    }

    public ProducerIterator(List<Producer<T>> producers, int bufferMinSize, int bufferMaxSize) {
        // TODO: Could we optimise IterableProducer by accepting ResourceIterator<Producer<T>> instead?
        assert !producers.isEmpty();
        this.producers = new LinkedList<>(producers);
        this.queue = new Queue(bufferMinSize, bufferMaxSize);
        this.state = State.EMPTY;
    }

    private void mayProduce() {
        synchronized (queue) {
            if (producers.isEmpty()) return;
            int available = queue.max - queue.size() - queue.pending;
            if (available > queue.max - queue.min) {
                queue.pending += available;
                assert !producers.isEmpty();
                Producer<T> producer = producers.peek();
                ExecutorService.forkJoinPool().submit(() -> producer.produce(queue, available));
            }
        }
    }

    private enum State {EMPTY, FETCHED, COMPLETED}

    @Override
    public boolean hasNext() {
        if (state == State.COMPLETED) return false;
        else if (state == State.FETCHED) return true;
        else mayProduce();

        Either<T, Done> result = queue.take();

        if (result.isFirst()) {
            next = result.first();
            state = State.FETCHED;
        } else {
            Done done = result.second();
            recycle();
            state = State.COMPLETED;
            if (done.error().isPresent()) {
                throw GraknException.of(done.error().get());
            }
        }

        return state == State.FETCHED;
    }

    @Override
    public T next() {
        if (!hasNext()) throw new NoSuchElementException();
        state = State.EMPTY;
        return next;
    }

    @Override
    public void recycle() {
        synchronized (queue) {
            producers.forEach(Producer::recycle);
        }
    }

    private static class Done {
        @Nullable
        private final Throwable error;

        private Done(Throwable error) {
            this.error = error;
        }

        public static Done success() {
            return new Done(null);
        }

        public static Done error(Throwable e) {
            return new Done(e);
        }

        public Optional<Throwable> error() {
            return Optional.ofNullable(error);
        }
    }

    @ThreadSafe
    private class Queue implements Producer.Queue<T> {

        private final ManagedBlockingQueue<Either<T, Done>> blockingQueue;
        private final AtomicBoolean isError;
        private final int min;
        private final int max;
        private int pending;

        private Queue(int bufferMinSize, int max) {
            this.min = bufferMinSize;
            this.max = max;
            this.blockingQueue = new ManagedBlockingQueue<>();
            this.isError = new AtomicBoolean(false);
            this.pending = 0;
        }

        @Override
        public synchronized void put(T item) {
            try {
                blockingQueue.put(Either.first(item));
                pending--;
                assert pending >= 0 || isError.get();
            } catch (InterruptedException e) {
                throw GraknException.of(e);
            }
        }

        @Override
        public synchronized void done() {
            done(null);
        }

        @Override
        public synchronized void done(@Nullable Throwable error) {
            assert !producers.isEmpty();
            producers.removeFirst();
            pending = 0;
            try {
                if (error != null) blockingQueue.put(Either.second(Done.error(error)));
                else if (producers.isEmpty()) blockingQueue.put(Either.second(Done.success()));
                else mayProduce();
            } catch (InterruptedException e) {
                throw GraknException.of(e);
            }
        }

        private Either<T, Done> take() {
            try {
                return blockingQueue.take();
            } catch (InterruptedException e) {
                throw GraknException.of(e);
            }
        }

        private int size() {
            return blockingQueue.size();
        }
    }
}

@@ -1141,37 +1141,8 @@
 fe;%0A
-import java.util.LinkedList;%0A
 impo
@@ -1160,16 +1160,16 @@
 l.List;%0A
+
 import j
@@ -1228,16 +1228,67 @@
 tional;%0A
+import java.util.concurrent.ConcurrentLinkedQueue;%0A
 import j
@@ -1587,16 +1587,111 @@
  = 64;%0A%0A
+    // TODO: why does 'producers' have to be ConcurrentLinkedQueue? see method recycle() below%0A
     priv
@@ -1700,26 +1700,37 @@
 e final 
+Concurrent
 Linked
-List
+Queue
 %3CProduce
@@ -2220,18 +2220,29 @@
 new 
+Concurrent
 Linked
-List
+Queue
 %3C%3E(p
@@ -3730,32 +3730,70 @@
 cle() %7B%0A        
+// TODO: If this method is wrapped in 
 synchronized (qu
@@ -3787,32 +3787,170 @@
 chronize
-d 
 (queue)
- %7B%0A    
+, we won't need producers to be ConcurrentLinkedQueue%0A        //       However, doing so would also cause a deadlock. Let's investigate this soon.%0A
         
@@ -3987,26 +3987,16 @@
 cycle);%0A
-        %7D%0A
     %7D%0A%0A 
@@ -5559,32 +5559,32 @@
 cers.isEmpty();%0A
+
             prod
@@ -5595,21 +5595,16 @@
 s.remove
-First
 ();%0A    

Reimplement ProducerIterator.recycle() with ConcurrentLinkedQueue