server/index.js

const express = require('express');
const path = require('node:path');
const { WebSocketServer, WebSocket } = require('ws');
const http = require('node:http');
require('dotenv').config();

const app = express();
const server = http.createServer(app);
const wss = new WebSocketServer({ server });

// Serve static files from the React app build directory
app.use(express.static(path.join(__dirname, '../build')));

// Make API key available to the WebSocket server but not to the client
const GEMINI_API_KEY = process.env.GEMINI_API_KEY;

if (!GEMINI_API_KEY) {
  console.error('GEMINI_API_KEY environment variable is not set!');
  process.exit(1);
}

// Create a WebSocket connection to Gemini for each client
const createGeminiWebSocket = (clientWs) => {
  const geminiWs = new WebSocket(
    `wss://generativelanguage.googleapis.com/ws/google.ai.generativelanguage.v1alpha.GenerativeService.BidiGenerateContent?key=${GEMINI_API_KEY}`
  );

  // Set up event handlers before connecting
  geminiWs.on('open', () => {
    console.log('Connected to Gemini API');
    // If there's a pending setup message, send it now
    if (geminiWs.pendingSetup) {
      console.log('Sending pending setup:', geminiWs.pendingSetup);
      geminiWs.send(JSON.stringify(geminiWs.pendingSetup));
      geminiWs.pendingSetup = null;
    }
  });

  geminiWs.on('message', (data) => {
    try {
      // Convert the message to a Blob before sending to client
      const message = data.toString();
      console.log('Received from Gemini:', message);
      
      // Create a Blob from the message
      const blob = Buffer.from(message);
      clientWs.send(blob, { binary: true });
    } catch (error) {
      console.error('Error handling Gemini message:', error);
    }
  });

  geminiWs.on('error', (error) => {
    console.error('Gemini WebSocket error:', error);
  });

  geminiWs.on('close', (code, reason) => {
    console.log('Gemini WebSocket closed:', code, reason.toString());
  });

  return geminiWs;
};

wss.on('connection', (ws) => {
  console.log('Client connected');
  let geminiWs = null;

  ws.on('message', async (message) => {
    try {
      const data = JSON.parse(message);
      console.log('Received from client:', data);

      // Initialize Gemini connection when receiving setup message
      if (data.setup) {
        console.log('Initializing Gemini connection with config:', data.setup);
        geminiWs = createGeminiWebSocket(ws);
        
        // Store setup message to send once connection is established
        if (geminiWs.readyState !== WebSocket.OPEN) {
          geminiWs.pendingSetup = data;
        } else {
          geminiWs.send(JSON.stringify(data));
        }
        return;
      }

      // Forward message to Gemini if connection exists
      if (geminiWs && geminiWs.readyState === WebSocket.OPEN) {
        console.log('Forwarding to Gemini:', data);
        geminiWs.send(JSON.stringify(data));
      } else if (geminiWs) {
        console.log('Waiting for Gemini connection to be ready...');
      } else {
        console.error('No Gemini connection established');
      }
    } catch (error) {
      console.error('Error processing message:', error);
    }
  });

  ws.on('close', () => {
    console.log('Client disconnected');
    if (geminiWs) {
      geminiWs.close();
    }
  });
});

// Handle any remaining requests by returning the React app
app.get('*', (req, res) => {
  res.sendFile(path.join(__dirname, '../build', 'index.html'));
});

const PORT = process.env.PORT || 3001;
server.listen(PORT, () => {
  console.log(`Server is running on port ${PORT}`);
});