import { WebSocket as ReconnectingWebSocket } from "partysocket";
import { useCallback, useEffect, useRef } from "react";
import { toast } from "sonner";
import { useAppStore } from "../stores/app-store";

function getWsUrl(): string {
	const proto = location.protocol === "https:" ? "wss:" : "ws:";
	const params = new URLSearchParams(location.search);
	const token = params.get("token") || "";
	const q = token ? `?token=${encodeURIComponent(token)}` : "";
	return `${proto}//${location.host}/ws${q}`;
}

export function useWebSocket() {
	const wsRef = useRef<ReconnectingWebSocket | null>(null);

	const sendJSON = useCallback((obj: Record<string, unknown>) => {
		const ws = wsRef.current;
		if (ws?.readyState === WebSocket.OPEN) {
			ws.send(JSON.stringify(obj));
		}
	}, []);

	const sendBinary = useCallback((data: Int16Array) => {
		const ws = wsRef.current;
		if (ws?.readyState === WebSocket.OPEN) {
			ws.send(data.buffer as ArrayBuffer);
		}
	}, []);

	useEffect(() => {
		useAppStore.getState().setConnectionStatus("connecting");

		const ws = new ReconnectingWebSocket(getWsUrl(), undefined, {
			minReconnectionDelay: 1000,
			maxReconnectionDelay: 16000,
		});
		ws.binaryType = "arraybuffer";

		ws.onopen = () => {
			useAppStore.getState().setConnectionStatus("connected");
		};

		ws.onmessage = (e: MessageEvent) => {
			if (typeof e.data !== "string") return;
			try {
				const msg = JSON.parse(e.data);
				const store = useAppStore.getState();
				switch (msg.type) {
					case "partial":
						store.setPreview(msg.text || "", false);
						break;
					case "final":
						store.setPreview(msg.text || "", true);
						if (msg.text) store.addHistory(msg.text);
						break;
					case "pasted":
						toast.success("已粘贴");
						break;
					case "error":
						toast.error(msg.message || "错误");
						break;
				}
			} catch {
				// Ignore malformed messages
			}
		};

		ws.onclose = () => {
			const store = useAppStore.getState();
			store.setConnectionStatus("disconnected");
			if (store.recording) store.setRecording(false);
			if (store.pendingStart) store.setPendingStart(false);
		};

		wsRef.current = ws;

		return () => {
			ws.close();
			wsRef.current = null;
		};
	}, []);

	return { sendJSON, sendBinary };
}