const WebSocket = require('ws');
const jwt = require('jsonwebtoken');
const { getDatabase } = require('../config/database');

// agentId -> WebSocket
const agentSockets = new Map();
// agentId -> WebSocket
const browserSockets = new Map();

// RDP: separate Maps damit Shell nicht kollidiert
const rdpAgentSockets = new Map();
const rdpBrowserSockets = new Map();

function setupWebSocketServer(httpServer) {
    const wss = new WebSocket.Server({ server: httpServer, path: '/ws' });

    wss.on('connection', (ws, req) => {
        let url;
        try {
            url = new URL(req.url, 'http://localhost');
        } catch {
            ws.close(1008, 'invalid url');
            return;
        }
        const type = url.searchParams.get('type');

        if (type === 'agent') {
            handleAgent(ws, url);
        } else if (type === 'shell') {
            handleBrowser(ws, url);
        } else if (type === 'rdp-agent') {
            handleRdpAgent(ws, url);
        } else if (type === 'rdp') {
            handleRdpBrowser(ws, url);
        } else {
            ws.close(1008, 'unknown type');
        }
    });

    console.log('[WS] WebSocket server ready at /ws');
    return wss;
}

function handleAgent(ws, url) {
    const key = url.searchParams.get('key');
    if (key !== process.env.AGENT_API_KEY) {
        ws.close(1008, 'unauthorized');
        return;
    }
    const hostname = url.searchParams.get('hostname');
    if (!hostname) { ws.close(1008, 'hostname required'); return; }

    const db = getDatabase();
    const agent = db.prepare('SELECT id FROM monitoring_agents WHERE hostname = ?').get(hostname);
    if (!agent) { ws.close(1008, 'agent not found'); return; }

    const agentId = agent.id;
    // Alte Verbindung schließen falls vorhanden
    const oldWs = agentSockets.get(agentId);
    if (oldWs && oldWs.readyState === WebSocket.OPEN) oldWs.close();

    agentSockets.set(agentId, ws);
    console.log(`[WS] Agent verbunden: ${hostname} (id=${agentId})`);

    // Falls Browser bereits wartet → Shell starten + melden
    const bws = browserSockets.get(agentId);
    if (bws?.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify({ type: 'start_shell' }));
        bws.send('\r\n\x1b[32m[Agent verbunden — Shell gestartet]\x1b[0m\r\n');
    }

    ws.on('message', (data) => {
        const bws = browserSockets.get(agentId);
        if (!bws || bws.readyState !== WebSocket.OPEN) return;
        // RTC-Signaling (offer/answer/ice) → transparent weiterleiten
        try {
            const msg = JSON.parse(data.toString());
            if (msg.type?.startsWith('rtc_')) { bws.send(data.toString()); return; }
        } catch { /* kein JSON → Shell-Output */ }
        bws.send(data.toString());
    });

    ws.on('close', () => {
        agentSockets.delete(agentId);
        const bws = browserSockets.get(agentId);
        if (bws?.readyState === WebSocket.OPEN) {
            bws.send('\r\n\x1b[31m[Agent getrennt]\x1b[0m\r\n');
        }
        console.log(`[WS] Agent getrennt: ${hostname}`);
    });
}

function handleBrowser(ws, url) {
    const token = url.searchParams.get('token');
    try {
        jwt.verify(token, process.env.JWT_SECRET);
    } catch {
        ws.close(1008, 'unauthorized');
        return;
    }

    const agentId = parseInt(url.searchParams.get('agentId'));
    if (!agentId) { ws.close(1008, 'agentId required'); return; }

    browserSockets.set(agentId, ws);

    const aws = agentSockets.get(agentId);
    if (aws?.readyState === WebSocket.OPEN) {
        aws.send(JSON.stringify({ type: 'start_shell' }));
        ws.send('\x1b[32m[Verbunden]\x1b[0m\r\n');
    } else {
        ws.send('\x1b[33m[Warte auf Agent-Verbindung...]\x1b[0m\r\n');
    }

    ws.on('message', (data) => {
        const aws = agentSockets.get(agentId);
        if (!aws || aws.readyState !== WebSocket.OPEN) return;
        // RTC-Signaling (offer/answer/ice) → transparent weiterleiten
        try {
            const msg = JSON.parse(data.toString());
            if (msg.type?.startsWith('rtc_')) { aws.send(data.toString()); return; }
        } catch { /* kein JSON → Tastatureingabe */ }
        aws.send(JSON.stringify({ type: 'input', data: data.toString() }));
    });

    ws.on('close', () => {
        browserSockets.delete(agentId);
        const aws = agentSockets.get(agentId);
        if (aws?.readyState === WebSocket.OPEN) {
            aws.send(JSON.stringify({ type: 'stop_shell' }));
        }
    });
}

function handleRdpAgent(ws, url) {
    const key = url.searchParams.get('key');
    if (key !== process.env.AGENT_API_KEY) { ws.close(1008, 'unauthorized'); return; }
    const hostname = url.searchParams.get('hostname');
    if (!hostname) { ws.close(1008, 'hostname required'); return; }

    const db = getDatabase();
    const agent = db.prepare('SELECT id FROM monitoring_agents WHERE hostname = ?').get(hostname);
    if (!agent) { ws.close(1008, 'agent not found'); return; }

    const agentId = agent.id;
    const oldWs = rdpAgentSockets.get(agentId);
    if (oldWs?.readyState === WebSocket.OPEN) oldWs.close();
    rdpAgentSockets.set(agentId, ws);
    console.log(`[WS-RDP] Agent verbunden: ${hostname}`);

    ws.on('message', (data) => {
        const bws = rdpBrowserSockets.get(agentId);
        if (bws?.readyState === WebSocket.OPEN) bws.send(data.toString());
    });

    ws.on('close', () => {
        rdpAgentSockets.delete(agentId);
        const bws = rdpBrowserSockets.get(agentId);
        if (bws?.readyState === WebSocket.OPEN) {
            bws.send(JSON.stringify({ type: 'rdp_disconnected' }));
        }
        console.log(`[WS-RDP] Agent getrennt: ${hostname}`);
    });
}

function handleRdpBrowser(ws, url) {
    const token = url.searchParams.get('token');
    try { jwt.verify(token, process.env.JWT_SECRET); }
    catch { ws.close(1008, 'unauthorized'); return; }

    const agentId = parseInt(url.searchParams.get('agentId'));
    if (!agentId) { ws.close(1008, 'agentId required'); return; }

    const oldBws = rdpBrowserSockets.get(agentId);
    if (oldBws?.readyState === WebSocket.OPEN) oldBws.close();
    rdpBrowserSockets.set(agentId, ws);

    ws.on('message', (data) => {
        const aws = rdpAgentSockets.get(agentId);
        if (aws?.readyState === WebSocket.OPEN) aws.send(data.toString());
    });

    ws.on('close', () => {
        rdpBrowserSockets.delete(agentId);
        const aws = rdpAgentSockets.get(agentId);
        if (aws?.readyState === WebSocket.OPEN) {
            aws.send(JSON.stringify({ type: 'rdp_stop' }));
        }
    });
}

// Ankündigung an Agent pushen (für sofortige Zustellung)
function pushAnnouncementToAgent(agentId, announcements) {
    const aws = agentSockets.get(agentId);
    if (aws?.readyState === WebSocket.OPEN) {
        aws.send(JSON.stringify({ type: 'announcements', announcements }));
        return true;
    }
    return false;
}

// Ankündigung an alle verbundenen Agents pushen
function broadcastAnnouncement(announcement) {
    let pushed = 0;
    for (const [agentId, aws] of agentSockets.entries()) {
        if (aws.readyState === WebSocket.OPEN) {
            aws.send(JSON.stringify({ type: 'announcements', announcements: [announcement] }));
            pushed++;
        }
    }
    return pushed;
}

function getConnectedAgents() {
    const ids = [];
    for (const [agentId, ws] of agentSockets.entries()) {
        if (ws.readyState === WebSocket.OPEN) ids.push(agentId);
    }
    return ids;
}

module.exports = { setupWebSocketServer, pushAnnouncementToAgent, broadcastAnnouncement, getConnectedAgents };
