package com.it_jaros.jscanner.scan.engine;

import java.time.Duration;
import java.util.Iterator;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;

public class ProducerThread {

    public static final Duration pollInterval = Duration.ofSeconds(1);

    private final ExecutorService executor;
    private final CancelledToken cancelledToken;

    public ProducerThread(ScanExecutionContext context) {
        this.executor = context.executorService();
        this.cancelledToken = context.cancelledToken();
    }

    /**
     * This method helps to cleanup the code a bit and remove redundancy
     * The producer for providing hosts and the one for providing ports
     * are similar and the small differences can be handled using a function
     *
     * @param queue
     * @param maxWorkers
     * @param taskFactory
     * @param <INPUT>
     * @param <OUTPUT>
     * @return
     */
    public <INPUT, OUTPUT> ProducerState<OUTPUT> startProducer(
            Iterator<INPUT> queue,
            int maxWorkers,
            Function<INPUT, Callable<OUTPUT>> taskFactory
    ) {
        final AtomicInteger inPipeline = new AtomicInteger(0);
        final AtomicBoolean running = new AtomicBoolean(true);
        final Semaphore activeWorkers = new Semaphore(maxWorkers);
        CompletionService<OUTPUT> completionService = new ExecutorCompletionService<>(executor);
        executor.submit(() -> {
            try {
                while (!cancelledToken.isCancelled() && queue.hasNext()) {
                    // get semaphore and remember if task got submitted
                    // so in case we fail to submit we release the semaphore
                    activeWorkers.acquire();
                    boolean isTaskSubmitted = false;
                    try {
                        // just in case something
                        // changed while waiting
                        if (cancelledToken.isCancelled()) {
                            break;
                        }

                        // get next item and create callable
                        // using lambda expression
                        final INPUT item = queue.next();
                        Callable<OUTPUT> task = taskFactory.apply(item);
                        inPipeline.incrementAndGet();
                        try {
                            completionService.submit(task);
                            isTaskSubmitted = true;
                        } catch (Throwable e) {
                            inPipeline.decrementAndGet();
                            throw e;
                        }
                    } finally {
                        if (!isTaskSubmitted) {
                            activeWorkers.release();
                        }
                    }
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            } finally {
                running.set(false);
            }
        });
        return new ProducerState<>(running, inPipeline, activeWorkers, completionService);
    }
}