type PartySocketEvent = {
	type: string;
	[key: string]: unknown;
};

type WebSocketLike = {
	send: (data: string) => void;
	close?: (code?: number, reason?: string) => void;
};

const partySockets = new Map<string, Map<string, Set<WebSocketLike>>>();
const userSockets = new Map<string, Set<WebSocketLike>>();

function getPartyUserSockets(partyId: string, userId: string) {
	const partyMap = partySockets.get(partyId);
	if (!partyMap) return null;
	return partyMap.get(userId) ?? null;
}

export function registerPartySocket(
	partyId: string,
	userId: string,
	ws: WebSocketLike,
) {
	let partyMap = partySockets.get(partyId);
	if (!partyMap) {
		partyMap = new Map();
		partySockets.set(partyId, partyMap);
	}

	let userSockets = partyMap.get(userId);
	if (!userSockets) {
		userSockets = new Set();
		partyMap.set(userId, userSockets);
	}

	userSockets.add(ws);
}

export function unregisterPartySocket(
	partyId: string,
	userId: string,
	ws: WebSocketLike,
) {
	const partyMap = partySockets.get(partyId);
	if (!partyMap) return;

	const userSockets = partyMap.get(userId);
	if (!userSockets) return;

	userSockets.delete(ws);

	if (userSockets.size === 0) {
		partyMap.delete(userId);
	}

	if (partyMap.size === 0) {
		partySockets.delete(partyId);
	}
}

export function registerUserSocket(userId: string, ws: WebSocketLike) {
	let sockets = userSockets.get(userId);
	if (!sockets) {
		sockets = new Set();
		userSockets.set(userId, sockets);
	}

	sockets.add(ws);
}

export function unregisterUserSocket(userId: string, ws: WebSocketLike) {
	const sockets = userSockets.get(userId);
	if (!sockets) return;

	sockets.delete(ws);

	if (sockets.size === 0) {
		userSockets.delete(userId);
	}
}

export function unregisterUserSocketFromAllParties(
	userId: string,
	ws: WebSocketLike,
) {
	for (const [partyId, partyMap] of partySockets) {
		const userSockets = partyMap.get(userId);
		if (!userSockets) continue;
		userSockets.delete(ws);
		if (userSockets.size === 0) {
			partyMap.delete(userId);
		}
		if (partyMap.size === 0) {
			partySockets.delete(partyId);
		}
	}
}

export function broadcastPartyEvent(partyId: string, event: PartySocketEvent) {
	const partyMap = partySockets.get(partyId);
	if (!partyMap) return;

	const payload = JSON.stringify(event);
	for (const userSockets of partyMap.values()) {
		for (const ws of userSockets) {
			ws.send(payload);
		}
	}
}

export function sendPartyEventToUser(
	partyId: string,
	userId: string,
	event: PartySocketEvent,
) {
	const userSockets = getPartyUserSockets(partyId, userId);
	if (!userSockets) return;

	const payload = JSON.stringify(event);
	for (const ws of userSockets) {
		ws.send(payload);
	}
}

export function sendDirectEventToUser(userId: string, event: PartySocketEvent) {
	const sockets = userSockets.get(userId);
	if (!sockets) return;

	const payload = JSON.stringify(event);
	for (const ws of sockets) {
		ws.send(payload);
	}
}

export function reassignUserSocketsToParty(
	userId: string,
	partyId: string | null,
) {
	for (const [existingPartyId, partyMap] of partySockets) {
		if (!partyMap.has(userId)) continue;
		partyMap.delete(userId);
		if (partyMap.size === 0) {
			partySockets.delete(existingPartyId);
		}
	}

	if (!partyId) return;
	const sockets = userSockets.get(userId);
	if (!sockets) return;

	let partyMap = partySockets.get(partyId);
	if (!partyMap) {
		partyMap = new Map();
		partySockets.set(partyId, partyMap);
	}

	partyMap.set(userId, new Set(sockets));
}