import { EventEmitter } from 'events';
import os from 'os';
import { Worker } from 'worker_threads';
import path from 'path';
import { fileURLToPath } from 'url';

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

export class WebSocketScanner extends EventEmitter {
    constructor(options = {}) {
        super();
        this.scanning = false;
        this.workers = [];
        // Default to number of CPU cores, but allow override
        this.maxThreads = options.maxThreads || os.cpus().length;
        this.totalIPs = 0;
        this.scannedIPs = 0;
    }

    /**
     * Scans the local network for websocket services on a specified port
     * @param {number} port - Port number to scan
     * @returns {Promise<Array>} Array of IP addresses where websocket service was found
     */
    async scanNetwork(port, protocol) {
        // Clean up any existing workers before starting a new scan
        this.cleanupWorkers();
        console.log("Cleaned up workers");
        
        if (this.scanning) {
            throw new Error('Scan already in progress');
        }

        this.scanning = true;
        this.scannedIPs = 0;
        const foundServices = [];
        
        // Get local network range
        const { startIP, endIP } = this.getLocalNetworkRange();
        const start = this.ipToNumber(startIP);
        const end = this.ipToNumber(endIP);

        // Calculate IP ranges for each worker
        this.totalIPs = end - start + 1;
        const ipsPerWorker = Math.ceil(this.totalIPs / this.maxThreads);
        const workerPromises = [];

        for (let i = 0; i < this.maxThreads; i++) {
            const workerStart = start + (i * ipsPerWorker);
            const workerEnd = Math.min(workerStart + ipsPerWorker - 1, end);
            
            if (workerStart > end) break;

            const workerStartIP = this.numberToIP(workerStart);
            const workerEndIP = this.numberToIP(workerEnd);

            const worker = new Worker(path.join(__dirname, 'websocketScannerWorker.js'));
            this.workers.push(worker);
            console.log("Created worker", i);

            const workerPromise = new Promise((resolve) => {
                worker.on('message', (message) => {
                    switch (message.type) {
                        case 'serviceFound':
                            foundServices.push({ ip: message.ip, hostname: message.hostname });
                            this.emit('serviceFound', { ip: message.ip, hostname: message.hostname });
                            break;
                        case 'scanProgress':
                            this.scannedIPs += message.increment;
                            const totalProgress = (this.scannedIPs / this.totalIPs) * 100;
                            this.emit('scanProgress', {
                                currentIP: message.currentIP,
                                progress: totalProgress
                            });
                            break;
                        case 'scanComplete':
                            resolve(message.results);
                            break;
                    }
                });
            });

            worker.postMessage({
                type: 'scan',
                startIP: workerStartIP,
                endIP: workerEndIP,
                port,
                protocol
            });

            workerPromises.push(workerPromise);
        }

        await Promise.all(workerPromises);
        this.cleanupWorkers();
        this.scanning = false;
        return foundServices;
    }

    cleanupWorkers() {
        this.workers.forEach(worker => worker.terminate());
        this.workers = [];
    }

    /**
     * Gets the local network IP range
     * @returns {Object} Object containing startIP and endIP
     */
    getLocalNetworkRange() {
        const interfaces = os.networkInterfaces();
        let localIP = null;
        let subnetMask = null;

        // Find the first non-internal IPv4 address
        for (const name of Object.keys(interfaces)) {
            for (const iface of interfaces[name]) {
                if (iface.family === 'IPv4' && !iface.internal) {
                    localIP = iface.address;
                    subnetMask = iface.netmask;
                    break;
                }
            }
            if (localIP) break;
        }

        if (!localIP) {
            throw new Error('Could not determine local network IP address');
        }

        // Convert IP and subnet mask to numbers
        const ipNum = this.ipToNumber(localIP);
        const maskNum = this.ipToNumber(subnetMask);

        // Calculate network address
        const networkNum = ipNum & maskNum;

        // Calculate broadcast address
        const broadcastNum = networkNum | (~maskNum >>> 0);

        // Start IP is network address + 1
        const startIP = this.numberToIP(networkNum + 1);
        // End IP is broadcast address - 1
        const endIP = this.numberToIP(broadcastNum - 1);

        return { startIP, endIP };
    }

    /**
     * Converts an IP address to a number
     * @param {string} ip - IP address to convert
     * @returns {number} Numeric representation of IP
     */
    ipToNumber(ip) {
        return ip.split('.')
            .reduce((acc, octet) => (acc << 8) + parseInt(octet), 0) >>> 0;
    }

    /**
     * Converts a number to an IP address
     * @param {number} num - Number to convert
     * @returns {string} IP address
     */
    numberToIP(num) {
        return [
            (num >>> 24) & 255,
            (num >>> 16) & 255,
            (num >>> 8) & 255,
            num & 255
        ].join('.');
    }

    /**
     * Stops the current scan
     */
    stopScan() {
        this.scanning = false;
        this.cleanupWorkers();
    }
}



// To stop scanning at any time:
// scanner.stopScan();