websocket-server.ts

import { WebSocketServer, WebSocket } from 'ws'
import http from 'http'
import {
	getElevenLabsWebSocket,
	cleanupElevenLabsWebSocket,
} from './lib/elevenlabs-websocket'
import {
	streamAIToSpeech,
	StreamingAIMessage,
} from './lib/streaming-ai-service'
import { transcribeAudio } from './lib/ai-service'
import { loadPrompts, buildSystemPrompt } from './lib/ai-service'
import { notepadService } from './lib/notepad-service'
import { vaultService } from './lib/vault-service'
import { checkConversationComplete } from './lib/message-service'

// Disable bufferutil to prevent bundler compatibility issues
process.env.WS_NO_BUFFER_UTIL = '1'

// WebSocket message types
interface ClientMessage {
	type: 'audio' | 'text' | 'initialize' | 'disconnect' | 'heartbeat'
	data?: string // Base64 audio or text content
	contextId?: string
	conversationId?: string
	isWelcomeMessage?: boolean
	elapsedTime?: string
	usePaywallPrompt?: boolean
	userName?: string // User's first name for personalization
}

interface ServerMessage {
	type: 'audio' | 'text' | 'transcription' | 'error' | 'complete'
	data?: string | Uint8Array
	contextId?: string
	conversationId?: string
	isComplete?: boolean
	transcribedText?: string
	responseText?: string
}

// Create HTTP server (nginx handles SSL termination)
const server = http.createServer()
const port = parseInt(process.env.WS_PORT || '8080')
const protocol = 'ws'

const wss = new WebSocketServer({
	server,
	perMessageDeflate: false,
	maxPayload: 50 * 1024 * 1024, // 50MB max payload
	skipUTF8Validation: true, // Skip UTF8 validation for performance
})

// Connection management
const connections = new Map<
	WebSocket,
	{
		userId?: string
		contextId: string
		elevenLabsService: any
		isAuthenticated: boolean
		hasSubscription: boolean
		conversationId?: string
		messageHistory: StreamingAIMessage[]
		isProcessing: boolean
		pingInterval?: NodeJS.Timeout
		lastPong: number
		isAlive: boolean
	}
>()

// Add process-level error handlers to prevent crashes
process.on('uncaughtException', (error) => {
	// Don't exit the process, just continue
})

process.on('unhandledRejection', (reason, promise) => {
	// Continue despite unhandled rejection
})

process.on('exit', (code) => {
	// Process exiting
})

process.on('SIGTERM', (signal) => {
	// Signal received
})

process.on('SIGINT', (signal) => {
	// Signal received
})

/**
 * Start heartbeat for a WebSocket connection
 */
const startHeartbeat = (ws: WebSocket) => {
	const connectionInfo = connections.get(ws)
	if (!connectionInfo) return

	// Clear any existing interval
	if (connectionInfo.pingInterval) {
		clearInterval(connectionInfo.pingInterval)
	}

	// Check for heartbeat every 30 seconds
	connectionInfo.pingInterval = setInterval(() => {
		// Validate WebSocket still exists and is in valid state
		if (!ws || ws.readyState !== WebSocket.OPEN) {
			stopHeartbeat(ws)
			return
		}

		const now = Date.now()
		const timeSinceLastActivity = now - connectionInfo.lastPong

		// If no heartbeat received in 60 seconds, close connection
		if (timeSinceLastActivity > 60000) {
			stopHeartbeat(ws)

			// Safely close the connection
			try {
				// Remove from connections map first
				connections.delete(ws)

				// Then close the WebSocket
				if (
					ws.readyState === WebSocket.OPEN ||
					ws.readyState === WebSocket.CONNECTING
				) {
					ws.close(1001, 'Heartbeat timeout')
				}
			} catch (error) {
				// Error terminating connection
			}
			return
		}

		// Request heartbeat from client
		sendMessage(ws, {
			type: 'complete',
			data: 'heartbeat_request',
			contextId: connectionInfo.contextId,
		})
	}, 30000) // 30 seconds

	// Mark connection as alive initially
	connectionInfo.isAlive = true
	connectionInfo.lastPong = Date.now()
}

/**
 * Stop heartbeat for a WebSocket connection
 */
const stopHeartbeat = (ws: WebSocket) => {
	const connectionInfo = connections.get(ws)
	if (connectionInfo && connectionInfo.pingInterval) {
		clearInterval(connectionInfo.pingInterval)
		connectionInfo.pingInterval = undefined
	}
}

/**
 * Send message to WebSocket client
 */
const sendMessage = (ws: WebSocket, message: ServerMessage) => {
	if (ws.readyState === WebSocket.OPEN) {
		try {
			ws.send(JSON.stringify(message))
		} catch (error) {
			// Error sending message
		}
	}
}

/**
 * Send error message to WebSocket client
 */
const sendErrorMessage = (ws: WebSocket, error: string) => {
	sendMessage(ws, {
		type: 'error',
		data: error,
	})
}

/**
 * Handle messages from clients
 */
const handleClientMessage = async (ws: WebSocket, message: ClientMessage) => {
	const connectionInfo = connections.get(ws)
	if (!connectionInfo) {
		sendErrorMessage(ws, 'Connection not found')
		return
	}

	// Update last activity on any message
	connectionInfo.lastPong = Date.now()
	connectionInfo.isAlive = true

	const {
		type,
		data,
		contextId,
		conversationId,
		isWelcomeMessage,
		elapsedTime,
		usePaywallPrompt,
		userName,
	} = message

	try {
		switch (type) {
			case 'initialize':
				await handleInitialize(ws, connectionInfo)
				break

			case 'audio':
				if (!data) {
					sendErrorMessage(ws, 'Audio data is required')
					return
				}
				await handleAudioMessage(ws, connectionInfo, data, {
					conversationId,
					isWelcomeMessage,
					elapsedTime,
					usePaywallPrompt,
					userName,
				})
				break

			case 'text':
				if (!data) {
					sendErrorMessage(ws, 'Text data is required')
					return
				}
				await handleTextMessage(ws, connectionInfo, data, {
					conversationId,
					isWelcomeMessage,
					elapsedTime,
					usePaywallPrompt,
					userName,
				})
				break

			case 'disconnect':
				await handleDisconnect(ws, connectionInfo)
				break

			case 'heartbeat':
				// Update last activity time
				connectionInfo.lastPong = Date.now()
				connectionInfo.isAlive = true
				sendMessage(ws, {
					type: 'complete',
					data: 'heartbeat_ack',
					contextId: connectionInfo.contextId,
				})
				break

			default:
				sendErrorMessage(ws, `Unknown message type: ${type}`)
		}
	} catch (error) {
		sendErrorMessage(
			ws,
			error instanceof Error ? error.message : 'Unknown error'
		)
	}
}

/**
 * Handle connection initialization
 */
const handleInitialize = async (ws: WebSocket, connectionInfo: any) => {
	try {
		// Skip authentication for WebSocket connections
		connectionInfo.userId = null
		connectionInfo.isAuthenticated = false
		connectionInfo.hasSubscription = true // Allow for preview/demo mode

		// Initialize ElevenLabs WebSocket connection
		if (!connectionInfo.elevenLabsService.connected) {
			try {
				await connectionInfo.elevenLabsService.connect()
				// Wait a bit for connection to be fully ready
				await new Promise((resolve) => setTimeout(resolve, 100))
			} catch (error) {
				console.error(
					'🔌 [Voice Stream WS] Failed to connect to ElevenLabs:',
					error
				)
				// Don't throw, just log - connection might work later
			}
		}

		sendMessage(ws, {
			type: 'complete',
			contextId: connectionInfo.contextId,
			data: 'Initialized successfully',
		})

		console.log(
			`🔌 [Voice Stream WS] Initialized connection for user: anonymous`
		)
	} catch (error) {
		sendErrorMessage(ws, 'Failed to initialize connection')
	}
}

/**
 * Handle audio input messages
 */
const handleAudioMessage = async (
	ws: WebSocket,
	connectionInfo: any,
	audioData: string,
	options: {
		conversationId?: string
		isWelcomeMessage?: boolean
		elapsedTime?: string
		usePaywallPrompt?: boolean
		userName?: string
	}
) => {
	try {
		// Convert base64 audio to Blob (works in server environments)
		const audioBuffer = Buffer.from(audioData, 'base64')
		// Create a Blob-like object that works with transcribeAudio
		const audioFile = new Blob([audioBuffer], {
			type: 'audio/webm',
		}) as any
		// Add name property for compatibility
		;(audioFile as any).name = 'audio.webm'

		// Transcribe audio
		const transcribedText = await transcribeAudio(audioFile)

		// Send transcription to client
		sendMessage(ws, {
			type: 'transcription',
			data: transcribedText,
			contextId: connectionInfo.contextId,
		})

		// Check if transcription is empty
		if (!transcribedText || transcribedText.trim() === '') {
			// Send completion message with empty transcription flag
			sendMessage(ws, {
				type: 'complete',
				data: 'No speech detected',
				contextId: connectionInfo.contextId,
				transcribedText: '',
				responseText: 'No speech detected',
			})
			return
		}

		// Process the transcribed text
		await processUserInput(ws, connectionInfo, transcribedText, options)
	} catch (error) {
		sendErrorMessage(ws, 'Failed to process audio')
	}
}

/**
 * Handle text input messages
 */
const handleTextMessage = async (
	ws: WebSocket,
	connectionInfo: any,
	text: string,
	options: {
		conversationId?: string
		isWelcomeMessage?: boolean
		elapsedTime?: string
		usePaywallPrompt?: boolean
		userName?: string
	}
) => {
	try {
		await processUserInput(ws, connectionInfo, text, options)
	} catch (error) {
		sendErrorMessage(ws, 'Failed to process text')
	}
}

/**
 * Process user input and generate streaming response
 */
const processUserInput = async (
	ws: WebSocket,
	connectionInfo: any,
	userInput: string,
	options: {
		conversationId?: string
		isWelcomeMessage?: boolean
		elapsedTime?: string
		usePaywallPrompt?: boolean
		userName?: string
	}
) => {
	const { userId, contextId, elevenLabsService } = connectionInfo
	const {
		conversationId,
		isWelcomeMessage,
		elapsedTime,
		usePaywallPrompt,
		userName,
	} = options

	// Check if already processing
	if (connectionInfo.isProcessing) {
		return
	}

	// Set processing flag
	connectionInfo.isProcessing = true

	try {
		// Handle welcome message
		if (isWelcomeMessage) {
			const welcomeMessage = userId
				? "Hey! I'm Dena, your AI advisor. I'm here to help you navigate the challenges of building and scaling your B2B startup. What's on your mind today?"
				: "Hi there! It's great to meet you. Quick question—how did you learn about me?"

			// Stream welcome message
			await streamResponse(ws, connectionInfo, [], welcomeMessage)
			return
		}

		// Build conversation context
		const messages = await buildConversationContext(
			userId,
			userInput,
			connectionInfo.messageHistory,
			elapsedTime,
			usePaywallPrompt,
			userName
		)

		// Simple approach: just start the new request without complex cleanup

		// Start streaming AI response with TTS
		try {
			await streamAIResponse(ws, connectionInfo, messages)
		} catch (error) {
			sendErrorMessage(ws, 'Failed to generate AI response')
		}
	} catch (error) {
		sendErrorMessage(ws, 'Failed to process input')
	} finally {
		// Clear processing flag
		connectionInfo.isProcessing = false
	}
}

/**
 * Build conversation context with prompts and user data
 */
const buildConversationContext = async (
	userId: string | undefined,
	userInput: string,
	messageHistory: StreamingAIMessage[],
	elapsedTime?: string,
	usePaywallPrompt?: boolean,
	userName?: string
): Promise<StreamingAIMessage[]> => {
	// Load prompts - support paywall prompt for initial paywall conversations
	const promptPaths = usePaywallPrompt
		? ['base_prompt.txt', 'initial_paywall_prompt.txt'] // For paywall conversation
		: userId
		? ['base_prompt.txt', 'advisory_coach_prompt.txt']
		: ['base_prompt.txt', 'preview_prompt.txt']

	const prompts = await loadPrompts(promptPaths)

	// Get notepad context if authenticated
	let notepadContext = ''
	if (userId) {
		const notepadItems = await notepadService.getNotepadItems(userId)
		if (notepadItems.length > 0) {
			notepadContext = `\n\nCURRENT NOTEPAD ITEMS:\n${notepadItems
				.map(
					(item) =>
						`- ${item.type.toUpperCase()}: ${item.content}${
							item.type === 'action_item' && !item.completed
								? ' [INCOMPLETE]'
								: ''
						}`
				)
				.join('\n')}`
		} else {
			notepadContext = '\n\nNOTEPAD: Currently empty'
		}
	}

	// Get vault context if authenticated
	let vaultContext = ''
	if (userId) {
		const vaultItems = await vaultService.getVaultItems(userId)
		if (vaultItems.length > 0) {
			vaultContext = `\n\nVAULT CLARITY SESSIONS:\n${vaultItems
				.map(
					(item, index) =>
						`Session ${index + 1} (${new Date(
							item.created_at
						).toLocaleDateString()}):\n` +
						`  Challenge: ${item.original_challenge}\n` +
						`  Product: ${item.product_description}\n` +
						`  Problem: ${item.problem_statement}\n` +
						`  Target Customer: ${item.ideal_customer_profile}\n` +
						`  Value Prop: ${item.value_proposition}\n` +
						`  Attempted Solutions: ${item.attempted_solutions}\n` +
						`  Recommendations: ${item.recommended_solutions}`
				)
				.join('\n\n')}`
		} else {
			vaultContext = '\n\nVAULT: No clarity sessions saved yet'
		}
	}

	// Build system prompt with user name for paywall conversations
	let systemPromptContext = notepadContext + vaultContext
	if (usePaywallPrompt && userName) {
		systemPromptContext += `\n\nUSER'S NAME: ${userName} (use this name naturally in conversation)`
	}

	const systemPrompt = buildSystemPrompt({
		basePrompt: prompts[0],
		extensions: prompts.slice(1),
		context: systemPromptContext,
		elapsedTime,
	})

	// Build conversation with history
	const messages: StreamingAIMessage[] = [
		{ role: 'system', content: systemPrompt },
		...messageHistory,
		{ role: 'user', content: userInput },
	]

	return messages
}

/**
 * Stream AI response with real-time TTS
 */
const streamAIResponse = async (
	ws: WebSocket,
	connectionInfo: any,
	messages: StreamingAIMessage[]
) => {
	const { contextId, elevenLabsService } = connectionInfo

	try {
		// Start the streaming pipeline
		const { textStream, audioStream } = await streamAIToSpeech(
			messages,
			elevenLabsService,
			contextId,
			{
				onTextChunk: (chunk: string) => {
					// Send text chunks to client in real-time
					sendMessage(ws, {
						type: 'text',
						data: chunk,
						contextId,
					})
				},
				onAudioChunk: (audioBuffer: Buffer) => {
					// Optimize audio chunk delivery for seamless playback
					const audioData = new Uint8Array(audioBuffer)

					// Skip tiny chunks that can cause glitches
					if (audioData.length < 64) {
						return
					}

					// Send audio chunks to client with proper buffering info
					sendMessage(ws, {
						type: 'audio',
						data: audioData,
						contextId,
					})
				},
				onComplete: (fullText: string) => {
					// Add user message and AI response to conversation history
					const userMessage = messages[messages.length - 1] // Last message is user input
					connectionInfo.messageHistory.push(userMessage)
					connectionInfo.messageHistory.push({
						role: 'assistant',
						content: fullText,
					})

					// Keep history manageable (last 20 messages)
					if (connectionInfo.messageHistory.length > 20) {
						connectionInfo.messageHistory =
							connectionInfo.messageHistory.slice(-20)
					}

					// Check if this is a conversation completion (for preview mode)
					const isComplete = checkConversationComplete(fullText)
					console.log('[WebSocket Server] Checking conversation completion:', {
						fullTextLength: fullText.length,
						last100Chars: fullText.slice(-100),
						isComplete,
						hasEndingPhrase: fullText.includes("Let's keep talking"),
					})

					// Send completion message with isComplete flag
					sendMessage(ws, {
						type: 'complete',
						data: fullText,
						contextId,
						responseText: fullText,
						isComplete,
					})
				},
				onError: (error: Error) => {
					sendErrorMessage(ws, error.message)
				},
			}
		)
	} catch (error) {
		sendErrorMessage(ws, 'Failed to generate response')
	}
}

/**
 * Stream a pre-generated response (for welcome messages)
 */
const streamResponse = async (
	ws: WebSocket,
	connectionInfo: any,
	messages: StreamingAIMessage[],
	responseText: string
) => {
	const { contextId, elevenLabsService } = connectionInfo

	try {
		// Send text in chunks to simulate streaming
		const words = responseText.split(' ')
		let currentText = ''

		for (const word of words) {
			currentText += word + ' '

			sendMessage(ws, {
				type: 'text',
				data: word + ' ',
				contextId,
			})

			// Small delay to simulate streaming
			await new Promise((resolve) => setTimeout(resolve, 50))
		}

		// Generate audio for the complete response
		if (!elevenLabsService.connected) {
			await elevenLabsService.connect()
		}

		// No listeners here - let ElevenLabs service handle everything internally
		// Just register our callback for this context
		elevenLabsService.setAudioCallback(contextId, (audioBuffer: Buffer) => {
			sendMessage(ws, {
				type: 'audio',
				data: new Uint8Array(audioBuffer),
				contextId,
			})
		})

		elevenLabsService.setCompleteCallback(contextId, () => {
			// Check if this is a conversation completion (for preview mode)
			const isComplete = checkConversationComplete(responseText)
			console.log('[WebSocket Server] Welcome message completion check:', {
				responseTextLength: responseText.length,
				isComplete,
			})

			sendMessage(ws, {
				type: 'complete',
				data: responseText,
				contextId,
				responseText,
				isComplete,
			})
		})

		// Send text to TTS
		await elevenLabsService.sendText(responseText, contextId, true)
	} catch (error) {
		sendErrorMessage(ws, 'Failed to stream response')
	}
}

/**
 * Handle disconnect messages
 */
const handleDisconnect = async (ws: WebSocket, connectionInfo: any) => {
	try {
		// Stop heartbeat first
		stopHeartbeat(ws)

		// Clean up ElevenLabs context
		if (connectionInfo.elevenLabsService) {
			await connectionInfo.elevenLabsService.closeContext(
				connectionInfo.contextId
			)
		}

		// Close WebSocket connection if still open
		if (ws.readyState === WebSocket.OPEN) {
			ws.close(1000, 'Client disconnect')
		}
	} catch (error) {}
}

// Handle new WebSocket connections
wss.on('connection', async (ws: WebSocket, request: any) => {
	// Generate unique context ID for this connection
	const contextId = `voice_${Date.now()}_${Math.random()
		.toString(36)
		.substr(2, 9)}`

	// Initialize ElevenLabs WebSocket service for this connection
	const elevenLabsService = getElevenLabsWebSocket(contextId)

	// Store connection info
	connections.set(ws, {
		contextId,
		elevenLabsService,
		isAuthenticated: false,
		hasSubscription: true,
		messageHistory: [],
		isProcessing: false, // Add processing flag
		isAlive: true,
		lastPong: Date.now(),
	})

	// Set up message handler
	ws.on('message', async (data: Buffer) => {
		try {
			const message: ClientMessage = JSON.parse(data.toString())
			await handleClientMessage(ws, message)
		} catch (error) {
			sendErrorMessage(ws, 'Invalid message format')
		}
	})

	// Note: Browser WebSocket clients don't expose ping/pong to JavaScript
	// So we'll handle heartbeat at application level instead

	// Handle connection close
	ws.on('close', (code, reason) => {
		const connectionInfo = connections.get(ws)
		if (connectionInfo) {
			// Stop heartbeat
			stopHeartbeat(ws)

			// Clean up ElevenLabs instance completely
			try {
				cleanupElevenLabsWebSocket(connectionInfo.contextId)
			} catch (error) {}
		}
		connections.delete(ws)
	})

	// Handle errors
	ws.on('error', (error) => {
		const connectionInfo = connections.get(ws)
		if (connectionInfo) {
			stopHeartbeat(ws)
			try {
				cleanupElevenLabsWebSocket(connectionInfo.contextId)
			} catch (cleanupError) {}
		}
		connections.delete(ws)
	})

	// Send initial connection confirmation
	sendMessage(ws, {
		type: 'complete',
		contextId,
		data: 'Connected to voice streaming service',
	})

	// Start heartbeat
	startHeartbeat(ws)
})

// Add error handling for the WebSocket server
wss.on('error', (error) => {
	// WebSocket server error
})

server.listen(port, '0.0.0.0', () => {
	// WebSocket server running
})

// Graceful shutdown
const gracefulShutdown = (signal: string) => {
	// Stop all heartbeats first
	for (const [ws, connectionInfo] of connections.entries()) {
		stopHeartbeat(ws)
		try {
			cleanupElevenLabsWebSocket(connectionInfo.contextId)
		} catch (error) {}
	}

	// Clear connections
	connections.clear()

	// Close WebSocket server
	wss.close((err) => {
		// Close HTTP server
		server.close((err) => {
			process.exit(0)
		})
	})

	// Force exit after 5 seconds if graceful shutdown fails
	setTimeout(() => {
		process.exit(1)
	}, 5000)
}

process.on('SIGINT', () => gracefulShutdown('SIGINT'))
process.on('SIGTERM', () => gracefulShutdown('SIGTERM'))