fenghuo/apps/backend/src/socket.ts

import { createBunWebSocket } from 'hono/bun';
import type { ServerWebSocket } from 'bun';
import { WSContext } from 'hono/ws';
import { MessageType } from '../../../packages/client/src/websocket/type';

// 定义消息类型接口
interface WSMessage {
	type: 'join' | 'leave' | 'message';
	roomId: string;
	data: any;
}

// 创建 WebSocket 实例并指定泛型类型
const { upgradeWebSocket, websocket } = createBunWebSocket<ServerWebSocket>();

// 存储房间和连接的映射关系
const rooms = new Map<string, Set<WSContext<ServerWebSocket>>>();

// WebSocket 处理器
const wsHandler = upgradeWebSocket((c) => {
	// 存储当前连接加入的房间
	const joinedRooms = new Set<string>();

	return {
		onOpen(_, ws) {
			console.log('WebSocket 连接已建立');
		},
		onMessage(message, ws) {
			try {
				const parsedMessage: WSMessage = JSON.parse(message.data);
				console.log('收到消息:', parsedMessage);

				switch (parsedMessage.type) {
					case 'join':
						// 加入房间
						if (!rooms.has(parsedMessage.roomId)) {
							rooms.set(parsedMessage.roomId, new Set());
						}
						rooms.get(parsedMessage.roomId)?.add(ws);
						joinedRooms.add(parsedMessage.roomId);

						// 发送加入成功消息
						ws.send(
							JSON.stringify({
								type: 'system',
								data: {
									text: `成功加入房间 ${parsedMessage.roomId}`,
									type: MessageType.TEXT,
								},
								roomId: parsedMessage.roomId,
							}),
						);
						break;

					case 'leave':
						// 离开房间
						rooms.get(parsedMessage.roomId)?.delete(ws);
						joinedRooms.delete(parsedMessage.roomId);
						if (rooms.get(parsedMessage.roomId)?.size === 0) {
							rooms.delete(parsedMessage.roomId);
						}
						break;

					case 'message':
						// 在指定房间内广播消息
						const room = rooms.get(parsedMessage.roomId);
						if (room) {
							const messageToSend = {
								type: 'message',
								data: parsedMessage.data,
								roomId: parsedMessage.roomId,
							};

							for (const conn of room) {
								if (conn.readyState === WebSocket.OPEN) {
									conn.send(JSON.stringify(messageToSend));
								}
							}
						}
						break;
				}
			} catch (error) {
				console.error('处理消息时出错:', error);
				ws.send(
					JSON.stringify({
						type: 'error',
						data: {
							text: '消息处理失败',
							type: MessageType.TEXT,
						},
					}),
				);
			}
		},
		onClose(_, ws) {
			console.log('WebSocket 连接已关闭');
			// 清理所有加入的房间
			for (const roomId of joinedRooms) {
				rooms.get(roomId)?.delete(ws);
				if (rooms.get(roomId)?.size === 0) {
					rooms.delete(roomId);
				}
			}
		},
		onError(error, ws) {
			console.error('WebSocket 错误:', error);
			// 清理所有加入的房间
			for (const roomId of joinedRooms) {
				rooms.get(roomId)?.delete(ws);
				if (rooms.get(roomId)?.size === 0) {
					rooms.delete(roomId);
				}
			}
		},
	};
});

export const wsConfig = {
	websocket,
};

export { wsHandler, rooms };