import type { Socket } from "bun";

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

type DeviceMessage =
	| {
			DeviceId: string;
	  }
	| {
			QuizResponse: number;
	  };

type DeviceQuestionData = {
	text: string;
	points: number;
	index: number;
	q_type: "Choice" | { Numeric: { min: number; max: number } };
};

type ProxyOutput =
	| { ConnectPrompt: string }
	| "WaitingForParty"
	| { Question: DeviceQuestionData }
	| "Results"
	| { Error: string };

type ApiMessage =
	| { type: "device_connected"; deviceId: string }
	| { type: "device_message"; deviceId: string; payload: unknown };

type QuizQuestion =
	| {
			type: "choice";
			text: string;
			points: number;
	  }
	| {
			type: "numeric";
			text: string;
			points: number;
			range: { min: number; max: number };
	  };

type QuizState = {
	status: "running" | "results";
	questionIndex: number;
	currentQuestion: QuizQuestion | null;
};

type PartyStatusEvent = {
	type: "party_status";
	party: { data?: QuizState } | null;
};

type QuizStateEvent = {
	type: "quiz_state";
	quiz: QuizState;
};

type ErrorEvent = {
	type: "error";
	message: string;
};

type DeviceLifecycleEvent =
	| { type: "device_connect_required" }
	| { type: "device_connected" };

type PartySocketEvent =
	| PartyStatusEvent
	| QuizStateEvent
	| ErrorEvent
	| DeviceLifecycleEvent;

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);
	console.log("Registered", socket.remoteAddress, deviceId);
}

function writeProxyOutput(socket: Socket, output: ProxyOutput) {
	socket.write(`${JSON.stringify(output)}\n`);
}

function sendApiMessage(message: ApiMessage) {
	if (apiSocket.readyState !== WebSocket.OPEN) return false;
	apiSocket.send(JSON.stringify(message));
	return true;
}

function toDeviceQuestionData(quizData: QuizState): DeviceQuestionData | null {
	if (!quizData.currentQuestion) return null;
	const question = quizData.currentQuestion;
	const q_type =
		question.type === "choice"
			? "Choice"
			: { Numeric: { min: question.range.min, max: question.range.max } };

	return {
		text: question.text,
		points: question.points,
		index: quizData.questionIndex,
		q_type,
	};
}

const listener = Bun.listen({
	port: 7070,
	hostname: "0.0.0.0",
	socket: {
		open(socket) {
			socket.setKeepAlive(true);
			console.log("Connection", socket.remoteAddress, socket.remotePort);
		},
		data(socket, buf) {
			const raw = new TextDecoder().decode(buf).trim();
			let data: DeviceMessage;
			try {
				data = JSON.parse(raw);
			} catch {
				return;
			}
			console.log("Data", socket.remoteAddress, data);
			if (!data) return;

			if ("DeviceId" in data) {
				registerSocket(socket, data.DeviceId);
				if (
					!sendApiMessage({ type: "device_connected", deviceId: data.DeviceId })
				) {
					writeProxyOutput(socket, {
						Error: "API device socket not connected.",
					});
				}
				return;
			}
			if ("QuizResponse" in data) {
				const deviceId = socketDeviceId(socket);
				if (!deviceId) return;
				sendApiMessage({
					type: "device_message",
					deviceId,
					payload: { QuizResponse: data.QuizResponse },
				});
				return;
			}
		},
		close(socket) {
			console.log("Connection", socket.remoteAddress);
			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;
	const event = message.event as PartySocketEvent;
	if (event.type === "device_connect_required") {
		writeProxyOutput(socket, { ConnectPrompt: message.deviceId });
		return;
	}

	if (event.type === "device_connected") {
		writeProxyOutput(socket, "WaitingForParty");
		return;
	}

	if (event.type === "error") {
		writeProxyOutput(socket, { Error: event.message });
		return;
	}

	if (event.type === "party_status") {
		const quizData = event.party?.data ?? null;
		if (!quizData) return;
		const question = toDeviceQuestionData(quizData);
		if (question) {
			writeProxyOutput(socket, { Question: question });
		} else if (quizData.status === "results") {
			writeProxyOutput(socket, "Results");
		}
		return;
	}

	if (event.type === "quiz_state") {
		const question = toDeviceQuestionData(event.quiz);
		if (question) {
			writeProxyOutput(socket, { Question: question });
		} else if (event.quiz.status === "results") {
			writeProxyOutput(socket, "Results");
		}
		return;
	}

	writeProxyOutput(socket, { Error: "Unsupported proxy event." });
};

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

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

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