import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import type {
	PartySocketEvent,
	PartySocketOutgoing,
} from "../../../api/src/party-types";

type Handler = (event: PartySocketEvent) => void;

const PING_INTERVAL_MS = 30_000;
const RECONNECT_BASE_MS = 1_000;
const RECONNECT_MAX_MS = 30_000;

export function usePartySocket({
	apiUrl,
	onMessage,
}: {
	apiUrl: string | null;
	onMessage: Handler | null;
}) {
	const [connectionState, setConnectionState] = useState<
		"disconnected" | "connecting" | "connected" | "reconnecting"
	>("disconnected");

	const wsRef = useRef<WebSocket | null>(null);
	const pingTimerRef = useRef<ReturnType<typeof setInterval> | null>(null);
	const reconnectTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
	const reconnectAttemptRef = useRef(0);
	const handlerRef = useRef(onMessage);

	const send = useCallback((message: PartySocketOutgoing) => {
		const ws = wsRef.current;
		if (!ws || ws.readyState !== WebSocket.OPEN) return;
		ws.send(JSON.stringify(message));
	}, []);

	useEffect(() => {
		handlerRef.current = onMessage;
	}, [onMessage]);

	const setupWs = useCallback(
		(ws: WebSocket) => {
			ws.onopen = () => {
				reconnectAttemptRef.current = 0;
				setConnectionState("connected");
				pingTimerRef.current = setInterval(() => {
					if (ws.readyState === WebSocket.OPEN) {
						ws.send(JSON.stringify({ type: "ping" }));
					}
				}, PING_INTERVAL_MS);
			};

			ws.onmessage = (event) => {
				const parsed = JSON.parse(event.data) as PartySocketEvent;
				handlerRef.current?.(parsed);
			};

			ws.onclose = () => {
				if (pingTimerRef.current) {
					clearInterval(pingTimerRef.current);
					pingTimerRef.current = null;
				}
				wsRef.current = null;
				setConnectionState("reconnecting");

				const delay = Math.min(
					RECONNECT_BASE_MS * 2 ** reconnectAttemptRef.current,
					RECONNECT_MAX_MS,
				);
				reconnectAttemptRef.current++;
				reconnectTimerRef.current = setTimeout(() => {
					if (!apiUrl) return;
					const protocol = apiUrl.startsWith("https") ? "wss" : "ws";
					const newWs = new WebSocket(
						`${protocol}://${apiUrl.replace(/https?:\/\//, "")}/api/party-socket/ws`,
					);
					wsRef.current = newWs;
					setupWs(newWs);
				}, delay);
			};
		},
		[apiUrl],
	);

	useEffect(() => {
		if (!apiUrl) {
			if (wsRef.current) {
				wsRef.current.close();
				wsRef.current = null;
			}
			if (pingTimerRef.current) {
				clearInterval(pingTimerRef.current);
				pingTimerRef.current = null;
			}
			if (reconnectTimerRef.current) {
				clearTimeout(reconnectTimerRef.current);
				reconnectTimerRef.current = null;
			}
			setConnectionState("disconnected");
			reconnectAttemptRef.current = 0;
			return;
		}

		setConnectionState("connecting");
		const protocol = apiUrl.startsWith("https") ? "wss" : "ws";
		const ws = new WebSocket(
			`${protocol}://${apiUrl.replace(/https?:\/\//, "")}/api/party-socket/ws`,
		);
		wsRef.current = ws;
		setupWs(ws);

		return () => {
			ws.close();
			wsRef.current = null;
			if (pingTimerRef.current) {
				clearInterval(pingTimerRef.current);
				pingTimerRef.current = null;
			}
			if (reconnectTimerRef.current) {
				clearTimeout(reconnectTimerRef.current);
				reconnectTimerRef.current = null;
			}
		};
	}, [apiUrl, setupWs]);

	const state = useMemo(
		() => ({
			connectionState,
			isConnected: connectionState === "connected",
			isConnecting: connectionState === "connecting",
			isReconnecting: connectionState === "reconnecting",
			send,
		}),
		[connectionState, send],
	);

	return state;
}