index.ts

import { WebSocketServer, WebSocket } from 'ws';
import express from 'express';
import path from 'path';
import { createServer } from 'http';

const app = express();
const port = process.env.PORT || 8080;

const frontendPath = path.resolve(__dirname, '../../fe/dist');
app.use(express.static(frontendPath));
app.use((_req, res) => {
    res.sendFile(path.resolve(frontendPath, 'index.html'));
});

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

interface User {
    socket: WebSocket;
    room: string;
    username: string;
}

interface Message {
    message: string;
    username: string;
    timestamp: number;
}

let allSockets: User[] = [];
const roomMessages: Map<string, Message[]> = new Map();

function broadcast(room: string, payload: object, exclude?: WebSocket) {
    const data = JSON.stringify(payload);
    allSockets.forEach((user) => {
        if (user.room === room && user.socket !== exclude) {
            try {
                user.socket.send(data);
            } catch (err) {
                console.error(`Failed to send to user ${user.username}:`, err);
            }
        }
    });
}

function broadcastUsersList(room: string) {
    const usersInRoom = allSockets.filter(u => u.room === room).map(u => u.username);
    const uniqueUsers = Array.from(new Set(usersInRoom));
    broadcast(room, {
        type: 'users_list',
        payload: { users: uniqueUsers }
    });
}

wss.on('connection', (socket) => {
    socket.on('message', (data) => {
        try {
            const parsed = JSON.parse(data.toString());

            if (parsed.type === 'join') {
                // Prevent double-join on the same socket
                if (allSockets.some((x) => x.socket === socket)) return;

                const roomId = typeof parsed.payload?.roomId === 'string' ? parsed.payload.roomId.trim().slice(0, 50) : '';
                const username = typeof parsed.payload?.username === 'string' ? parsed.payload.username.trim().slice(0, 30) : '';
                if (!roomId || !username) return;

                allSockets.push({ socket, room: roomId, username });

                // Send history to the new user
                const history = roomMessages.get(roomId) || [];
                socket.send(JSON.stringify({
                    type: 'history',
                    payload: { messages: history },
                }));

                // Notify others in the room
                broadcast(roomId, {
                    type: 'system',
                    payload: { message: `${username} joined the room` },
                }, socket);
                
                broadcastUsersList(roomId);
            }

            if (parsed.type === 'chat') {
                const currentUser = allSockets.find((x) => x.socket === socket);
                if (!currentUser) return;

                const rawMessage = typeof parsed.payload?.message === 'string' ? parsed.payload.message.trim().slice(0, 2000) : '';
                if (!rawMessage) return;

                const msg: Message = {
                    message: rawMessage,
                    username: currentUser.username,
                    timestamp: Date.now(),
                };

                if (!roomMessages.has(currentUser.room)) {
                    roomMessages.set(currentUser.room, []);
                }
                roomMessages.get(currentUser.room)!.push(msg);

                broadcast(currentUser.room, { type: 'chat', payload: msg });
            }

            if (parsed.type === 'typing') {
                const currentUser = allSockets.find((x) => x.socket === socket);
                if (!currentUser) return;

                broadcast(currentUser.room, {
                    type: 'typing',
                    payload: { username: currentUser.username, isTyping: !!parsed.payload?.isTyping },
                }, socket);
            }
        } catch (e) {
            console.error('Error processing message:', e);
        }
    });

    socket.on('close', () => {
        const user = allSockets.find((x) => x.socket === socket);
        if (user) {
            allSockets = allSockets.filter((x) => x.socket !== socket);
            broadcast(user.room, {
                type: 'system',
                payload: { message: `${user.username} left the room` },
            });
            broadcastUsersList(user.room);
        }
    });
});

server.listen(port, () => {
    console.log(`Server running on port ${port}`);
});