import { WebSocket } from 'ws';
import { randomUUID } from 'crypto';
import { createRequire } from 'module';
import { sessions } from './session-manager.js';
import { auth } from './auth-manager.js';
import { MessageRelay } from './message-relay.js';
import { generateSessionKey, encodeKey, parseMessage, encryptAndSend } from './encryption.js';
import { verifyEntraToken, jose } from './entra.js';
const require = createRequire(import.meta.url);
const serverPkg = require('../package.json');
const clientRelay = new MessageRelay();
export function handleWebConnection(ws, req) {
    const url = new URL(req.url || '/', `http://${req.headers.host}`);
    const sessionId = url.searchParams.get('sessionId');
    const authToken = url.searchParams.get('authToken');
    const clientId = randomUUID();
    if (!sessionId) {
        try {
            ws.send(JSON.stringify({ type: 'error', message: 'Missing sessionId' }));
        }
        catch { /* swallow */ }
        ws.close();
        return;
    }
    // Check if agent exists for this session
    const agent = sessions.getAgentBySession(sessionId);
    // Password-protected session?
    const requiresAuth = auth.requiresAuth(sessionId);
    if (requiresAuth) {
        // Check saved auth token first
        if (authToken && auth.verifyAuthToken(authToken, sessionId)) {
            const refreshedToken = auth.generateAuthToken(sessionId);
            completeConnection(ws, clientId, sessionId, agent, refreshedToken);
            return;
        }
        // Check lockout
        if (auth.isLocked(sessionId)) {
            ws.send(JSON.stringify({
                type: 'auth_locked',
                message: 'Too many failed attempts. Try again in 15 minutes.',
            }));
            ws.close();
            return;
        }
        // Require authentication
        auth.setPending(clientId, sessionId);
        try {
            ws.send(JSON.stringify({ type: 'auth_required', sessionId }));
        }
        catch { /* swallow */ }
        ws.on('message', (data) => {
            handlePendingAuthMessage(clientId, ws, data.toString());
        });
        ws.on('close', () => {
            auth.removePending(clientId);
        });
        ws.on('error', (err) => {
            console.error(`[Web] WebSocket error during auth for ${clientId.slice(0, 8)}:`, err.message);
        });
        return;
    }
    // No auth required — proceed directly
    completeConnection(ws, clientId, sessionId, agent);
}
function handlePendingAuthMessage(clientId, ws, raw) {
    const sessionId = auth.getPending(clientId);
    if (!sessionId)
        return;
    let msg;
    try {
        msg = JSON.parse(raw);
    }
    catch {
        return;
    }
    if (msg.type !== 'authenticate' || typeof msg.password !== 'string')
        return;
    const agent = sessions.getAgentBySession(sessionId);
    const sessionAuth = auth.getSessionAuth(sessionId);
    if (!sessionAuth?.passwordHash || !sessionAuth?.passwordSalt) {
        // Password removed while authenticating
        ws.send(JSON.stringify({ type: 'error', message: 'Session no longer available.' }));
        auth.removePending(clientId);
        ws.close();
        return;
    }
    // Check lockout (may have been triggered by another client)
    if (auth.isLocked(sessionId)) {
        ws.send(JSON.stringify({
            type: 'auth_locked',
            message: 'Too many failed attempts. Try again in 15 minutes.',
        }));
        auth.removePending(clientId);
        ws.close();
        return;
    }
    const valid = auth.verifyPassword(msg.password, sessionAuth.passwordHash, sessionAuth.passwordSalt);
    if (!valid) {
        const { locked, remaining } = auth.recordFailure(sessionId);
        if (locked) {
            ws.send(JSON.stringify({
                type: 'auth_locked',
                message: 'Too many failed attempts. Try again in 15 minutes.',
            }));
            auth.removePending(clientId);
            ws.close();
        }
        else {
            ws.send(JSON.stringify({
                type: 'auth_failed',
                message: 'Incorrect password.',
                attemptsRemaining: remaining,
            }));
        }
        return;
    }
    // Success
    auth.clearFailures(sessionId);
    auth.removePending(clientId);
    const token = auth.generateAuthToken(sessionId);
    // Replace pending auth message handler with normal encrypted handler
    ws.removeAllListeners('message');
    ws.removeAllListeners('close');
    completeConnection(ws, clientId, sessionId, agent, token);
}
function completeConnection(ws, clientId, sessionId, agent, authToken) {
    const sessionKey = generateSessionKey();
    const client = {
        ws,
        clientId,
        sessionId,
        sessionKey,
        connectedAt: new Date(),
        isAlive: true,
    };
    sessions.registerClient(clientId, client);
    // Build connected payload — hide agent details for Entra-protected sessions
    // until the web client passes identity verification
    // Entra protection is active only when ownerOid has been established (agent verified)
    // agent.entra alone just means the agent intends to use Entra — not yet verified
    const sessionOwner = sessions.sessionOwners.get(sessionId);
    const isEntraProtected = !!(agent?.ownerOid || sessionOwner);
    const agentInfo = agent && !isEntraProtected ? {
        agentId: agent.agentId,
        name: agent.name,
        hostname: agent.hostname,
        workDir: agent.workDir,
        version: agent.version,
    } : null;
    const payload = {
        type: 'connected',
        clientId,
        sessionKey: encodeKey(sessionKey),
        serverVersion: serverPkg.version,
        agent: agentInfo,
        entraRequired: isEntraProtected,
        latestAgentVersion: serverPkg.latestAgentVersion || '',
        msInternal: agent?.msInternal || false,
    };
    if (authToken) {
        payload.authToken = authToken;
    }
    try {
        ws.send(JSON.stringify(payload));
    }
    catch (err) {
        console.error(`[Web] Failed to send connected payload to ${clientId.slice(0, 8)}:`, err.message);
    }
    console.log(`[Web] Client ${clientId.slice(0, 8)} connected to session ${sessionId}, agent: ${agent ? agent.name : 'none'}${authToken ? ' (authenticated)' : ''}`);
    // If session is Entra-protected, tell web client it must verify identity
    if (isEntraProtected) {
        encryptAndSend(ws, { type: 'entra_auth_required' }, sessionKey).catch(() => { });
    }
    ws.on('message', (data) => {
        clientRelay.enqueue(clientId, () => handleWebMessage(clientId, data.toString()));
    });
    ws.on('close', () => {
        console.log(`[Web] Client ${clientId.slice(0, 8)} disconnected`);
        clientRelay.cleanup(clientId);
        sessions.removeClient(clientId);
    });
    ws.on('pong', () => {
        client.isAlive = true;
    });
    ws.on('error', (err) => {
        console.error(`[Web] WebSocket error for ${clientId.slice(0, 8)}:`, err.message);
    });
}
async function handleWebMessage(clientId, raw) {
    const client = sessions.getClient(clientId);
    if (!client)
        return;
    const msg = await parseMessage(raw, client.sessionKey);
    if (!msg) {
        console.error(`[Web] Failed to parse/decrypt message from ${clientId.slice(0, 8)} (key=${client.sessionKey ? 'set' : 'null'})`);
        return;
    }
    if (msg.type !== 'ping') {
        console.log(`[Web] ${clientId.slice(0, 8)} → ${msg.type}`);
    }
    // Find the agent for this session
    const agent = sessions.getAgentBySession(client.sessionId);
    // Handle entra_auth from web client
    if (msg.type === 'entra_auth' && typeof msg.token === 'string') {
        const sessionOwnerForAuth = sessions.sessionOwners.get(client.sessionId);
        if (!agent && !sessionOwnerForAuth) {
            // No agent connected and no prior Entra owner — don't pre-verify, wait for agent
            await encryptAndSend(client.ws, {
                type: 'entra_auth_failed',
                reason: 'agent_pending',
            }, client.sessionKey);
            return;
        }
        // Determine the ownerOid to check against (session-level or agent-level)
        const ownerOid = sessionOwnerForAuth?.oid || agent?.ownerOid;
        if (!ownerOid && !agent?.entra) {
            // Session doesn't require Entra verification — non-Entra agent
            await encryptAndSend(client.ws, { type: 'entra_auth_ok', email: '' }, client.sessionKey);
            client.entraVerified = true;
            return;
        }
        if (!ownerOid && agent?.entra) {
            // Agent has entra=1 but hasn't verified yet — tell web client to wait
            await encryptAndSend(client.ws, {
                type: 'entra_auth_failed',
                reason: 'agent_pending',
            }, client.sessionKey);
            return;
        }
        try {
            const identity = await verifyEntraToken(msg.token);
            if (ownerOid && identity.oid !== ownerOid) {
                console.error(`[Web] Entra oid mismatch for client ${clientId.slice(0, 8)}: session owner ${ownerOid}, client ${identity.oid}`);
                await encryptAndSend(client.ws, {
                    type: 'entra_auth_failed',
                    reason: 'oid_mismatch',
                }, client.sessionKey);
                client.ws.close(4403, 'Entra identity mismatch');
                return;
            }
            client.entraVerified = true;
            console.log(`[Web] Entra identity verified for client ${clientId.slice(0, 8)}: ${identity.email}`);
            await encryptAndSend(client.ws, {
                type: 'entra_auth_ok',
                email: identity.email,
                msInternal: agent?.msInternal || false,
                agent: agent ? {
                    agentId: agent.agentId,
                    name: agent.name,
                    hostname: agent.hostname,
                    workDir: agent.workDir,
                    version: agent.version,
                } : null,
            }, client.sessionKey);
        }
        catch (err) {
            const reason = err instanceof jose.errors.JWTExpired ? 'token_expired' : 'invalid_token';
            console.error(`[Web] Entra auth failed for client ${clientId.slice(0, 8)}:`, err.message);
            await encryptAndSend(client.ws, {
                type: 'entra_auth_failed',
                reason,
            }, client.sessionKey);
        }
        return;
    }
    // Block non-keepalive messages from unverified web clients on Entra-protected sessions
    // Only gate when ownerOid is actually established, not just when agent.entra is set
    // Silently drop (don't send error) to avoid spamming the UI during verification
    const sessionOwnerBlock = sessions.sessionOwners.get(client.sessionId);
    if ((agent?.ownerOid || sessionOwnerBlock) && !client.entraVerified && msg.type !== 'ping') {
        return;
    }
    if (!agent || agent.ws.readyState !== WebSocket.OPEN) {
        await encryptAndSend(client.ws, { type: 'error', message: 'Agent not connected' }, client.sessionKey);
        return;
    }
    // Block forwarding to unverified agent on Entra-protected sessions
    if (agent.ownerOid && !agent.entraVerified) {
        await encryptAndSend(client.ws, { type: 'error', message: 'Agent is verifying identity, please wait' }, client.sessionKey);
        return;
    }
    // Forward web client message to agent, encrypted with agent's session key
    await encryptAndSend(agent.ws, msg, agent.sessionKey);
}
//# sourceMappingURL=ws-client.js.map