package com.it_jaros.jscanner;

import java.time.Duration;
import java.util.BitSet;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.stream.Collectors;

public class App {

    private static Thread shutdownHook;
    private static final CountDownLatch shutdownLatch = new CountDownLatch(1);
    private static final int shutdownTimeoutInMillis = 10000;

    public static void main(String[] args) {
        try {
            run(args);
        } catch (Exception e) {
            System.out.printf("ERROR: %s", e.getMessage());
        }
    }

    private static void run(String[] args) throws Exception {
        ScanOptions options = parseOptions(args);
        if (options == null) return;
        try (Scanner scanner = new Scanner(options)) {
            Scan scan = Scan.create(options);
            addShutdownHook(scanner);
            printScanResults(scanner.runScan(scan), scan.getPeakConcurrentConnects(), options.showFilteredPorts());
        } finally {
            shutdownLatch.countDown();
            removeShutdownHook();
        }
    }

    private static ScanOptions parseOptions(String[] args) {
        CliParser cliParser = new CliParser();
        try {
            return cliParser.parseArgs(args);
        } catch (IllegalArgumentException e) {
            System.out.printf("Error while parsing arguments: %s", e.getMessage());
            System.out.println();
        }

        cliParser.usageAndExit(1);
        return null;
    }

    private static void printScanResults(List<ScanResult> result, int max, boolean showFilteredPorts) {
        printStats(max);
        result.forEach((ScanResult r) -> printPortsResults(r, showFilteredPorts));
    }

    private static void printStats(int max) {
        System.out.println("--------------------");
        System.out.println("Stats:");
        System.out.println("Peak concurrent connects: " + max);
        System.out.println("--------------------");
    }

    private static void printPortsResults(ScanResult result, boolean showFilteredPorts) {
        if (result.openPorts().isEmpty() && (!showFilteredPorts || result.filteredPorts().isEmpty())) {
            return;
        }

        System.out.println(result.host());
        if (!result.openPorts().isEmpty()) {
            System.out.printf("\t (%d) open:\t%s%n", result.openPorts().cardinality(), map(result.openPorts(), result.bannerRecognition()));
        }
        if (showFilteredPorts && !result.filteredPorts().isEmpty()) {
            System.out.printf("\t (%d) filtered:\t%s%n", result.filteredPorts().cardinality(), map(result.filteredPorts()));
        }
        System.out.println();
    }

    private static String map(BitSet ports) {
        return ports.stream().mapToObj(String::valueOf).collect(Collectors.joining(","));
    }

    private static String map(BitSet ports, java.util.HashMap<Integer, ServiceType> bannerRecognition) {
        return ports.stream().mapToObj(port -> {
            ServiceType serviceType = bannerRecognition.getOrDefault(port, ServiceType.UNKNOWN);

            if (serviceType == ServiceType.UNKNOWN) {
                return String.valueOf(port);
            }

            return port + "/" + serviceType.name().toLowerCase();
        }).collect(Collectors.joining(","));
    }

    /**
     * Notice user that the signal was received
     * Start shutdown of scanner
     *
     * @param scanner
     */
    private static void addShutdownHook(Scanner scanner) {
        shutdownHook = new Thread(() -> {
            System.err.printf("%nCtrl-C received. Shutting down... Waiting for timeout of %d seconds%n", shutdownTimeoutInMillis / 1000);
            shutdown(scanner);
        });
        Runtime.getRuntime().addShutdownHook(shutdownHook);
    }

    /**
     * Try to gracefully shutdown scanner
     * and wait
     *
     * @param scanner
     */
    private static void shutdown(Scanner scanner) {
        try {
            scanner.cancel();
            scanner.awaitTermination(Duration.ofMillis(shutdownTimeoutInMillis));
            shutdownLatch.await(shutdownTimeoutInMillis, TimeUnit.MILLISECONDS);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        } catch (TimeoutException e) {
            // Time ran out
        }
    }

    private static void removeShutdownHook() {
        if (shutdownHook == null) {
            return;
        }

        try {
            Runtime.getRuntime().removeShutdownHook(shutdownHook);
        } catch (IllegalStateException ignored) {
            // JVM is already shutting down
        }
    }
}