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: string }
	| { Question: DeviceQuestionData }
	| "Results"
	| { Error: string };

type ApiMessage =
	| { type: "device_status_request"; 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"; username: string };

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

const sockets = new Map<string, Socket>();
const socketIds = new WeakMap<Socket, string>();
const API_SOCKET_URL = "ws://localhost:4000/api/dev-socket/ws";
let apiSocket: WebSocket | null = null;
let apiReconnectTimer: ReturnType<typeof setTimeout> | null = null;
const pendingDeviceStatus = new Set<string>();

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 || apiSocket.readyState !== WebSocket.OPEN) return false;
	console.log("API send", message.type, message.deviceId);
	apiSocket.send(JSON.stringify(message));
	return true;
}

function requestDeviceStatus(deviceId: string) {
	pendingDeviceStatus.add(deviceId);
	if (sendApiMessage({ type: "device_status_request", deviceId })) {
		console.log("Requested device status", deviceId);
		pendingDeviceStatus.delete(deviceId);
		return;
	}
	console.log("Queued device status request", deviceId);
}

function flushPendingDeviceStatus() {
	for (const deviceId of pendingDeviceStatus) {
		if (
			sockets.has(deviceId) &&
			sendApiMessage({ type: "device_status_request", deviceId })
		) {
			console.log("Flushed device status request", deviceId);
			pendingDeviceStatus.delete(deviceId);
		}
	}
}

function disconnectDeviceClients(reason: string) {
	console.log("Disconnecting device clients", reason, sockets.size);
	for (const socket of sockets.values()) {
		socket.end();
	}
	sockets.clear();
	pendingDeviceStatus.clear();
}

function scheduleApiReconnect() {
	if (apiReconnectTimer) return;
	apiReconnectTimer = setTimeout(() => {
		apiReconnectTimer = null;
		connectApiSocket();
	}, 500);
}

function connectApiSocket() {
	if (
		apiSocket?.readyState === WebSocket.OPEN ||
		apiSocket?.readyState === WebSocket.CONNECTING
	) {
		return;
	}

	console.log("Connecting to API device socket");
	apiSocket = new WebSocket(API_SOCKET_URL);
	apiSocket.onmessage = handleApiMessage;
	apiSocket.onerror = (error) => {
		console.error("API device socket error", error);
	};
	apiSocket.onclose = () => {
		console.log("API device socket closed; reconnecting");
		apiSocket = null;
		disconnectDeviceClients("api socket closed");
		scheduleApiReconnect();
	};
	apiSocket.onopen = () => {
		console.log("Connected to API device socket");
		flushPendingDeviceStatus();
	};
}

function toDeviceQuestionData(quizData: QuizState): DeviceQuestionData | null {
	if (!quizData.currentQuestion) return null;
	if (quizData.status === "results") 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);
				console.log(
					"Requesting device status",
					data.DeviceId,
					"apiState",
					apiSocket?.readyState,
				);
				requestDeviceStatus(data.DeviceId);
				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);
				pendingDeviceStatus.delete(deviceId);
			}
		},
	},
});

function handleApiMessage(e: MessageEvent) {
	let message: ApiEnvelope;
	try {
		message = JSON.parse(e.data) as ApiEnvelope;
	} catch {
		return;
	}

	console.log("API recv", message.type);
	if (message.type !== "device_event") return;
	const socket = sockets.get(message.deviceId);
	if (!socket) {
		console.log("No TCP socket for API event", message.deviceId);
		return;
	}
	const event = message.event as PartySocketEvent;
	console.log("API device event", message.deviceId, event.type);
	if (event.type === "device_connect_required") {
		console.log("Writing connect prompt", message.deviceId);
		writeProxyOutput(socket, { ConnectPrompt: message.deviceId });
		return;
	}

	if (event.type === "device_connected") {
		console.log("Writing waiting-for-party", message.deviceId);
		writeProxyOutput(socket, { WaitingForParty: event.username });
		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;
		if (quizData.status === "results") {
			writeProxyOutput(socket, "Results");
			return;
		}
		const question = toDeviceQuestionData(quizData);
		if (question) {
			writeProxyOutput(socket, { Question: question });
		}
		return;
	}

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

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

connectApiSocket();

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