package com.it_jaros.networkScanner;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.SocketException;
import java.net.SocketTimeoutException;
import java.util.ArrayList;
import java.util.BitSet;
import java.util.List;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.LockSupport;

public class Scanner {

    private final Counter counter = new Counter();
    private final ProgressBar progressBar = new ProgressBar();
    private final Semaphore maxWorkers = new Semaphore(100000);
    private final Semaphore socketLimit;
    private final int timeoutInMillis;
    private final long delayInNanos;

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

    public List<ScanResult> scan(ScanOptions target) {
        return scanTargets(target.targets(), target.ports());
    }

    public List<ScanResult> scanTargets(List<String> targets, String ports) {
        List<ScanResult> results = new ArrayList<>();
        progressBar.start();
        try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor();) {
            List<Future<ScanResult>> futures = new ArrayList<>();
            for (String target : targets) {
                futures.add(executor.submit(() -> scanTarget(target, ports)));
            }

            futures.forEach((f) -> {
                try {
                    results.add(f.get());
                } catch (InterruptedException | ExecutionException e) {
                    e.printStackTrace();
                }
            });
        }

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

    private ScanResult scanTarget(String target, String ports) {
        PortRange portRange = new PortRange(ports);
        AtomicLong portSlotFactory = new AtomicLong(System.nanoTime());
        BitSet openPorts = new BitSet(PortRange.MAX_PORT);

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

        try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor();) {
            List<Future<Integer>> futures = new ArrayList<>();
            while (portRange.hasNext()) {
                maxWorkers.acquireUninterruptibly();
                final int currentPort = portRange.next();
                futures.add(executor.submit(() -> {
                    try {
                        waitForSlot(portSlotFactory);
                        int port = scanPort(target, currentPort);
                        if (port != -1) {
                            progress.open().incrementAndGet();
                        }
                        return port;
                    } finally {
                        maxWorkers.release();
                        progress.done().incrementAndGet();
                    }
                }));
            }

            futures.forEach(f -> {
                try {
                    int port = f.get();
                    if (port != -1) {
                        // openPorts is not threadsafe, so we do not update it outside the worker threads
                        openPorts.set(port);
                    }
                } catch (InterruptedException | ExecutionException e) {
                    e.printStackTrace();
                }
            });
        }

        return new ScanResult(target, openPorts);
    }

    private int scanPort(String target, final int currentPort) {
        if (isTargetPortOpen(target, currentPort)) {
            return currentPort;
        }

        return -1;
    }

    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 boolean isTargetPortOpen(String target, int port) {
        socketLimit.acquireUninterruptibly();
        counter.inc();
        try (Socket socket = new Socket()) {
            socket.connect(new InetSocketAddress(target, port), timeoutInMillis);
            return true;
        } catch (SocketException | SocketTimeoutException ignored) {
        } catch (IOException e) {
            System.out.println(e.getClass().getSimpleName() + ": " + e.getMessage());
        } finally {
            counter.dec();
            socketLimit.release();
        }

        return false;
    }
}