import { useLocalStorage } from 'foxact/use-local-storage' import { useCallback, useEffect, useRef } from 'react' import { mutate, type MutatorCallback } from 'swr' import useSWRSubscription from 'swr/subscription' import { type Message, type MihomoWebSocket } from 'tauri-plugin-mihomo-api' export const RECONNECT_DELAY_MS = 100 type NextFn = (error?: any, data?: T | MutatorCallback) => void interface HandlerContext { next: NextFn scheduleReconnect: () => Promise isMounted: () => boolean } interface HandlerResult { handleMessage: (data: string) => void onConnected?: (ws: MihomoWebSocket) => Promise | void cleanup?: () => void } interface UseMihomoWsSubscriptionOptions { storageKey: string buildSubscriptKey: (date: number) => string | null fallbackData: T connect: () => Promise keepPreviousData?: boolean setupHandlers: (ctx: HandlerContext) => HandlerResult } export const useMihomoWsSubscription = ( options: UseMihomoWsSubscriptionOptions, ) => { const { storageKey, buildSubscriptKey, fallbackData, connect, keepPreviousData = true, setupHandlers, } = options // eslint-disable-next-line @eslint-react/purity const [date, setDate] = useLocalStorage(storageKey, Date.now()) const subscriptKey = buildSubscriptKey(date) const subscriptionCacheKey = subscriptKey ? `$sub$${subscriptKey}` : null const wsRef = useRef(null) const wsFirstConnectionRef = useRef(true) const timeoutRef = useRef | null>(null) const response = useSWRSubscription( subscriptKey, (_key, { next }) => { let isMounted = true const clearReconnectTimer = () => { if (timeoutRef.current) { clearTimeout(timeoutRef.current) timeoutRef.current = null } } const closeSocket = async () => { if (wsRef.current) { await wsRef.current.close() wsRef.current = null } } const scheduleReconnect = async () => { if (!isMounted) return clearReconnectTimer() await closeSocket() if (!isMounted) return timeoutRef.current = setTimeout(connectWs, RECONNECT_DELAY_MS) } const { handleMessage: handleTextMessage, onConnected, cleanup, } = setupHandlers({ next, scheduleReconnect, isMounted: () => isMounted, }) const cleanupAll = () => { clearReconnectTimer() cleanup?.() void closeSocket() } const handleMessage = (msg: Message) => { if (msg.type !== 'Text') return handleTextMessage(msg.data) } async function connectWs() { try { const ws_ = await connect() if (!isMounted) { await ws_.close() return } wsRef.current = ws_ clearReconnectTimer() if (onConnected) { await onConnected(ws_) if (!isMounted) { await ws_.close() return } } ws_.addListener(handleMessage) } catch (ignoreError) { if (!wsRef.current && isMounted) { timeoutRef.current = setTimeout(connectWs, RECONNECT_DELAY_MS) } } } if (wsFirstConnectionRef.current || !wsRef.current) { wsFirstConnectionRef.current = false cleanupAll() void connectWs() } return () => { isMounted = false wsFirstConnectionRef.current = true cleanupAll() } }, { fallbackData, keepPreviousData, }, ) useEffect(() => { if (subscriptionCacheKey) { mutate(subscriptionCacheKey) } }, [subscriptionCacheKey]) const refresh = useCallback(() => { setDate(Date.now()) }, [setDate]) return { response, refresh, subscriptionCacheKey, wsRef } }