const express = require('express');
const http = require('http');
const WebSocket = require('ws');
const pty = require('node-pty');
const path = require('path');
const os = require('os');
const url = require('url');
const { execSync } = require('child_process');

const app = express();
const server = http.createServer(app);
const wss = new WebSocket.Server({ noServer: true });

const PORT = process.env.PORT || 7860;
const PASSWORD = process.env.PASS || 'admin';

// Resource limits
const MAX_CPU_PERCENT = 95;
const MAX_MEM_PERCENT = 95;
const TOTAL_MEM_MB = Math.round(os.totalmem() / 1024 / 1024);
const MAX_MEM_MB = Math.round(TOTAL_MEM_MB * MAX_MEM_PERCENT / 100);

console.log(`Password: ${PASSWORD ? 'SET' : 'DEFAULT'}`);
console.log(`Limits: CPU ${MAX_CPU_PERCENT}%, Memory ${MAX_MEM_MB}MB / ${TOTAL_MEM_MB}MB`);

app.use(express.json());
app.use(express.static(path.join(__dirname, 'public')));

// Track terminal PIDs for resource monitoring
const terminals = {};
const terminalPids = new Set();

function getOrCreateTerminal(sessionId = 'main') {
    if (terminals[sessionId] && !terminals[sessionId].killed) {
        return terminals[sessionId];
    }

    console.log(`Creating terminal: ${sessionId}`);
    const term = pty.spawn('/bin/bash', ['--login'], {
        name: 'xterm-256color',
        cols: 80,
        rows: 24,
        cwd: '/root',
        env: {
            ...process.env,
            TERM: 'xterm-256color',
            COLORTERM: 'truecolor',
            LANG: 'C.UTF-8',
            HOME: '/root',
            SHELL: '/bin/bash'
        }
    });

    // Track PID
    terminalPids.add(term.pid);

    term.outputBuffer = '';
    term.maxBuffer = 50000;
    term.clients = new Set();
    term.killed = false;

    term.onData((data) => {
        term.outputBuffer += data;
        if (term.outputBuffer.length > term.maxBuffer) {
            term.outputBuffer = term.outputBuffer.slice(-term.maxBuffer);
        }
        term.clients.forEach(ws => {
            if (ws.readyState === WebSocket.OPEN) {
                try { ws.send(JSON.stringify({ type: 'output', data })); } catch {}
            }
        });
    });

    term.onExit(() => {
        console.log(`Terminal ${sessionId} exited`);
        term.killed = true;
        terminalPids.delete(term.pid);
        const oldClients = new Set(term.clients);
        setTimeout(() => {
            delete terminals[sessionId];
            const newTerm = getOrCreateTerminal(sessionId);
            oldClients.forEach(ws => {
                if (ws.readyState === WebSocket.OPEN) {
                    newTerm.clients.add(ws);
                    ws._term = newTerm;
                    ws.send(JSON.stringify({ type: 'output', data: '\r\n\x1b[33m[Session restarted]\x1b[0m\r\n' }));
                }
            });
        }, 500);
    });

    terminals[sessionId] = term;
    return term;
}

// Pre-create terminal
getOrCreateTerminal('main');

function validToken(token) {
    return token && token.trim() === PASSWORD.trim();
}

app.post('/api/auth', (req, res) => {
    const { password } = req.body;
    if (password && password.trim() === PASSWORD.trim()) {
        return res.json({ success: true, token: PASSWORD.trim() });
    }
    res.status(401).json({ error: 'Invalid password' });
});

app.get('/api/auth/check', (req, res) => {
    res.json({ authenticated: validToken(req.query.token) });
});

// Get resource usage for terminal processes only
function getTerminalResourceUsage() {
    let cpuPercent = 0;
    let memMB = 0;
    
    try {
        // Get Node.js process memory
        const nodeMemMB = Math.round(process.memoryUsage().rss / 1024 / 1024);
        memMB += nodeMemMB;
        
        // Get child process stats
        terminalPids.forEach(pid => {
            try {
                // Get process stats
                const stat = execSync(`ps -p ${pid} -o %cpu,%mem --no-headers 2>/dev/null || true`, { encoding: 'utf8', timeout: 1000 });
                const parts = stat.trim().split(/\s+/);
                if (parts.length >= 2) {
                    cpuPercent += parseFloat(parts[0]) || 0;
                    const memPercent = parseFloat(parts[1]) || 0;
                    memMB += Math.round((memPercent / 100) * TOTAL_MEM_MB);
                }
                
                // Also get child processes
                const children = execSync(`pgrep -P ${pid} 2>/dev/null || true`, { encoding: 'utf8', timeout: 1000 });
                children.trim().split('\n').filter(p => p).forEach(childPid => {
                    try {
                        const childStat = execSync(`ps -p ${childPid} -o %cpu,%mem --no-headers 2>/dev/null || true`, { encoding: 'utf8', timeout: 1000 });
                        const childParts = childStat.trim().split(/\s+/);
                        if (childParts.length >= 2) {
                            cpuPercent += parseFloat(childParts[0]) || 0;
                            const childMemPercent = parseFloat(childParts[1]) || 0;
                            memMB += Math.round((childMemPercent / 100) * TOTAL_MEM_MB);
                        }
                    } catch {}
                });
            } catch {}
        });
    } catch (e) {
        // Fallback to basic stats
        cpuPercent = Math.round(os.loadavg()[0] * 50); // Rough estimate
        memMB = Math.round(process.memoryUsage().rss / 1024 / 1024);
    }
    
    // Cap at limits
    cpuPercent = Math.min(Math.round(cpuPercent), MAX_CPU_PERCENT);
    const memPercent = Math.min(Math.round((memMB / MAX_MEM_MB) * 100), 100);
    
    return {
        cpu: {
            percent: cpuPercent,
            limit: MAX_CPU_PERCENT
        },
        memory: {
            used: memMB,
            limit: MAX_MEM_MB,
            total: TOTAL_MEM_MB,
            percent: memPercent
        }
    };
}

// Check and enforce resource limits
function enforceResourceLimits() {
    const usage = getTerminalResourceUsage();
    
    // If approaching limits, try to free resources
    if (usage.memory.percent > 90) {
        console.log(`[Enforcer] Memory at ${usage.memory.percent}% - triggering GC`);
        if (global.gc) {
            try { global.gc(); } catch {}
        }
        
        // Clear old terminal buffers
        Object.values(terminals).forEach(term => {
            if (term.outputBuffer.length > 10000) {
                term.outputBuffer = term.outputBuffer.slice(-10000);
            }
        });
    }
    
    return usage;
}

app.get('/api/stats', (req, res) => {
    if (!validToken(req.query.token)) return res.status(401).json({});
    res.json(enforceResourceLimits());
});

server.on('upgrade', (req, socket, head) => {
    const parsed = url.parse(req.url, true);
    if (!validToken(parsed.query.token)) {
        socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
        socket.destroy();
        return;
    }
    wss.handleUpgrade(req, socket, head, ws => wss.emit('connection', ws, req));
});

wss.on('connection', (ws, req) => {
    const sessionId = url.parse(req.url, true).query.session || 'main';
    const term = getOrCreateTerminal(sessionId);

    term.clients.add(ws);
    ws._term = term;
    ws.isAlive = true;

    if (term.outputBuffer) {
        ws.send(JSON.stringify({ type: 'output', data: term.outputBuffer }));
    }
    ws.send(JSON.stringify({ type: 'connected', session: sessionId, cols: term.cols, rows: term.rows }));

    ws.on('message', raw => {
        try {
            const msg = JSON.parse(raw.toString());
            if (msg.type === 'input' && ws._term && !ws._term.killed) {
                ws._term.write(msg.data);
            } else if (msg.type === 'resize' && ws._term && !ws._term.killed) {
                try {
                    ws._term.resize(
                        Math.max(2, Math.min(500, +msg.cols || 80)),
                        Math.max(2, Math.min(200, +msg.rows || 24))
                    );
                } catch {}
            } else if (msg.type === 'ping') {
                ws.isAlive = true;
                ws.send(JSON.stringify({ type: 'pong' }));
            }
        } catch {}
    });

    ws.on('close', () => { if (ws._term) ws._term.clients.delete(ws); });
    ws.on('error', () => { if (ws._term) ws._term.clients.delete(ws); });
    ws.on('pong', () => { ws.isAlive = true; });
});

// Heartbeat + resource check
setInterval(() => {
    wss.clients.forEach(ws => {
        if (!ws.isAlive) {
            if (ws._term) ws._term.clients.delete(ws);
            return ws.terminate();
        }
        ws.isAlive = false;
        ws.ping();
    });
    
    // Periodic resource enforcement
    enforceResourceLimits();
}, 30000);

// More frequent resource check
setInterval(enforceResourceLimits, 10000);

app.get('/', (req, res) => res.sendFile(path.join(__dirname, 'public', 'index.html')));

server.listen(PORT, '0.0.0.0', () => console.log(`Server running on port ${PORT}`));

process.on('SIGTERM', () => {
    Object.values(terminals).forEach(t => { if (!t.killed) t.kill(); });
    process.exit(0);
});

process.on('uncaughtException', (err) => {
    console.error('Uncaught exception:', err.message);
});

process.on('unhandledRejection', (err) => {
    console.error('Unhandled rejection:', err);
});