socketServer.ts

import { Server } from 'socket.io';
import { createServer } from 'http';

const httpServer = createServer();
const io = new Server(httpServer, {
  cors: {
    origin: [process.env.NEXT_PUBLIC_APP_URL || 'http://localhost:3000'],
    methods: ['GET', 'POST']
  },
  allowEIO3: true, // Allow older Engine.IO versions (if necessary)
  transports: ['websocket', 'polling'],
});

const generatingLLMChats = new Map()
const generatingWorkspaces = new Map()

io.on('connection', (socket) => {
  console.log(`Client connected: ${socket.id}`);

  socket.on('join-conversation', (conversationId) => {
    console.log(`User joined conversation ${conversationId}`);
    socket.join(conversationId);
  });

  socket.on('leave-conversation', (conversationId) => {
    console.log(`User left conversation ${conversationId}`);
    socket.leave(conversationId);
  });

  socket.on('new-message', (message) => {
    socket.to(message.conversationId).emit('new-message', message);
  });

  socket.on('generation-status-update', (llmChatId, isGenerating) => {
    console.log(`Generation status update for ${llmChatId}: ${isGenerating}`);
    if (isGenerating) {
      generatingLLMChats.set(llmChatId, true);
    } else {
      generatingLLMChats.delete(llmChatId);
    }
    socket.to(llmChatId).emit('generation-status-update', isGenerating);
  });

  socket.on('join-llm-chat', (llmChatId) => {
    socket.join(llmChatId);

    if (generatingLLMChats.has(llmChatId)) {
      console.log(`User ${socket.id} joined LLM chat ${llmChatId} and there's a generation in progress`);
      socket.emit('generation-status-update', true);
    } else {
      console.log(`User ${socket.id} joined LLM chat ${llmChatId} and there's no generation in progress`);
    }
  });

  socket.on('leave-llm-chat', (llmChatId) => {
    console.log(`User left LLM chat ${llmChatId}`);
    socket.leave(llmChatId);
  });

  socket.on('new-llm-message', (message) => {
    // console.log(`New LLM message received: ${message.content}`);
    console.log(`New LLM message received from ${socket.id}`);
    socket.to(message.llmChatId).emit('new-llm-message', message);
  });

  socket.on('switch-llm-model', (llmChatId, model) => {
    console.log(`User switched LLM model ${model} for ${llmChatId}`);
    socket.to(llmChatId).emit('switch-llm-model', model);
  });

  socket.on('workspace-generation-status-update', (workspaceId, isGenerating) => {
    console.log(`Generation status update for ${workspaceId}: ${isGenerating}`);
    if (isGenerating) {
      generatingWorkspaces.set(workspaceId, true);
    } else {
      generatingWorkspaces.delete(workspaceId);
    }
    socket.to(workspaceId).emit('workspace-generation-status-update', isGenerating);
  });

  socket.on('join-workspace', (workspaceId) => {
    socket.join(workspaceId);

    if (generatingWorkspaces.has(workspaceId)) {
      console.log(`User ${socket.id} joined workspace ${workspaceId} and there's a generation in progress`);
      socket.emit('workspace-generation-status-update', true);
    } else {
      console.log(`User ${socket.id} joined workspace ${workspaceId} and there's no generation in progress`);
    }
  });

  socket.on('leave-workspace', (workspaceId) => {
    console.log(`User left workspace ${workspaceId}`);
    socket.leave(workspaceId);
  });

  socket.on('new-workspace-message', (message) => {
    console.log(`New workspace message received from ${socket.id}`);
    socket.to(message.workspaceId).emit('new-workspace-message', message);
  });

  socket.on('switch-workspace-model', (workspaceId, model) => {
    console.log(`User switched workspace model ${model} for ${workspaceId}`);
    socket.to(workspaceId).emit('switch-workspace-model', model);
  });

  socket.on('disconnect', () => {
    console.log(`Client disconnected: ${socket.id}`);
  });
});

const PORT = parseInt(process.env.SOCKET_PORT || '4000', 10);
httpServer.listen(PORT, '0.0.0.0', () => {
  console.log(`Socket.io server running on port ${PORT}`);
});