"use strict";
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
    return new (P || (P = Promise))(function (resolve, reject) {
        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
        step((generator = generator.apply(thisArg, _arguments || [])).next());
    });
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.handleWebSocketMessage = exports.createWebSocketServer = exports.broadcastToSession = exports.sessions = exports.SessionState = void 0;
const ws_1 = require("ws");
const LLMService_1 = require("../services/LLMService");
const EncryptionService_1 = require("../services/EncryptionService");
// Initialize Encryption Service
const encryptionService = new EncryptionService_1.EncryptionService(process.env.ENCRYPTION_KEY || '');
// Define the SessionState enum
var SessionState;
(function (SessionState) {
    SessionState["SETUP"] = "SETUP";
    SessionState["GATHERING"] = "GATHERING";
    SessionState["HARMONIZING"] = "HARMONIZING";
    SessionState["FINAL"] = "FINAL";
    SessionState["ERROR"] = "ERROR";
})(SessionState = exports.SessionState || (exports.SessionState = {}));
exports.sessions = new Map();
// Initialize LLM Service (API key from environment)
const llmService = new LLMService_1.LLMService(process.env.GEMINI_API_KEY || '');
// Structured logging function
const logEvent = (eventName, sessionId, details = {}) => {
    console.log(JSON.stringify(Object.assign({ timestamp: new Date().toISOString(), eventName, sessionId }, details)));
};
// Metrics recording function
const recordMetric = (metricName, value, sessionId, details = {}) => {
    console.log(JSON.stringify(Object.assign({ timestamp: new Date().toISOString(), metricName, value, sessionId }, details)));
};
// Helper to create a serializable version of the session state
const getSerializableSession = (sessionData, currentClientId = null) => {
    const filteredResponses = new Map();
    sessionData.responses.forEach((response, clientId) => {
        if (clientId === currentClientId) {
            // For the current client, decrypt and send their own full response
            const decryptedWants = response.wants.map((d) => encryptionService.decrypt(d));
            const decryptedAccepts = response.accepts.map((d) => encryptionService.decrypt(d));
            const decryptedNoGoes = response.noGoes.map((d) => encryptionService.decrypt(d));
            const decryptedAfraidToAsk = encryptionService.decrypt(response.afraidToAsk);
            filteredResponses.set(clientId, { wants: decryptedWants, accepts: decryptedAccepts, noGoes: decryptedNoGoes, afraidToAsk: decryptedAfraidToAsk });
        }
        else {
            // For other clients, only send non-AfraidToAsk parts (wants, accepts, noGoes without AfraidToAsk)
            const decryptedWants = response.wants.map((d) => encryptionService.decrypt(d));
            const decryptedAccepts = response.accepts.map((d) => encryptionService.decrypt(d));
            const decryptedNoGoes = response.noGoes.map((d) => encryptionService.decrypt(d));
            filteredResponses.set(clientId, { wants: decryptedWants, accepts: decryptedAccepts, noGoes: decryptedNoGoes, afraidToAsk: "" }); // Hide afraidToAsk for other clients
        }
    });
    return Object.assign(Object.assign({}, sessionData), { responses: Object.fromEntries(filteredResponses), clients: Array.from(sessionData.clients.keys()) });
};
const broadcastToSession = (sessionId, message, excludeClientId = null) => {
    const sessionData = exports.sessions.get(sessionId);
    if (sessionData) {
        sessionData.clients.forEach((client, clientId) => {
            if (clientId !== excludeClientId && client.readyState === ws_1.WebSocket.OPEN) {
                const serializableMessage = Object.assign(Object.assign({}, message), { payload: Object.assign(Object.assign({}, message.payload), { session: getSerializableSession(sessionData, clientId) }) });
                client.send(JSON.stringify(serializableMessage));
            }
        });
    }
};
exports.broadcastToSession = broadcastToSession;
const createWebSocketServer = (server) => {
    const wss = new ws_1.WebSocketServer({ server });
    wss.on('connection', (ws, req) => {
        const url = new URL(req.url || '', `http://${req.headers.host}`);
        const sessionId = url.pathname.split('/').pop();
        if (!sessionId) {
            ws.close(1008, 'Invalid session ID');
            return;
        }
        if (!exports.sessions.has(sessionId)) {
            exports.sessions.set(sessionId, {
                state: SessionState.SETUP,
                topic: null,
                description: null,
                expectedResponses: 0,
                submittedCount: 0,
                responses: new Map(),
                clients: new Map(),
                finalResult: null,
            });
        }
        const sessionData = exports.sessions.get(sessionId);
        console.log(`Client connecting to session: ${sessionId}`);
        // Set up a ping interval to keep the connection alive
        const pingInterval = setInterval(() => {
            if (ws.readyState === ws_1.WebSocket.OPEN) {
                ws.ping();
            }
        }, 30000); // Send ping every 30 seconds
        ws.on('message', (message) => __awaiter(void 0, void 0, void 0, function* () {
            const parsedMessage = JSON.parse(message.toString());
            const { type, clientId, payload } = parsedMessage;
            if (!clientId) {
                console.error(`Received message without clientId in session ${sessionId}. Type: ${type}`);
                ws.send(JSON.stringify({ type: 'ERROR', payload: { message: 'clientId is required' } }));
                return;
            }
            if (!sessionData.clients.has(clientId)) {
                sessionData.clients.set(clientId, ws);
                console.log(`Client ${clientId} registered for session: ${sessionId}. Total clients: ${sessionData.clients.size}`);
                ws.send(JSON.stringify({ type: 'STATE_UPDATE', payload: { session: getSerializableSession(sessionData, clientId) } }));
            }
            console.log(`Received message from ${clientId} in session ${sessionId}:`, type);
            yield (0, exports.handleWebSocketMessage)(ws, sessionId, parsedMessage);
        }));
        ws.on('close', () => {
            clearInterval(pingInterval); // Clear the interval when the connection closes
            let disconnectedClientId = null;
            for (const [clientId, clientWs] of sessionData.clients.entries()) {
                if (clientWs === ws) {
                    disconnectedClientId = clientId;
                    break;
                }
            }
            if (disconnectedClientId) {
                sessionData.clients.delete(disconnectedClientId);
                console.log(`Client ${disconnectedClientId} disconnected from session: ${sessionId}. Remaining clients: ${sessionData.clients.size}`);
            }
            else {
                console.log(`An unregistered client disconnected from session: ${sessionId}.`);
            }
            if (sessionData.clients.size === 0) {
                exports.sessions.delete(sessionId);
                logEvent('session_purged', sessionId);
                console.log(`Session ${sessionId} closed and state cleared.`);
            }
        });
        ws.on('error', (error) => {
            console.error(`WebSocket error in session ${sessionId}:`, error);
        });
    });
    return wss;
};
exports.createWebSocketServer = createWebSocketServer;
const handleWebSocketMessage = (ws, sessionId, parsedMessage) => __awaiter(void 0, void 0, void 0, function* () {
    const { type, clientId, payload } = parsedMessage;
    if (!clientId) {
        console.error(`Received message without clientId in session ${sessionId}. Type: ${type}`);
        ws.send(JSON.stringify({ type: 'ERROR', payload: { message: 'clientId is required' } }));
        return;
    }
    const sessionData = exports.sessions.get(sessionId);
    if (!sessionData.clients.has(clientId)) {
        sessionData.clients.set(clientId, ws);
        console.log(`Client ${clientId} registered for session: ${sessionId}. Total clients: ${sessionData.clients.size}`);
        ws.send(JSON.stringify({ type: 'STATE_UPDATE', payload: { session: getSerializableSession(sessionData, clientId) } }));
    }
    console.log(`Received message from ${clientId} in session ${sessionId}:`, type);
    switch (type) {
        case 'REGISTER_CLIENT':
            console.log(`Client ${clientId} registered successfully for session ${sessionId}.`);
            break;
        case 'PING':
            // Respond to client pings with a pong
            if (ws.readyState === ws_1.WebSocket.OPEN) {
                ws.pong();
            }
            break;
        case 'SETUP_SESSION':
            if (sessionData.state === SessionState.SETUP) {
                const { expectedResponses, topic, description } = payload;
                if (typeof expectedResponses !== 'number' || expectedResponses <= 0) {
                    ws.send(JSON.stringify({ type: 'ERROR', payload: { message: 'Invalid expectedResponses' } }));
                    return;
                }
                sessionData.expectedResponses = expectedResponses;
                sessionData.topic = topic || 'Untitled Session';
                sessionData.description = description || null;
                sessionData.state = SessionState.GATHERING;
                (0, exports.broadcastToSession)(sessionId, { type: 'STATE_UPDATE', payload: {} });
                console.log(`Session ${sessionId} moved to GATHERING with topic "${sessionData.topic}" and ${expectedResponses} expected responses.`);
            }
            else {
                ws.send(JSON.stringify({ type: 'ERROR', payload: { message: `Session is not in SETUP state. Current state: ${sessionData.state}` } }));
            }
            break;
        case 'SUBMIT_RESPONSE':
            if (sessionData.state === SessionState.GATHERING) {
                if (sessionData.responses.has(clientId)) {
                    ws.send(JSON.stringify({ type: 'ERROR', payload: { message: 'You have already submitted a response for this session.' } }));
                    return;
                }
                const { wants, accepts, noGoes, afraidToAsk } = payload.response;
                if ([...wants, ...accepts, ...noGoes].some(desire => desire.length > 500) || afraidToAsk.length > 500) {
                    ws.send(JSON.stringify({ type: 'ERROR', payload: { message: 'One of your desires or afraidToAsk exceeds the 500 character limit.' } }));
                    return;
                }
                const hasContradictionsGist = yield llmService.checkForInnerContradictions(payload.response);
                if (hasContradictionsGist) {
                    ws.send(JSON.stringify({ type: 'ERROR', payload: { message: `Your submission contains inner contradictions: ${hasContradictionsGist} Please resolve them and submit again.` } }));
                    return;
                }
                const encryptedWants = wants.map((d) => encryptionService.encrypt(d));
                const encryptedAccepts = accepts.map((d) => encryptionService.encrypt(d));
                const encryptedNoGoes = noGoes.map((d) => encryptionService.encrypt(d));
                const encryptedAfraidToAsk = encryptionService.encrypt(afraidToAsk);
                sessionData.responses.set(clientId, { wants: encryptedWants, accepts: encryptedAccepts, noGoes: encryptedNoGoes, afraidToAsk: encryptedAfraidToAsk });
                sessionData.submittedCount++;
                logEvent('response_submitted', sessionId, { clientId, submittedCount: sessionData.submittedCount });
                console.log(`Client ${clientId} submitted response. Submitted count: ${sessionData.submittedCount}/${sessionData.expectedResponses}`);
                if (sessionData.submittedCount === sessionData.expectedResponses) {
                    sessionData.state = SessionState.HARMONIZING;
                    (0, exports.broadcastToSession)(sessionId, { type: 'STATE_UPDATE', payload: {} });
                    logEvent('session_harmonizing', sessionId, { expectedResponses: sessionData.expectedResponses });
                    console.log(`Session ${sessionId} moved to HARMONIZING. Triggering LLM analysis.`);
                    // Perform LLM analysis asynchronously
                    (() => __awaiter(void 0, void 0, void 0, function* () {
                        let durationMs = 0; // Declare here
                        try {
                            logEvent('llm_analysis_started', sessionId);
                            const startTime = process.hrtime.bigint();
                            const allDecryptedDesires = Array.from(sessionData.responses.values()).map(encryptedResponse => {
                                const decryptedWants = encryptedResponse.wants.map((d) => encryptionService.decrypt(d));
                                const decryptedAccepts = encryptedResponse.accepts.map((d) => encryptionService.decrypt(d));
                                const decryptedNoGoes = encryptedResponse.noGoes.map((d) => encryptionService.decrypt(d));
                                const decryptedAfraidToAsk = encryptionService.decrypt(encryptedResponse.afraidToAsk);
                                return { wants: decryptedWants, accepts: decryptedAccepts, noGoes: decryptedNoGoes, afraidToAsk: decryptedAfraidToAsk };
                            });
                            const decision = yield llmService.analyzeDesires(allDecryptedDesires);
                            sessionData.finalResult = decision;
                            sessionData.state = SessionState.FINAL;
                            (0, exports.broadcastToSession)(sessionId, { type: 'STATE_UPDATE', payload: {} });
                            logEvent('llm_analysis_completed', sessionId, { result: decision });
                            recordMetric('llm_analysis_duration', durationMs, sessionId, { status: 'success' });
                            recordMetric('llm_analysis_availability', 'available', sessionId);
                            console.log(`Analysis complete for session ${sessionId}. Result:`, decision);
                        }
                        catch (error) {
                            console.error(`Error during analysis for session ${sessionId}:`, error.message);
                            sessionData.state = SessionState.ERROR;
                            (0, exports.broadcastToSession)(sessionId, { type: 'STATE_UPDATE', payload: {} });
                            logEvent('llm_analysis_error', sessionId, { error: error.message });
                            recordMetric('llm_analysis_availability', 'unavailable', sessionId, { error: error.message });
                        }
                    }))();
                }
                else {
                    // Only broadcast the latest count if the session is not yet harmonizing
                    (0, exports.broadcastToSession)(sessionId, { type: 'STATE_UPDATE', payload: {} });
                }
            }
            else {
                ws.send(JSON.stringify({ type: 'ERROR', payload: { message: `Session is not in GATHERING state. Current state: ${sessionData.state}` } }));
            }
            break;
        default:
            console.warn(`Unknown message type: ${type} from client ${clientId} in session ${sessionId}`);
            ws.send(JSON.stringify({ type: 'ERROR', payload: { message: `Unknown message type: ${type}` } }));
            break;
    }
});
exports.handleWebSocketMessage = handleWebSocketMessage;
