const cluster = require('cluster');
const numCPUs = require('os').cpus().length;
const net = require('net');

if (cluster.isMaster) {
    console.log(`Master ${process.pid} is running`);
    console.log(`Total Workers ${numCPUs * 6}`);
    for (let i = 0; i < numCPUs * 4; i++) {
        cluster.fork();
    }

    cluster.on('exit', (worker, code, signal) => {
        console.log(`Worker ${worker.process.pid} died`);
    });
} else {
    const fs = require('fs');
    const http = require('http');
    const httpProxy = require('http-proxy');
    const HyperDHTServer = require('hyperdht');
    const b32 = require("hi-base32");
    const agent = new http.Agent({ maxSockets: Number.MAX_VALUE });
    const content = fs.readFileSync('404.txt', 'utf8');
    const DEBUG = 0;
    const CONINFO = 0;
    const dhtServer = new HyperDHTServer();

    const startServer = async () => {
        console.log(`Worker ${process.pid} started`);
        await dhtServer.ready();

        const proxy = httpProxy.createProxyServer({
            ws: true,
            agent: agent,
            timeout: 360000
        });

        const server = http.createServer(async function (req, res) {
            try {
                const split = req.headers.host.split('.');
                const publicKey = Buffer.from(b32.decode.asBytes(split[0].toUpperCase()));

                if (publicKey.length < 32) {
                    console.log("Invalid Connection!");
                    res.writeHead(418, { 'Content-Type': 'text/html' });
                    res.end(content);
                    return;
                }

                // Create a unique tunnel for each request
                const port = await getAvailablePort();
                const tunnelServer = net.createServer(function (servsock) {
                    const socket = dhtServer.connect(publicKey);
                    let open = { local: true, remote: true };

                    servsock.on('data', (d) => socket.write(d));
                    socket.on('data', (d) => servsock.write(d));

                    const remoteend = () => {
                        if (open.remote) socket.end();
                        open.remote = false;
                    };
                    const localend = () => {
                        if (open.local) servsock.end();
                        open.local = false;
                    };

                    servsock.on('error', remoteend);
                    servsock.on('finish', remoteend);
                    servsock.on('end', remoteend);
                    socket.on('finish', localend);
                    socket.on('error', localend);
                    socket.on('end', localend);
                });

                tunnelServer.listen(port, "127.0.0.1", () => {
                    if (DEBUG === 1 || CONINFO === 1) console.log(`Tunnel server listening on port ${port}`);

                    proxy.web(req, res, {
                        target: 'http://127.0.0.1:' + port
                    }, function (e) {
                        console.log("Proxy Web Error: ", e);
                        res.writeHead(404, { 'Content-Type': 'text/html' });
                        res.end(content);
                    });

                    // Close tunnel and release resources after response is sent
                    res.on('finish', () => {
                        tunnelServer.close(() => {
                            if (DEBUG === 1 || CONINFO === 1) console.log("Tunnel closed after request completion.");
                        });
                    });
                });
            } catch (e) {
                console.error("Error Occurred: ", e);
            }
        });

        // Add WebSocket upgrade handling
        server.on('upgrade', async function (req, socket, head) {
            try {
                const split = req.headers.host.split('.');
                const publicKey = Buffer.from(b32.decode.asBytes(split[0].toUpperCase()));

                if (publicKey.length < 32) {
                    console.log("Invalid Connection!");
                    socket.destroy();
                    return;
                }

                // Create a unique tunnel for each WebSocket connection
                const port = await getAvailablePort();
                const tunnelServer = net.createServer(function (servsock) {
                    const socketConn = dhtServer.connect(publicKey);
                    let open = { local: true, remote: true };

                    servsock.on('data', (d) => socketConn.write(d));
                    socketConn.on('data', (d) => servsock.write(d));

                    const remoteend = () => {
                        if (open.remote) socketConn.end();
                        open.remote = false;
                    };
                    const localend = () => {
                        if (open.local) servsock.end();
                        open.local = false;
                    };

                    servsock.on('error', remoteend);
                    servsock.on('finish', remoteend);
                    servsock.on('end', remoteend);
                    socketConn.on('finish', localend);
                    socketConn.on('error', localend);
                    socketConn.on('end', localend);
                });

                tunnelServer.listen(port, "127.0.0.1", () => {
                    if (DEBUG === 1 || CONINFO === 1) console.log(`WebSocket tunnel server listening on port ${port}`);

                    proxy.ws(req, socket, head, {
                        target: 'http://127.0.0.1:' + port
                    }, function (e) {
                        console.log("Proxy WS Error: ", e);
                        socket.destroy();
                    });

                    // Clean up tunnel after WebSocket disconnects
                    socket.on('close', () => {
                        tunnelServer.close(() => {
                            if (DEBUG === 1 || CONINFO === 1) console.log("WebSocket tunnel closed after disconnect.");
                        });
                    });
                });
            } catch (e) {
                console.error("Error Occurred: ", e);
                socket.destroy();
            }
        });

        server.listen(8081, () => {
            console.log(`Worker ${process.pid} listening on port 8081`);
        });
    };

    startServer().catch(console.error);

    async function getAvailablePort() {
        return new Promise((resolve, reject) => {
            const tryPort = () => {
                const port = 1337 + Math.floor(Math.random() * 1000);
                const tester = net.createServer()
                    .once('error', (err) => {
                        if (err.code === 'EADDRINUSE') {
                            tryPort(); // Port in use, try another
                        } else {
                            reject(err);
                        }
                    })
                    .once('listening', () => {
                        tester.close(() => resolve(port)); // Port is available
                    })
                    .listen(port, '127.0.0.1');
            };
            tryPort();
        });
    }
}