server.ts

import { createServer } from "http";
import { Server } from "socket.io";
import { Message } from "@/lib/socket";

interface ThreadMessage extends Message {
  parentId?: string;
  reactions?: Record<string, string[]>; // emoji -> array of userIds
}

const httpServer = createServer();
const io = new Server(httpServer, {
  cors: {
    origin: "http://localhost:3000",
    methods: ["GET", "POST"],
  },
});

// Store messages by channel
const messagesByChannel: Record<string, ThreadMessage[]> = {
  general: [],
  random: [],
  introductions: [],
};

// Store direct messages between pairs of users
// Key format: "dm:user1:user2" where user1 and user2 are sorted alphabetically
const directMessages: Record<string, ThreadMessage[]> = {};

// Store self-messages for each user
// Key format: "self:username"
const selfMessages: Record<string, ThreadMessage[]> = {};

// Store thread replies
const threadReplies: Record<string, ThreadMessage[]> = {};

// Track user socket connections
const userSockets = new Map<string, string>(); // username -> socketId
const socketToUser = new Map<string, string>(); // socketId -> username
const users = new Set<string>();

// Helper function to create a consistent DM channel ID for any two users
function getDMChannelId(user1: string, user2: string): string {
  if (user1 === user2) {
    return `self:${user1}`;
  }
  const sortedUsers = [user1, user2].sort();
  return `dm:${sortedUsers[0]}:${sortedUsers[1]}`;
}

io.on("connection", (socket) => {
  console.log("Client connected:", socket.id);

  socket.on("userJoined", (username: string) => {
    console.log(`User joined: ${username} with socket ${socket.id}`);
    users.add(username);
    userSockets.set(username, socket.id);
    socketToUser.set(socket.id, username);
    io.emit("users", Array.from(users));
  });

  socket.on("joinChannel", (channelId: string) => {
    const currentUser = socketToUser.get(socket.id);
    console.log(`Client ${socket.id} (${currentUser}) joining channel:`, channelId);
    
    if (!currentUser) {
      console.log("No user found for socket");
      return;
    }

    if (channelId.startsWith("dm:") || channelId.startsWith("self:")) {
      // Handle self-messages
      if (channelId === `self:${currentUser}` || channelId === `dm:${currentUser}`) {
        const selfChannelId = `self:${currentUser}`;
        const messages = selfMessages[selfChannelId] || [];
        console.log(`Sending self-message history for ${currentUser}:`, messages);
        socket.emit("channelHistory", messages);
        socket.emit("channelUpdate", selfChannelId);
        return;
      }

      // Handle DMs
      const parts = channelId.split(":");
      if (parts.length !== 3) {
        const targetUser = parts[1];
        // Create the proper DM channel ID for these two users
        const dmChannelId = getDMChannelId(currentUser, targetUser);
        console.log(`Converting ${channelId} to ${dmChannelId}`);
        
        socket.join(dmChannelId);
        const messages = directMessages[dmChannelId] || [];
        console.log(`Sending DM history for ${dmChannelId}:`, messages);
        socket.emit("channelHistory", messages);
        
        // Update the client with the correct channel ID
        socket.emit("channelUpdate", dmChannelId);
      } else {
        // Already in correct format
        socket.join(channelId);
        const messages = directMessages[channelId] || [];
        console.log(`Sending DM history for ${channelId}:`, messages);
        socket.emit("channelHistory", messages);
      }
    } else {
      // For regular channels
      socket.join(channelId);
      const messages = messagesByChannel[channelId] || [];
      // Add thread replies count to each message
      const messagesWithReplies = messages.map(message => ({
        ...message,
        threadReplies: threadReplies[message.id] || []
      }));
      socket.emit("channelHistory", messagesWithReplies);
    }
  });

  socket.on("leaveChannel", (channelId: string) => {
    console.log(`Client ${socket.id} leaving channel:`, channelId);
    socket.leave(channelId);
  });

  socket.on("message", (message: ThreadMessage) => {
    const senderId = socketToUser.get(socket.id);
    console.log(`Message received from ${senderId}:`, message);

    if (!senderId) {
      console.log("Sender not found");
      return;
    }

    // If this is a thread reply, store it in threadReplies
    if (message.parentId) {
      if (!threadReplies[message.parentId]) {
        threadReplies[message.parentId] = [];
      }
      threadReplies[message.parentId].push(message);
      
      // Find the parent message's channel and broadcast to all clients in that channel
      let parentChannel = message.channelId;
      for (const [channel, messages] of Object.entries(messagesByChannel)) {
        if (messages.some(m => m.id === message.parentId)) {
          parentChannel = channel;
          break;
        }
      }
      
      // Broadcast the thread reply to all clients in the channel
      io.to(parentChannel).emit("message", message);
      return;
    }

    if (message.channelId.startsWith("self:")) {
      // Handle self-message
      const selfChannelId = `self:${senderId}`;
      if (!selfMessages[selfChannelId]) {
        selfMessages[selfChannelId] = [];
      }
      selfMessages[selfChannelId].push(message);
      socket.emit("message", message);
      return;
    }

    if (message.channelId.startsWith("dm:")) {
      // Handle direct message
      const parts = message.channelId.split(":");
      let dmChannelId = message.channelId;

      if (parts.length !== 3) {
        // Convert to proper format if needed
        const targetUser = parts[1];
        dmChannelId = getDMChannelId(senderId, targetUser);
        console.log(`Converting message channel from ${message.channelId} to ${dmChannelId}`);
        message.channelId = dmChannelId;
      }

      // Check if this is a self-message
      if (dmChannelId.startsWith("self:")) {
        if (!selfMessages[dmChannelId]) {
          selfMessages[dmChannelId] = [];
        }
        selfMessages[dmChannelId].push(message);
        socket.emit("message", message);
        return;
      }

      // Store the message
      if (!directMessages[dmChannelId]) {
        directMessages[dmChannelId] = [];
      }
      directMessages[dmChannelId].push(message);

      // Get recipient's socket ID
      const [_, user1, user2] = dmChannelId.split(":");
      const recipientId = user1 === senderId ? user2 : user1;
      const recipientSocketId = userSockets.get(recipientId);
      
      console.log(`Sending DM from ${senderId} to ${recipientId} in channel ${dmChannelId}`);

      // Send to both users through the room only
      io.to(dmChannelId).emit("message", message);
    } else {
      // Handle regular channel message
      if (!messagesByChannel[message.channelId]) {
        messagesByChannel[message.channelId] = [];
      }
      messagesByChannel[message.channelId].push(message);
      // For regular channel messages, broadcast to everyone in the channel
      io.to(message.channelId).emit("message", {
        ...message,
        threadReplies: threadReplies[message.id] || []
      });
    }
  });

  socket.on("showThread", (parentMessage: ThreadMessage) => {
    // Get thread replies for this message
    const replies = threadReplies[parentMessage.id] || [];
    // Emit both the parent message and its replies
    socket.emit("threadSelected", {
      ...parentMessage,
      replies
    });
  });

  socket.on("getThreadReplies", (parentId: string) => {
    console.log(`Getting thread replies for message: ${parentId}`);
    const replies = threadReplies[parentId] || [];
    socket.emit("threadReplies", replies);
  });

  socket.on("addReaction", ({ messageId, emoji, channelId }: { messageId: string; emoji: string; channelId: string }) => {
    const userId = socketToUser.get(socket.id);
    if (!userId) return;

    // Find the message in the appropriate store
    let message: ThreadMessage | undefined;
    let messages: ThreadMessage[] | undefined;

    if (channelId.startsWith("dm:")) {
      messages = directMessages[channelId];
    } else if (channelId.startsWith("self:")) {
      messages = selfMessages[channelId];
    } else {
      messages = messagesByChannel[channelId];
    }

    if (!messages) return;

    // Find and update the message
    const messageIndex = messages.findIndex(m => m.id === messageId);
    if (messageIndex === -1) return;

    message = messages[messageIndex];
    if (!message.reactions) {
      message.reactions = {};
    }

    // Add or remove the reaction
    if (!message.reactions[emoji]) {
      message.reactions[emoji] = [];
    }

    const userIndex = message.reactions[emoji].indexOf(userId);
    if (userIndex === -1) {
      // Add reaction
      message.reactions[emoji].push(userId);
    } else {
      // Remove reaction
      message.reactions[emoji].splice(userIndex, 1);
      // Clean up empty reactions
      if (message.reactions[emoji].length === 0) {
        delete message.reactions[emoji];
      }
    }

    // Broadcast the updated message
    io.to(channelId).emit("message", {
      ...message,
      threadReplies: threadReplies[message.id] || []
    });
  });

  socket.on("disconnect", () => {
    console.log("Client disconnected:", socket.id);
    const username = socketToUser.get(socket.id);
    if (username) {
      users.delete(username);
      userSockets.delete(username);
      socketToUser.delete(socket.id);
      io.emit("users", Array.from(users));
    }
  });
});

const PORT = 3001;
httpServer.listen(PORT, () => {
  console.log(`Socket.IO server running on port ${PORT}`);
});