const http  = require('http');
const https = require('https');
const WebSocket = require('ws');
const fs   = require('fs');
const path = require('path');
const dotenv = require('dotenv');
const axios  = require('axios');

// --- Logging -----------------------------------------------------------

const LOG_FILE_PATH = path.join(__dirname, '..', 'logs', 'console_output.log');
const TZ_NAME = Intl.DateTimeFormat('en', { timeZoneName: 'short' })
    .formatToParts(new Date())
    .find(p => p.type === 'timeZoneName').value;

const LOG_MAX_BYTES = () => {
    const v = parseInt(process.env.LOG_MAX_BYTES || '', 10);
    return isNaN(v) ? 10 * 1024 * 1024 : v;
};

let _logStream = null;

function rotateLog() {
    const maxBytes = LOG_MAX_BYTES();
    let size = 0;
    try { size = fs.statSync(LOG_FILE_PATH).size; } catch (_) { return; }
    if (size < maxBytes) return;

    if (_logStream) { _logStream.end(); _logStream = null; }
    const rotated = LOG_FILE_PATH + '.1';
    try { fs.unlinkSync(rotated); } catch (_) {}
    try { fs.renameSync(LOG_FILE_PATH, rotated); } catch (_) {}
    openLogStream();
    logToFile(`Log rotated - previous log saved to ${rotated}`);
}

function openLogStream() {
    _logStream = fs.createWriteStream(LOG_FILE_PATH, { flags: 'a' });
    _logStream.on('error', (err) => process.stderr.write(`Log stream error: ${err.message}\n`));
}

const logToFile = (message) => {
    const timestamp = new Date().toISOString().replace('T', ' ').slice(0, 19);
    // Strip newlines so a forged payload cannot inject fake log lines.
    const logMessage = `[${timestamp} ${TZ_NAME}] ${message.replace(/[\r\n]/g, ' ')}\n`;
    if (_logStream) {
        _logStream.write(logMessage);
    } else {
        process.stderr.write(logMessage);
    }
};

// --- Env loading -------------------------------------------------------

function loadEnv() {
    const paths = [
        path.resolve(__dirname, '.env'),
        path.resolve(__dirname, '../.env'),
        path.resolve(__dirname, '../../.env'),
    ];
    for (const envPath of paths) {
        if (fs.existsSync(envPath)) {
            dotenv.config({ path: envPath });
            logToFile(`Loaded .env from: ${envPath}`);
            return;
        }
    }
    logToFile('No .env file found in specified paths');
}

// --- Config validation -------------------------------------------------

function validateConfig() {
    const required = ['HOST_NAME', 'API_TOKEN'];
    const missing = required.filter(k => !process.env[k]);
    if (missing.length) {
        logToFile(`FATAL: Missing required environment variables: ${missing.join(', ')}`);
        process.exit(1);
    }
    if (process.env.API_TLS_VERIFY?.toUpperCase() !== 'TRUE') {
        logToFile('WARNING: API_TLS_VERIFY is not TRUE - TLS certificate verification disabled for API calls. Set API_TLS_VERIFY=TRUE to enable verification.');
    }
    if (!process.env.PROXMOX_CA_CERT) {
        logToFile('WARNING: PROXMOX_CA_CERT not set - TLS certificate verification disabled (rejectUnauthorized: false). Set PROXMOX_CA_CERT=/path/to/ca.crt to enable verification.');
    }
    if (!process.env.PROXMOX_ALLOWED_HOSTS) {
        logToFile('WARNING: PROXMOX_ALLOWED_HOSTS not set - wssUrl hostname will not be validated against an allowlist. Set PROXMOX_ALLOWED_HOSTS=node1.example.com,node2.example.com to restrict allowed Proxmox hosts.');
    }
}

// --- Env-dependent constants -------------------------------------------
// Initialised after loadEnv() + validateConfig() in the Startup section.

let PROTOCOL;
let EXACT_ORIGIN;
let ALLOWED_ORIGINS;
let PROXMOX_ALLOWED_HOSTNAMES;

function initEnvConstants() {
    PROTOCOL = process.env.HTTPS_DISABLED?.toUpperCase() === 'TRUE' ? 'http' : 'https';

    const hostName = process.env.HOST_NAME || '';
    EXACT_ORIGIN = `${PROTOCOL}://${hostName}`;

    // Multi-portal deployments list additional allowed origins here.
    const allowedOriginsEnv = process.env.ALLOWED_ORIGINS || '';
    ALLOWED_ORIGINS = new Set(
        allowedOriginsEnv.split(',').map(o => o.trim()).filter(Boolean)
    );

    const allowedHostsEnv = process.env.PROXMOX_ALLOWED_HOSTS || '';
    PROXMOX_ALLOWED_HOSTNAMES = new Set(
        allowedHostsEnv.split(',').map(h => h.trim()).filter(Boolean)
    );
}

// --- TLS agents --------------------------------------------------------
// Two separate agents: the API agent defaults to rejectUnauthorized: false
// (set API_TLS_VERIFY=TRUE to enable strict verification). The Proxmox
// agent also defaults to rejectUnauthorized: false for self-signed certs
// (set PROXMOX_CA_CERT to enable CA-based verification).

let _apiHttpsAgent = null;
const getApiHttpsAgent = () => {
    if (!_apiHttpsAgent) {
        const rejectUnauthorized = process.env.API_TLS_VERIFY?.toUpperCase() === 'TRUE';
        _apiHttpsAgent = new https.Agent({ keepAlive: true, maxSockets: 10, rejectUnauthorized });
    }
    return _apiHttpsAgent;
};

// CA cert is read once at first connection - restart after cert rotation.
let _proxmoxCaAgent = null;
const getProxmoxCaAgent = () => {
    if (!_proxmoxCaAgent) {
        _proxmoxCaAgent = process.env.PROXMOX_CA_CERT
            ? new https.Agent({ ca: fs.readFileSync(process.env.PROXMOX_CA_CERT) })
            : new https.Agent({ rejectUnauthorized: false });
    }
    return _proxmoxCaAgent;
};

// --- Rate limiting -----------------------------------------------------

const connCountByIp = new Map();
const connRateByIp  = new Map();

const MAX_CONN_PER_IP  = 20;
const MAX_RATE_PER_IP  = 30;
const MAX_GLOBAL_CONN  = 100;

// Backpressure high-watermark: pause the source socket when the outbound
// send buffer exceeds this to prevent unbounded heap growth.
const RELAY_HWM = 1024 * 1024; // 1 MiB

setInterval(() => {
    const cutoff = Date.now() - 60000;
    for (const [ip, timestamps] of connRateByIp) {
        const fresh = timestamps.filter(t => t > cutoff);
        if (fresh.length === 0) connRateByIp.delete(ip);
        else connRateByIp.set(ip, fresh);
    }
}, 60000).unref();

// --- wssUrl validation (SSRF prevention) -------------------------------

function isPrivateHost(hostname) {
    let h = hostname.toLowerCase();

    // URL.hostname wraps IPv6 in brackets (e.g. '[::1]') - strip them.
    if (h.startsWith('[') && h.endsWith(']')) h = h.slice(1, -1);

    // IPv4-mapped IPv6: Node normalises [::ffff:10.0.0.1] to ::ffff:a00:1
    // (hex). Convert back to dotted-decimal so the IPv4 checks catch it.
    if (h.startsWith('::ffff:')) {
        const inner = h.slice(7);
        if (/^\d+\.\d+\.\d+\.\d+$/.test(inner)) {
            h = inner;
        } else {
            const parts = inner.split(':');
            if (parts.length === 2) {
                const hi = parseInt(parts[0], 16);
                const lo = parseInt(parts[1], 16);
                if (!isNaN(hi) && !isNaN(lo)) {
                    h = `${(hi >> 8) & 0xff}.${hi & 0xff}.${(lo >> 8) & 0xff}.${lo & 0xff}`;
                }
            }
        }
    }

    if (/^127\./.test(h) || /^10\./.test(h) || /^192\.168\./.test(h)) return true;
    if (/^172\.(1[6-9]|2\d|3[01])\./.test(h))                          return true;
    if (/^169\.254\./.test(h))                                           return true;
    if (h === '::1' || /^fe80:/i.test(h))                               return true;
    if (/^fc[0-9a-f]{2}:|^fd[0-9a-f]{2}:/i.test(h))                    return true;
    if (h === 'localhost' || h === 'metadata.google.internal')           return true;
    return false;
}

function validateWssUrl(wssUrl) {
    let parsed;
    try { parsed = new URL(wssUrl); }
    catch (e) { throw new Error(`Invalid wssUrl - not a valid URL: ${e.message}`); }

    if (parsed.protocol !== 'ws:' && parsed.protocol !== 'wss:') {
        throw new Error(`Invalid wssUrl scheme '${parsed.protocol}' - only ws: and wss: are permitted`);
    }

    if (PROXMOX_ALLOWED_HOSTNAMES.size > 0 && !PROXMOX_ALLOWED_HOSTNAMES.has(parsed.hostname)) {
        throw new Error(`wssUrl hostname '${parsed.hostname}' is not in PROXMOX_ALLOWED_HOSTS`);
    }
}

// --- Authentication ----------------------------------------------------

const authenticateAndObtainTicket = async (id) => {
    const apiUrl = `${PROTOCOL}://${process.env.HOST_NAME}/api-v1/virtual-machine/console-view?id=${id}&console=false`;

    const config = {
        headers: { 'x-api-key': process.env.API_TOKEN },
        httpsAgent: getApiHttpsAgent(),
        timeout: 10000,
    };

    logToFile('Calling console-view API: ' + apiUrl);
    const vncResponse = await axios.get(apiUrl, config);

    const vnc = vncResponse.data?.vncwebsocket;
    if (!vnc?.token || !vnc?.wssUrl) {
        throw new Error(`Unexpected API response shape - missing token or wssUrl`);
    }

    validateWssUrl(vnc.wssUrl);

    const logData = {
        ...vncResponse.data,
        vncwebsocket: {
            ...vnc,
            token:          '********',
            password:       '********',
            vnc_ticket:     '********',
            wssUrl:         '********',
            haproxy_cookie: '********',
        },
    };
    logToFile('Ticket obtained: ' + JSON.stringify(logData));

    // Strip CR/LF/semicolons to prevent CRLF injection and cookie injection
    // into the Proxmox WebSocket upgrade Cookie header.
    return {
        authToken:     (vnc.token          || '').replace(/[\r\n;]/g, ''),
        wssUrl:        vnc.wssUrl,
        haproxyCookie: (vnc.haproxy_cookie || '').replace(/[\r\n;]/g, ''),
    };
};

// --- WebSocket server --------------------------------------------------

const UUID_RE = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i;

// maxPayload 256 KiB - VNC frames are 1–64 KiB; the 100 MiB default would
// let unauthenticated clients exhaust memory before auth runs.
const server = http.createServer((req, res) => {
    if (req.method === 'GET' && req.url === '/healthz') {
        res.writeHead(200, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({
            status: 'ok',
            clients: wss.clients.size,
            maxClients: MAX_GLOBAL_CONN,
            uptimeSeconds: Math.floor(process.uptime()),
        }));
        return;
    }
    res.writeHead(426, { 'Upgrade': 'websocket' });
    res.end('WebSocket connections only');
});
const wss = new WebSocket.Server({ server, maxPayload: 262144 });

// --- Connection handler ------------------------------------------------

wss.on('connection', async (ws, req) => {

    // --- Rate limiting (before origin check and auth) ---
    // Only trust XFF when the direct peer is the local reverse proxy.
    // Take the rightmost XFF value (the one appended by the trusted proxy).
    const directPeer       = req.socket.remoteAddress;
    const isFromLocalProxy = directPeer === '127.0.0.1' || directPeer === '::1';
    const rawIp = isFromLocalProxy
        ? (req.headers['x-forwarded-for']?.split(',').pop().trim() || directPeer)
        : directPeer;
    const clientIp = rawIp.replace(/[^\x20-\x7E]/g, '_');

    if (wss.clients.size >= MAX_GLOBAL_CONN) {
        logToFile(`Rejected connection from ${clientIp}: server at capacity (${wss.clients.size})`);
        ws.close(4004, 'Server at capacity');
        return;
    }

    const currentCount = connCountByIp.get(clientIp) || 0;
    if (currentCount >= MAX_CONN_PER_IP) {
        logToFile(`Rejected connection from ${clientIp}: too many concurrent connections (${currentCount})`);
        ws.close(4004, 'Too many connections');
        return;
    }

    const now = Date.now();
    const recentTimestamps = (connRateByIp.get(clientIp) || []).filter(t => now - t < 60000);
    if (recentTimestamps.length >= MAX_RATE_PER_IP) {
        logToFile(`Rejected connection from ${clientIp}: rate limit exceeded (${recentTimestamps.length}/min)`);
        ws.close(4004, 'Rate limit exceeded');
        return;
    }
    recentTimestamps.push(now);
    connRateByIp.set(clientIp, recentTimestamps);
    connCountByIp.set(clientIp, currentCount + 1);

    let pingInterval = null;
    let proxmoxWs = null;

    let cleaned = false;
    const cleanup = () => {
        if (cleaned) return;
        cleaned = true;
        clearInterval(pingInterval);
        const count = connCountByIp.get(clientIp) || 1;
        if (count <= 1) connCountByIp.delete(clientIp);
        else connCountByIp.set(clientIp, count - 1);
        if (proxmoxWs) proxmoxWs.close();
    };

    ws.on('close', () => {
        logToFile('Client connection closed');
        cleanup();
    });

    // --- Origin validation (CSWSH prevention) ---
    const origin = req.headers.origin;
    if (origin) {
        const exactMatch = origin === EXACT_ORIGIN;
        const explicitMatch = ALLOWED_ORIGINS.size > 0 && ALLOWED_ORIGINS.has(origin);
        if (!exactMatch && !explicitMatch) {
            logToFile(`Rejected connection from disallowed origin: ${origin}`);
            ws.close(4003, 'Origin not allowed');
            return;
        }
    }

    // --- UUID input validation ---
    const id = new URL(req.url, 'http://localhost').searchParams.get('id');
    if (!id || !UUID_RE.test(id)) {
        logToFile(`Rejected connection with invalid VM ID: ${id}`);
        ws.close(4002, 'Invalid VM ID');
        return;
    }

    // --- Ping-pong keepalive ---
    ws.isAlive = true;
    ws.on('pong', () => { ws.isAlive = true; });

    pingInterval = setInterval(() => {
        if (!ws.isAlive) {
            logToFile(`Terminating unresponsive client connection from ${clientIp}`);
            clearInterval(pingInterval);
            ws.terminate();
            return;
        }
        ws.isAlive = false;
        ws.ping();
    }, 30000);

    let clientMessageLogged  = false;
    let proxmoxMessageLogged = false;

    ws.on('error', (err) => {
        logToFile(`Client WebSocket error from ${clientIp}: ${err.message}`);
    });

    try {
        logToFile('Authenticating and obtaining ticket...');
        const credentials = await authenticateAndObtainTicket(id);

        if (ws.readyState !== WebSocket.OPEN) {
            logToFile(`Client disconnected during auth from ${clientIp} - aborting`);
            clearInterval(pingInterval);
            return;
        }

        logToFile('Authentication successful');

        // Do NOT use encodeURIComponent - Proxmox VNC tickets contain ':',
        // '@', '=' which would be over-encoded, causing Proxmox to reject.
        const cookieParts = [];
        if (credentials.haproxyCookie) cookieParts.push(credentials.haproxyCookie);
        cookieParts.push(`PVEAuthCookie=${credentials.authToken}`);

        const wsOptions = {
            headers: { Cookie: cookieParts.join('; ') },
            agent:   getProxmoxCaAgent(),
        };

        logToFile('Creating Proxmox WebSocket');
        proxmoxWs = new WebSocket(credentials.wssUrl, wsOptions);

        if (ws.readyState !== WebSocket.OPEN) {
            logToFile(`Client disconnected before proxmoxWs created (${clientIp}) - terminating`);
            proxmoxWs.terminate();
            return;
        }

        // Buffer messages arriving while proxmoxWs is CONNECTING so VNC/RFB
        // negotiation frames aren't silently dropped under network latency.
        const messageBuffer = [];

        ws.on('message', (message) => {
            if (!clientMessageLogged) {
                logToFile('Received message from client');
                clientMessageLogged = true;
            }
            if (proxmoxWs.readyState === WebSocket.OPEN) {
                proxmoxWs.send(message, (err) => {
                    if (err) logToFile(`Error relaying to Proxmox: ${err.message}`);
                });
                if (proxmoxWs.bufferedAmount > RELAY_HWM) {
                    ws.pause();
                    const drainTimer = setTimeout(() => ws.resume(), 5000);
                    proxmoxWs.once('drain', () => { clearTimeout(drainTimer); ws.resume(); });
                }
            } else {
                messageBuffer.push(message);
            }
        });

        proxmoxWs.on('open', () => {
            logToFile('Proxmox WebSocket open - relay active');
            for (const msg of messageBuffer) {
                proxmoxWs.send(msg, (err) => {
                    if (err) logToFile(`Error flushing buffered message to Proxmox: ${err.message}`);
                });
            }
            messageBuffer.length = 0;
        });

        proxmoxWs.on('message', (message) => {
            if (!proxmoxMessageLogged) {
                logToFile('Received message from Proxmox VNC');
                proxmoxMessageLogged = true;
            }
            if (ws.readyState === WebSocket.OPEN) {
                ws.send(message, (err) => {
                    if (err) logToFile(`Error relaying to client: ${err.message}`);
                });
                if (ws.bufferedAmount > RELAY_HWM) {
                    proxmoxWs.pause();
                    const drainTimer = setTimeout(() => proxmoxWs.resume(), 5000);
                    ws.once('drain', () => { clearTimeout(drainTimer); proxmoxWs.resume(); });
                }
            }
        });

        proxmoxWs.on('close', () => {
            logToFile('Proxmox WebSocket closed');
            ws.close();
        });

        proxmoxWs.on('error', (error) => {
            logToFile(`Proxmox VNC WebSocket error: ${error.message}`);
        });

    } catch (error) {
        logToFile(`Error in WebSocket connection: ${error.message}`);
        clearInterval(pingInterval);
        ws.close(4001, 'Authentication failed');
    }
});

// --- Graceful shutdown -------------------------------------------------

function gracefulShutdown(signal) {
    logToFile(`Received ${signal} - shutting down gracefully`);
    server.close();
    for (const client of wss.clients) {
        client.close(1001, 'Server shutting down');
    }
    setTimeout(() => {
        if (_logStream) _logStream.end(() => process.exit(0));
        else process.exit(0);
    }, 5000).unref();
}
process.on('SIGTERM', () => gracefulShutdown('SIGTERM'));
process.on('SIGINT',  () => gracefulShutdown('SIGINT'));

process.on('unhandledRejection', (reason) => {
    logToFile(`Unhandled promise rejection: ${reason}`);
});

// --- Startup -----------------------------------------------------------

fs.mkdirSync(path.join(__dirname, '..', 'logs'), { recursive: true });

rotateLog();
if (!_logStream) openLogStream();

loadEnv();
validateConfig();
initEnvConstants();

// Warm up TLS agents so the first connection doesn't block on readFileSync.
getApiHttpsAgent();
getProxmoxCaAgent();

setInterval(rotateLog, 60 * 60 * 1000).unref();

const PORT = process.env.CONSOLE_PORT || 8099;

// Bind to localhost only - reverse proxy handles TLS and forwards here.
server.listen(PORT, '127.0.0.1', () => {
    logToFile(`Server running on 127.0.0.1:${PORT}`);
});
