package com.it_jaros.jscanner;

import java.io.IOException;
import java.io.InputStream;
import java.net.*;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.util.function.Function;

public class Scanner implements AutoCloseable {

    private static final int READ_BUFFER_SIZE = 1024;

    private volatile boolean cancelled = false;

    private final ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor();
    private final Semaphore socketLimit;
    private final boolean bannerRecognition;
    private final boolean disableOnlineCheck;
    private final int maxHostsLimit;
    private final int maxWorkersPerHost;
    private final int timeoutInMillis;
    private final long delayInNanos;

    public Scanner(
            int socketLimit,
            int timeoutInMillis,
            int delayInMillis,
            int maxWorkersPerHost,
            int maxHostsLimit,
            boolean disableOnlineCheck,
            boolean bannerRecognition
    ) {
        this.delayInNanos = TimeUnit.MILLISECONDS.toNanos(Math.max(0, delayInMillis));
        this.bannerRecognition = bannerRecognition;
        this.disableOnlineCheck = disableOnlineCheck;
        this.maxHostsLimit = maxHostsLimit;
        this.maxWorkersPerHost = maxWorkersPerHost;
        this.socketLimit = new Semaphore(socketLimit);
        this.timeoutInMillis = timeoutInMillis;
    }

    public Scanner(ScanOptions options) {
        this(options.socketLimit(), options.timeoutInMillis(), options.delayInMillis(), options.maxWorkersPerHost(), options.maxHostsLimit(), options.disableOnlineCheck(), options.bannerRecognition());
    }

    /**
     * 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");
        }

        scan.start();

        // start producer thread
        scan.producerStart();
        final ProducerState<ScanResult> state = startProducer(
                scan.getHosts().iterator(),
                maxHostsLimit,
                host -> new ScanHostTask(scan, host)
        );

        // 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.println("runScan: " + 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(10, 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();
        }
    }

    /**
     * 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
     */
    private <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 (!cancelled && 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 (cancelled) {
                            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);
    }

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

    public void cancel() {
        if (!cancelled) {
            cancelled = true;
        }
        executor.shutdown();
    }

    public void cancelNow() {
        if (!cancelled) {
            cancelled = true;
        }
        executor.shutdownNow();
    }

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


    private final class ScanHostTask implements Callable<ScanResult> {
        private final PortScanRateLimiter rateLimiter = new PortScanRateLimiter();

        private final Scan scan;
        private final String host; // input parameter

        ScanHostTask(Scan scan, String host) {
            this.scan = scan;
            this.host = host;
        }

        @Override
        public ScanResult call() {
            try {
                scan.hostStart();
                if (cancelled) {
                    return ScanResult.empty(host);
                }
                return scanHostPorts();
            } finally {
                scan.hostFinish();
            }
        }

        private ScanResult scanHostPorts() {
            if (!disableOnlineCheck && !checkHostOnline()) {
                // Unreachable host
                return ScanResult.empty(host);
            }

            final PortRange portRange = new PortRange(scan.getPorts());
            // producer thread
            scan.producerStart();
            ProducerState<PortResult> state = startProducer(
                    portRange.iterator(),
                    maxWorkersPerHost,
                    port -> new ScanPortTask(host, port, scan, rateLimiter)
            );

            // consumer is the main thread
            // we run as long as the producer is running or as long as things are in pipeline to be processed
            // only exception is when cancelled is set
            final PortResultAccumulator accumulator = new PortResultAccumulator(host);
            while (!cancelled && (state.running().get() || state.inPipeline().get() > 0)) {
                try {
                    PollState<PortResult> poll = getPortResult(state);
                    if (poll instanceof PollState.Success<PortResult>(PortResult value)) {
                        accumulator.add(value);
                    } else if (poll instanceof PollState.Failure(Throwable error)) {
                        System.err.printf("scanHostPorts: %s -> %s%n", error.getClass().getSimpleName(), error.getMessage());
                    }
                } catch (InterruptedException ignored) {
                    Thread.currentThread().interrupt();
                }
            }
            scan.producerStop();

            return accumulator.build();
        }

        private boolean checkHostOnline() {
            try {
                return InetAddress.getByName(host).isReachable(timeoutInMillis);
            } catch (IOException e) {
                // we ignore this error because it means that the host is probably not online
            }

            return false;
        }

        private PollState<PortResult> getPortResult(ProducerState<PortResult> state) throws InterruptedException {
            Future<PortResult> portResultFuture = state.completionService().poll(10, TimeUnit.MILLISECONDS);
            if (portResultFuture == null) {
                return new PollState.Unavailable<>();
            }

            PortResult portResult;
            try {
                portResult = portResultFuture.get();
                return new PollState.Success<>(portResult);
            } catch (ExecutionException e) {
                Throwable cause = e.getCause();
                return new PollState.Failure<>(cause);
            } finally {
                state.activeWorkers().release();
                state.inPipeline().decrementAndGet();
            }
        }
    }

    /**
     * Per-host rate limiter. Each ScanHostTask creates its own instance and
     * shares it with all its ScanPortTasks via constructor.
     *
     * Java allows one inner class to access another's private members, so this works.
     */
    private final class PortScanRateLimiter {
        private final Object lock = new Object();
        private volatile long nextAllowedTime;

        void apply() throws InterruptedException {
            if (delayInNanos <= 0) {
                return;
            }
            synchronized (lock) {
                if (cancelled) {
                    return;
                }
                long now = System.nanoTime();
                if (nextAllowedTime > now) {
                    Thread.sleep(Duration.ofNanos(nextAllowedTime - now));
                    now = System.nanoTime();  // re-read after waking
                }
                nextAllowedTime = now + delayInNanos;
            }
        }
    }

    private final class ScanPortTask implements Callable<PortResult> {
        private final Scan scan;
        private final String host;
        private final int port;
        private final PortScanRateLimiter portScanRateLimiter;  // per-host shared limiter

        private ScanPortTask(String host, int port, Scan scan, PortScanRateLimiter portScanRateLimiter) {
            this.host = host;
            this.port = port;
            this.scan = scan;
            this.portScanRateLimiter = portScanRateLimiter;
        }

        @Override
        public PortResult call() throws Exception {
            try {
                scan.portStart();
                portScanRateLimiter.apply();
                if (cancelled) {
                    return PortResult.empty();
                }
                return checkPort();
            } finally {
                scan.portFinish();
            }
        }

        private PortResult checkPort() throws InterruptedException {
            PortResult result = new PortResult();
            result.setPort(port);
            result.setState(PortState.UNKNOWN);

            socketLimit.acquire();
            scan.socketStart();
            try(
                    SocketChannel socketChannel = SocketChannel.open();
                    Selector isOpen = Selector.open()
            ) {
                socketChannel.configureBlocking(false);
                SelectionKey isOpenKey = socketChannel.register(isOpen, SelectionKey.OP_CONNECT);
                socketChannel.connect(new InetSocketAddress(host, port));
                long deadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(timeoutInMillis);
                long waitInMillis = Math.min(timeoutInMillis, 5);
                while (isOpenKey.isValid() && !isOpenKey.isConnectable()) {
                    if (isOpen.select(waitInMillis) == 1) {
                        if (socketChannel.finishConnect()) {
                            result.setState(PortState.OPEN);
                            //if (bannerRecognition) {
                            //    result.setBanner(getBanner(socket));
                            //}
                        } else {
                            result.setState(PortState.FILTERED);
                        }
                        break;
                    }

                    if (deadlineNanos - System.nanoTime() < 0) {
                        result.setState(PortState.FILTERED);
                        break;
                    }
                }
            } catch (SocketTimeoutException ignored) {
                result.setState(PortState.FILTERED);
            } catch (ConnectException ignored) {
                result.setState(PortState.CLOSED);
            } catch (NoRouteToHostException ignored) {
                // NoRouteToHostException: this can be safely ignored because the port is closed if a host is unreachable
                // Will happen a lot when scanning for open ports, so not needed
            } catch (IOException e) {
                result.setException(e);
            } finally {
                scan.socketFinish();
                socketLimit.release();
            }

            return result;
        }

        private String getBanner(Socket socket) {
            byte[] buffer = new byte[READ_BUFFER_SIZE];
            try (InputStream input = socket.getInputStream()) {
                socket.setSoTimeout(timeoutInMillis);
                int bytesRead = input.read(buffer);
                if (bytesRead <= 0) {
                    return null;
                }
                return new String(buffer, 0, bytesRead, StandardCharsets.UTF_8).trim();
            } catch (IOException e) {
                // we ignore this failure
            }
            return null;
        }
    }

    private final class PortResultAccumulator {
        private final String host;
        private final BitSet openPorts = new BitSet(PortRange.MAX_PORT);
        private final BitSet filteredPorts = new BitSet(PortRange.MAX_PORT);
        private final Map<Integer, ServiceType> serviceTypes = new HashMap<>();
        private final List<ScanFailure> scanFailures = new ArrayList<>();

        private PortResultAccumulator(String host) {
            this.host = host;
        }

        void add(PortResult portResult) {
            switch (portResult.getState()) {
                case OPEN -> {
                    openPorts.set(portResult.getPort());
                    serviceTypes.put(portResult.getPort(), ServiceDetector.detect(portResult.getBanner()));
                }
                case FILTERED -> {
                    filteredPorts.set(portResult.getPort());
                }
                default -> {
                    // intentional no-op for uncovered port states
                }
            }
            Exception e = portResult.getException();
            if (e != null) {
                scanFailures.add(new ScanFailure(portResult.getPort(), ExceptionInfo.from(e)));
            }
        }

        ScanResult build() {
            return new ScanResult(host, openPorts, filteredPorts, Collections.unmodifiableMap(serviceTypes), Collections.unmodifiableList(scanFailures));
        }
    }
}