package com.it_jaros.jscanner;

import java.io.IOException;
import java.io.InputStream;
import java.net.*;
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.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.LockSupport;
import java.util.function.Function;
import java.util.stream.Stream;

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.openSocketLimit(), options.timeoutInMillis(), options.delayInMillis(), options.maxWorkersPerHost(), options.maxHostsLimit(), options.disableHostCheck(), options.bannerRecognition());
    }


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

        scan.start();

        // start producer thread
        final ProducerState state = startProducer(
                scan.getHosts().iterator(),
                maxHostsLimit,
                host -> () -> scanHostPorts(host, scan)
        );
        CompletionService<ScanResult> completionService = state.completionService();

        // the main thread is the consumer
        // Let the consumer run as long as
        // if there is something in queue and the scanner is not canceled
        //   -> the producer did not have time yet to add new tasks
        // or if in queue is still some workers left that need to finish
        while ((state.running().get() && !cancelled) || state.inPipeline().get() > 0) {
            try {
                // let's check for results and give add them to our scan data holder object
                Future<ScanResult> finishedHost = completionService.poll(10, TimeUnit.MILLISECONDS);
                if (finishedHost == null) {
                    continue;
                }

                // No matter what happens we have to free the resources after getting ScanResult
                try {
                    ScanResult result = finishedHost.get();
                    // we only want results that give value and not cost RAM for nothing
                    if (!result.openPorts().isEmpty() || !result.filteredPorts().isEmpty()) {
                        scan.addScanResult(result);
                    }
                } finally {
                    state.activeWorkers().release();
                    state.inPipeline().decrementAndGet();
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            } catch (ExecutionException e) {
                System.out.printf("%s -> %s%n", e.getClass().getSimpleName(), e.getMessage());
            }
        }

        scan.stop();
        return scan.getResults();
    }

    private ProducerState startProducer(Iterator<String> queue, int maxWorkers, Function<String, Callable<ScanResult>> taskFactory) {
        final AtomicInteger inPipeline = new AtomicInteger(0);
        final AtomicBoolean running = new AtomicBoolean(true);
        final Semaphore activeWorkers = new Semaphore(maxWorkers);
        CompletionService<ScanResult> completionService = new ExecutorCompletionService<>(executor);
        executor.submit(() -> {
            try {
                while (queue.hasNext() && !cancelled) {
                    activeWorkers.acquire();
                    // just in case while waiting something changed
                    if (cancelled) {
                        break;
                    }

                    // here we are filling the completion service
                    // host <-> virtual thread
                    final String item = queue.next();
                    inPipeline.incrementAndGet();
                    completionService.submit(taskFactory.apply(item));
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            } finally {
                running.set(false);
            }
        });
        return new ProducerState(running, inPipeline, activeWorkers, completionService);
    }

    private ScanResult scanHostPorts(final String host, final Scan scan) {
        PortRange portRange = new PortRange(scan.getPorts());
        BitSet openPorts = new BitSet(PortRange.MAX_PORT);
        BitSet filteredPorts = new BitSet(PortRange.MAX_PORT);
        HashMap<Integer, ServiceType> serviceTypes = new HashMap<>();

        // Give progressbar the current progress object which is then updated in the sub virtual threads
        Progress progress = new Progress(host, portRange.getTotal(), new AtomicInteger(), new AtomicInteger(), new AtomicInteger());
        scan.addHostProgress(progress);

        if (!disableOnlineCheck && !checkHostOnline(host)) {
            // Visually show that this host is basically done
            progress.done().set(progress.total());
            return new ScanResult(host, openPorts, filteredPorts, serviceTypes);
        }

        CompletionService<PortResult> completionService = new ExecutorCompletionService<>(executor);
        List<Throwable> errors = new ArrayList<>();
        final AtomicLong portSlotFactory = new AtomicLong(System.nanoTime());
        final Semaphore activePerHostWorkers = new Semaphore(maxWorkersPerHost);
        final AtomicInteger inPipeline = new AtomicInteger(0);
        final AtomicBoolean producerWorking = new AtomicBoolean(true);

        // producer thread
        executor.submit(() -> {
            try {
                while (portRange.hasNext() && !cancelled) {
                    activePerHostWorkers.acquire();
                    // Here we are filling the completion service
                    // port <-> virtual thread
                    final int currentPort = portRange.next();
                    inPipeline.incrementAndGet();
                    try {
                        completionService.submit(() -> {
                            waitForSlot(portSlotFactory);
                            return checkPort(host, currentPort, scan);
                        });
                    } catch (RuntimeException e) {
                        inPipeline.decrementAndGet();
                        throw e;
                    }
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            } finally {
                producerWorking.set(false);
            }
        });

        // consumer is the main thread
        while ((producerWorking.get() && !cancelled) || inPipeline.get() > 0) {
            try {
                Future<PortResult> portResultFuture = completionService.poll(10, TimeUnit.MILLISECONDS);
                if (portResultFuture == null) {
                    continue;
                }

                try {
                    PortResult portResult = portResultFuture.get();
                    switch (portResult.getState()) {
                        case OPEN -> {
                            // bitset is not thread-safe, so it is set
                            // outside the other virtual threads that update progress
                            openPorts.set(portResult.getPort());
                            progress.open().incrementAndGet();
                            serviceTypes.put(portResult.getPort(), ServiceDetector.detect(portResult.getBanner()));
                        }
                        case FILTERED -> {
                            filteredPorts.set(portResult.getPort());
                            progress.filtered().incrementAndGet();
                        }
                        default -> {
                            // sonarcube glücklich machen
                        }
                    }
                } finally {
                    activePerHostWorkers.release();
                    inPipeline.decrementAndGet();
                    progress.done().incrementAndGet();
                }
            } catch (InterruptedException ignored) {
                Thread.currentThread().interrupt();
            } catch (ExecutionException e) {
                errors.add(e);
            }
        }

        if (!errors.isEmpty()) {
            System.out.printf("Errors happened during scan of host %s%nErrors:%s -> %s", host, errors.size(), errors);
        }

        return new ScanResult(host, openPorts, filteredPorts, serviceTypes);
    }

    private boolean checkHostOnline(String host) {
        try {
            return InetAddress.getByName(host).isReachable(timeoutInMillis);
        } catch (IOException e) {
            // something went wrong
        }

        return false;
    }

    private void waitForSlot(AtomicLong scanSlotFactory) {
        if (delayInNanos > 0) {
            long slot = scanSlotFactory.getAndAdd(delayInNanos);
            long wait = slot - System.nanoTime();
            if (wait > 0)
                LockSupport.parkNanos(wait);
        }
    }

    private PortResult checkPort(String host, int port, Scan scan) throws InterruptedException {
        PortResult result = new PortResult();
        result.setPort(port);
        result.setState(PortState.UNKNOWN);

        // don't allow more sockets then specified
        socketLimit.acquire();
        // count how many ports are concurrently checked
        scan.incSocketCounter();
        try (Socket socket = new Socket()) {
            socket.connect(new InetSocketAddress(host, port), timeoutInMillis);
            result.setState(PortState.OPEN);
            if (bannerRecognition) {
                result.setBanner(getBanner(socket));
            }
        } catch (SocketTimeoutException ignored) {
            result.setState(PortState.FILTERED);
        } catch (ConnectException ignored) {
            result.setState(PortState.CLOSED);
        } catch (IOException 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
        } finally {
            scan.decSocketCounter();
            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;
    }

    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();
    }

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