import { WebSocket } from 'ws';
import { randomUUID } from 'crypto';
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 { handleAgentTunnelMessage, updateProxyConfig, cleanupAgentTunnels } from './tunnel.js';
import { verifyEntraToken, jose } from './entra.js';
const agentRelay = new MessageRelay();
export function handleAgentConnection(ws, req) {
    const url = new URL(req.url || '/', `http://${req.headers.host}`);
    const agentId = url.searchParams.get('id') || randomUUID();
    const name = url.searchParams.get('name') || `Agent-${agentId.slice(0, 8)}`;
    const workDir = url.searchParams.get('workDir') || 'unknown';
    const hostname = url.searchParams.get('hostname') || '';
    const version = url.searchParams.get('version') || '';
    const password = url.searchParams.get('password') || '';
    const entra = url.searchParams.get('entra') === '1';
    const msInternal = url.searchParams.get('msInternal') === '1';
    // Hash password if provided (agent sends plaintext over WSS)
    let passwordHash = null;
    let passwordSalt = null;
    if (password) {
        const h = auth.hashPassword(password);
        passwordHash = h.hash;
        passwordSalt = h.salt;
    }
    // Reuse requested sessionId (agent reconnecting) or generate a new one
    const requestedSessionId = url.searchParams.get('sessionId');
    const sessionId = requestedSessionId || sessions.generateSessionId();
    if (!requestedSessionId) {
        console.log(`[Agent] No sessionId in query params for ${name} — generated new: ${sessionId}`);
    }
    // Preserve ownerOid from session-level store (survives agent disconnect/removal)
    const sessionOwner = requestedSessionId ? sessions.sessionOwners.get(requestedSessionId) : undefined;
    // Also check the previous agent object (in case it hasn't been removed yet)
    const previousAgent = requestedSessionId ? sessions.getAgentBySession(requestedSessionId) : undefined;
    const previousOwnerOid = sessionOwner?.oid || previousAgent?.ownerOid;
    const previousOwnerEmail = sessionOwner?.email || previousAgent?.ownerEmail;
    const sessionKey = generateSessionKey();
    const agent = {
        ws,
        agentId,
        name,
        hostname,
        workDir,
        version,
        sessionId,
        sessionKey,
        connectedAt: new Date(),
        isAlive: true,
        passwordHash,
        passwordSalt,
        entra,
        entraVerified: false,
        msInternal,
        ownerOid: previousOwnerOid,
        ownerEmail: previousOwnerEmail,
    };
    sessions.registerAgent(agentId, agent);
    // Persist password auth per session so it survives agent disconnects
    if (passwordHash && passwordSalt) {
        auth.setSessionPassword(sessionId, passwordHash, passwordSalt);
    }
    console.log(`[Agent] Registered: ${name} (${agentId}), session: ${sessionId}${requestedSessionId ? ' (reconnect)' : ''}${passwordHash ? ' (password protected)' : ''}`);
    // Send registration with session key (this initial message is plain text)
    try {
        ws.send(JSON.stringify({
            type: 'registered',
            agentId,
            sessionId,
            sessionKey: encodeKey(sessionKey),
        }));
    }
    catch (err) {
        console.error(`[Agent] Failed to send registration to ${name}:`, err.message);
    }
    // Notify any web clients already connected to this session (reconnect scenario)
    const isEntraSession = !!previousOwnerOid;
    for (const client of sessions.getClientsForSession(sessionId)) {
        if (client.ws.readyState === WebSocket.OPEN) {
            encryptAndSend(client.ws, {
                type: 'agent_reconnected',
                agent: { agentId, name, hostname, workDir, version },
                entraVerifying: isEntraSession, // tell web client agent is re-verifying
            }, client.sessionKey).catch((err) => {
                console.error(`[Agent] Failed to notify client of reconnect:`, err.message);
            });
        }
    }
    // Entra verification timeout: only enforce when session already has an owner
    // (reconnect to a previously verified session). Don't timeout on first connect
    // with entra=1 alone — old agents use --entra without entra_auth logic.
    if (previousOwnerOid) {
        setTimeout(() => {
            const current = sessions.getAgent(agentId);
            if (current && current.ws === ws && !current.entraVerified) {
                console.error(`[Agent] Entra verification timeout for ${name} — closing`);
                ws.close(4408, 'Entra verification timeout');
            }
        }, 30_000);
    }
    ws.on('message', (data) => {
        agentRelay.enqueue(agentId, () => handleAgentMessage(agentId, data.toString()));
    });
    ws.on('close', () => {
        console.log(`[Agent] Disconnected: ${name} (${agentId})`);
        agentRelay.cleanup(agentId);
        // Only clean up if this WebSocket is still the current one for this agent.
        // On reconnect, the new connection overwrites the agents Map entry before
        // the old connection's close event fires — deleting would remove the new entry.
        const current = sessions.getAgent(agentId);
        if (current && current.ws === ws) {
            sessions.removeAgent(agentId);
            cleanupAgentTunnels(agentId);
            // Notify connected web clients that agent is gone
            for (const client of sessions.getClientsForSession(sessionId)) {
                if (client.ws.readyState === WebSocket.OPEN) {
                    encryptAndSend(client.ws, { type: 'agent_disconnected' }, client.sessionKey).catch((err) => {
                        console.error(`[Agent] Failed to notify client of disconnect:`, err.message);
                    });
                }
            }
        }
    });
    ws.on('pong', () => {
        agent.isAlive = true;
    });
    ws.on('error', (err) => {
        console.error(`[Agent] WebSocket error for ${name} (${agentId}):`, err.message);
    });
}
async function handleAgentMessage(agentId, raw) {
    const agent = sessions.getAgent(agentId);
    if (!agent)
        return;
    const msg = await parseMessage(raw, agent.sessionKey);
    if (!msg) {
        console.error(`[Agent] Failed to parse/decrypt message from ${agentId} (key=${agent.sessionKey ? 'set' : 'null'})`);
        return;
    }
    // Log non-streaming messages for debugging
    if (msg.type !== 'claude_output' && msg.type !== 'pong') {
        console.log(`[Agent] ${agent.name} → ${msg.type}`);
    }
    // Intercept workdir_changed to keep server state in sync
    if (msg.type === 'workdir_changed' && typeof msg.workDir === 'string') {
        agent.workDir = msg.workDir;
        console.log(`[Agent] ${agent.name} changed workDir to: ${msg.workDir}`);
    }
    // Block all messages from unverified agents on Entra-protected sessions (except entra_auth itself)
    // Only gate when ownerOid is established — agent.entra alone doesn't mean verified identity exists
    if (agent.ownerOid && !agent.entraVerified && msg.type !== 'entra_auth') {
        console.warn(`[Agent] Blocked ${msg.type} from unverified agent ${agent.name}`);
        return;
    }
    // Handle entra_auth: verify JWT and bind identity to session
    if (msg.type === 'entra_auth' && typeof msg.token === 'string') {
        try {
            const identity = await verifyEntraToken(msg.token);
            if (agent.ownerOid && agent.ownerOid !== identity.oid) {
                // Reconnect with different identity — reject
                console.error(`[Agent] Entra oid mismatch for session ${agent.sessionId}: expected ${agent.ownerOid}, got ${identity.oid}`);
                await encryptAndSend(agent.ws, {
                    type: 'entra_auth_failed',
                    reason: 'oid_mismatch',
                }, agent.sessionKey);
                agent.ws.close(4403, 'Entra identity mismatch');
                return;
            }
            agent.ownerOid = identity.oid;
            agent.ownerEmail = identity.email;
            agent.entraVerified = true;
            // Persist ownerOid at session level so it survives agent disconnect
            sessions.setSessionOwner(agent.sessionId, identity.oid, identity.email);
            console.log(`[Agent] Entra identity verified for ${agent.name}: ${identity.email} (oid: ${identity.oid})`);
            await encryptAndSend(agent.ws, {
                type: 'entra_auth_ok',
                email: identity.email,
            }, agent.sessionKey);
            // Request proxy config now that agent is verified (initial send was blocked)
            await encryptAndSend(agent.ws, { type: 'get_proxy_config' }, agent.sessionKey);
            // Notify existing web clients that Entra verification is now required
            for (const client of sessions.getClientsForSession(agent.sessionId)) {
                if (client.ws.readyState !== WebSocket.OPEN)
                    continue;
                if (!client.entraVerified) {
                    await encryptAndSend(client.ws, { type: 'entra_auth_required' }, client.sessionKey);
                }
                else {
                    // Agent (re-)verified — notify verified clients that agent is back
                    await encryptAndSend(client.ws, {
                        type: 'agent_reconnected',
                        agent: { agentId: agent.agentId, name: agent.name, hostname: agent.hostname, workDir: agent.workDir, version: agent.version },
                    }, client.sessionKey);
                }
            }
        }
        catch (err) {
            const reason = err instanceof jose.errors.JWTExpired ? 'token_expired' : 'invalid_token';
            console.error(`[Agent] Entra auth failed for ${agent.name}:`, err.message);
            await encryptAndSend(agent.ws, {
                type: 'entra_auth_failed',
                reason,
            }, agent.sessionKey);
            // Notify web clients that agent token is expired
            if (reason === 'token_expired') {
                for (const client of sessions.getClientsForSession(agent.sessionId)) {
                    if (client.ws.readyState === WebSocket.OPEN && client.entraVerified) {
                        await encryptAndSend(client.ws, { type: 'entra_agent_token_expired' }, client.sessionKey);
                    }
                }
            }
        }
        return; // entra_auth is not forwarded to web clients
    }
    // Intercept proxy_config_updated to cache config on server side
    if (msg.type === 'proxy_config_updated' && msg.config) {
        updateProxyConfig(agent.agentId, msg.config);
        // Fall through — also forward to web clients
    }
    // Route tunnel messages to tunnel handler (do NOT forward to web clients)
    if (handleAgentTunnelMessage(msg)) {
        return;
    }
    // Forward agent messages to all web clients connected to this session
    let relayCount = 0;
    for (const client of sessions.getClientsForSession(agent.sessionId)) {
        if (client.ws.readyState !== WebSocket.OPEN)
            continue;
        // Skip unverified web clients on Entra-protected sessions (only when ownerOid established)
        if (agent.ownerOid && !client.entraVerified)
            continue;
        await encryptAndSend(client.ws, msg, client.sessionKey);
        relayCount++;
    }
    if (relayCount === 0 && msg.type !== 'claude_output' && msg.type !== 'pong') {
        console.warn(`[Agent] No web clients to relay ${msg.type} for session ${agent.sessionId}`);
    }
}
//# sourceMappingURL=ws-agent.js.map