import type { Socket } from "bun";

type ApiEnvelope =
	| { type: "hello" }
	| { type: "device_event"; deviceId: string; event: unknown };

const sockets = new Map<string, Socket>();
const socketIds = new WeakMap<Socket, string>();
const apiSocket = new WebSocket("ws://localhost:4000/api/dev-socket/ws");

function socketDeviceId(socket: Socket) {
	return socketIds.get(socket);
}

function registerSocket(socket: Socket, deviceId: string) {
	const existing = sockets.get(deviceId);
	if (existing && existing !== socket) existing.end();
	sockets.set(deviceId, socket);
	socketIds.set(socket, deviceId);
}

const listener = Bun.listen({
	port: 7070,
	hostname: "0.0.0.0",
	socket: {
		open(socket) {
			socket.setKeepAlive(true);
		},
		data(socket, buf) {
			const raw = new TextDecoder().decode(buf).trim();
			if (!raw) return;

			const currentDeviceId = socketDeviceId(socket);
			if (!currentDeviceId) {
				registerSocket(socket, raw);
				return;
			}

			apiSocket?.send(
				JSON.stringify({
					type: "device_message",
					deviceId: currentDeviceId,
					payload: raw,
				}),
			);
		},
		close(socket) {
			const deviceId = socketDeviceId(socket);
			if (deviceId && sockets.get(deviceId) === socket) {
				sockets.delete(deviceId);
			}
		},
	},
});

apiSocket.onmessage = (e) => {
	let message: ApiEnvelope;
	try {
		message = JSON.parse(e.data) as ApiEnvelope;
	} catch {
		return;
	}

	if (message.type !== "device_event") return;
	const socket = sockets.get(message.deviceId);
	if (!socket) return;
	socket.write(`${JSON.stringify(message.event)}\n`);
};

apiSocket.onerror = (error) => {
	console.error(error);
};

apiSocket.onopen = () => {
	console.log("Connected to API device socket");
};

console.log(`Started on :${listener.port}`);