package com.it_jaros.jscanner;

import java.io.BufferedReader;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.net.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.BitSet;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.LockSupport;
import java.util.stream.Stream;

public class Scanner {

    private final Counter counter = new Counter();
    private final ProgressBar progressBar = new ProgressBar();
    private final Semaphore socketLimit;
    private final int timeoutInMillis;
    private final long delayInNanos;
    private final int maxWorkersPerHost;
    private final int maxHostsLimit;
    private final boolean disableHostCheck;

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

    public List<ScanResult> scan(ScanOptions options) {
        if (options.hostsFile() != null) {
            return scanHosts(options.hostsFile(), options.ports());
        }
        return scanHosts(options.hostsArgv(), options.ports());
    }

    public List<ScanResult> scanHosts(String sourceFile, String ports) {
        try (BufferedReader reader = Files.newBufferedReader(Path.of(sourceFile))) {
            return scanHosts(reader.lines(), ports);
        } catch (FileNotFoundException | NoSuchFileException e) {
            System.out.printf("Could not open file %s%n", sourceFile);
        } catch (IOException e) {
            System.out.printf("Error while trying to read hosts: %s -> %s %n", e.getClass().getSimpleName(), e.getMessage());
        }

        return List.of();
    }

    public List<ScanResult> scanHosts(List<String> hostsArgv, String ports) {
        try (Stream<String> stream = hostsArgv.stream()){
            return scanHosts(stream, ports);
        } catch (IOException e) {
            System.out.printf("Error while trying to read hosts: %s -> %s %n", e.getClass().getSimpleName(), e.getMessage());
        }

        return List.of();
    }

    public List<ScanResult> scanHosts(Stream<String> hosts, String ports) throws IOException {
        List<ScanResult> results = new ArrayList<>();
        progressBar.start();
        try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor()) {
            List<Future<ScanResult>> futures = new ArrayList<>();
            Semaphore maxHosts = new Semaphore(maxHostsLimit);
            hosts.forEach((final String host) -> {
                maxHosts.acquireUninterruptibly();
                futures.add(executor.submit(() -> {
                    try {
                        return scanHost(host, ports);
                    } finally {
                        maxHosts.release();
                    }
                }));
            });

            futures.forEach((Future<ScanResult> f) -> {
                try {
                    results.add(f.get());
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                } catch (ExecutionException e) {
                    System.out.printf("%s -> %s%n", e.getClass().getSimpleName(), e.getMessage());
                }
            });
        }

        progressBar.stop();
        System.out.println("--------------------");
        System.out.println("Finished scan!");
        System.out.println("Stats:");
        System.out.println("Peak concurrent connects: " + counter.max());
        System.out.println("--------------------");
        return results;
    }

    private ScanResult scanHost(String host, String ports) {
        PortRange portRange = new PortRange(ports);
        AtomicLong portSlotFactory = new AtomicLong(System.nanoTime());
        Semaphore maxWorkers = new Semaphore(maxWorkersPerHost);
        BitSet openPorts = new BitSet(PortRange.MAX_PORT);
        BitSet filteredPorts = new BitSet(PortRange.MAX_PORT);
        HashMap<Integer, ServiceType> bannerRecognition = new HashMap<Integer, ServiceType>();

        Progress progress = new Progress(host, portRange.getTotal(), new AtomicInteger(), new AtomicInteger(), new AtomicInteger());
        progressBar.submit(progress);

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

        try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor()) {
            List<Future<PortResult>> futures = new ArrayList<>();
            while (portRange.hasNext()) {
                maxWorkers.acquireUninterruptibly();
                final int currentPort = portRange.next();
                futures.add(executor.submit(() -> {
                    try {
                        waitForSlot(portSlotFactory);
                        return checkPort(host, currentPort);
                    } finally {
                        maxWorkers.release();
                        progress.done().incrementAndGet();
                    }
                }));
            }

            List<Throwable> errors = new ArrayList<>();
            for (Future<PortResult> f : futures) {
                try {
                    // bitset is not thread-safe, so it is set
                    // outside the other virtual threads that update progress
                    PortResult portResult = f.get();
                    if (portResult.getState().equals(PortState.OPEN)) {
                        openPorts.set(portResult.getPort());
                        bannerRecognition.put(portResult.getPort(), ServiceDetector.detect(portResult.getBanner()));
                    } else if (portResult.getState().equals(PortState.FILTERED)) {
                        filteredPorts.set(portResult.getPort());
                    }
                } catch (InterruptedException e) {
                    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, bannerRecognition);
    }

    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) {
        socketLimit.acquireUninterruptibly();
        counter.inc();
        PortResult result = new PortResult();
        result.setPort(port);
        result.setState(PortState.UNKNOWN);
        try (Socket socket = new Socket()) {
            socket.connect(new InetSocketAddress(host, port), timeoutInMillis);
            result.setBanner(getBanner(socket));
            result.setState(PortState.OPEN);
        } catch (NoRouteToHostException ignored) {
            // this can be safely ignored because the port is closed if a host is unreachable
        } catch (SocketTimeoutException ignored) {
            result.setState(PortState.FILTERED);
        } catch (ConnectException ignored) {
            result.setState(PortState.CLOSED);
        } catch (IOException ignored) {
            // Will happen a lot when scanning for open ports, so not needed
        }  finally {
            counter.dec();
            socketLimit.release();
        }

        return result;
    }

    private String getBanner(Socket socket) {
        int READ_BUFFER_SIZE = 1024;
        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;
    }
}