package com.it_jaros.jscanner.scan;

import com.it_jaros.jscanner.scan.domain.ScanResult;
import com.it_jaros.jscanner.scan.engine.*;

import java.time.Duration;
import java.util.concurrent.*;
import java.util.function.Consumer;

public class Scanner implements AutoCloseable {

    private final CancelledToken cancelledToken = new CancelledToken();
    private final ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor();
    private final ScanOptions scanOptions;
    private final Semaphore socketLimit;

    public Scanner(ScanOptions options) {
        this.scanOptions = options;
        this.socketLimit = new Semaphore(scanOptions.socketLimit());
    }

    /**
     * Starts a given scan.
     *
     * @param scan
     */
    public void runScan(final Scan scan, final Consumer<ScanResult> consumer) {
        if (scan == null) {
            throw new IllegalArgumentException("Scan argument cannot be null");
        }

        long delayInNanos = TimeUnit.MILLISECONDS.toNanos(Math.max(0, scanOptions.delayInMillis()));
        ScanExecutionContext context = new ScanExecutionContext(
                scanOptions,
                executor,
                socketLimit,
                cancelledToken,
                scan,
                new PortScanRateLimiter(cancelledToken, delayInNanos)
        );

        scan.start();
        scan.producerStart();
        final ProducerState<ScanResult> state = new ProducerThread(context).startProducer(
                scan.getHosts().iterator(),
                scanOptions.maxHostsLimit(),
                host -> new ScanHostTask(scan, host, context)
        );

        // the main thread is the consumer
        // Let the consumer run as long as the producer runs
        // or if still tasks are pending in pipeline
        // we do not listen to canceled here because we want
        // all results (also partial) collected for the consumer
        // with whatever is there already
        while (state.running().get() || state.inPipeline().get() > 0) {
            try {
                PollState<ScanResult> poll = getHostResult(state);
                if (poll instanceof PollState.Success<ScanResult>(ScanResult value)) {
                    consumer.accept(value);
                } else if (poll instanceof PollState.Failure<ScanResult>(Throwable error)) {
                    System.err.printf("runScan(): ScanHostTask() failed with error %s -> %s%n", error.getClass().getSimpleName(), error.getMessage());
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
        scan.producerStop();
        scan.stop();
    }

    private PollState<ScanResult> getHostResult(ProducerState<ScanResult> state) throws InterruptedException {
        Future<ScanResult> finishedHost = state.completionService().poll(ProducerThread.pollInterval.toMillis(), TimeUnit.MILLISECONDS);
        if (finishedHost == null) {
            return new PollState.Unavailable<>();
        }

        try {
            ScanResult result = finishedHost.get();
            return new PollState.Success<>(result);
        } catch (ExecutionException e) {
            return new PollState.Failure<>(e.getCause());
        } finally {
            state.activeWorkers().release();
            state.inPipeline().decrementAndGet();
        }
    }

    public boolean awaitTermination(Duration duration) throws InterruptedException {
        return executor.awaitTermination(duration.toMillis(), TimeUnit.MILLISECONDS);
    }

    public void cancel() {
        this.cancelledToken.signal();
        executor.shutdown();
    }

    public void cancelNow() {
        this.cancelledToken.signal();
        executor.shutdownNow();
    }

    @Override
    public void close() throws Exception {
        cancel();
    }

}