"use client"; import { createContext, type ReactNode, useCallback, useContext, useEffect, useState, } from "react"; import useWebSocket, { ReadyState } from "react-use-websocket"; import { type TelegramError, type WebSocketMessage, WebSocketMessageType, } from "@/lib/websocket-types"; import { useToast } from "./use-toast"; import { useDebounce } from "use-debounce"; import { getWsUrl } from "@/lib/api"; import { useParams } from "next/navigation"; const WS_URL = `${getWsUrl()}`; interface WebsocketContextType { sendMessage: (message: WebSocketMessage) => void; lastJsonMessage: WebSocketMessage | null; connectionStatus: string; isReady: boolean; accountDownloadSpeed: number; } const WebSocketContext = createContext( undefined, ); interface WebSocketProviderProps { children: ReactNode; } export const WebSocketProvider: React.FC = ({ children, }) => { const params = useParams(); const [isReady, setIsReady] = useState(false); const [accountDownloadSpeed, setAccountDownloadSpeed] = useState({ speed: 0, lastDownloadedSize: 0, lastTimestamp: 0, }); const { toast } = useToast(); const [debounceSpeed] = useDebounce(accountDownloadSpeed.speed, 300, { leading: true, maxWait: 1000, }); const { sendMessage, lastJsonMessage, readyState } = useWebSocket( `${WS_URL}?telegramId=${(params.accountId as string) ?? ""}`, { shouldReconnect: (closeEvent) => true, reconnectAttempts: 3, reconnectInterval: 3000, }, ); const connectionStatus = { [ReadyState.CONNECTING]: "Connecting", [ReadyState.OPEN]: "Open", [ReadyState.CLOSING]: "Closing", [ReadyState.CLOSED]: "Closed", [ReadyState.UNINSTANTIATED]: "Uninstantiated", }[readyState]; useEffect(() => { setIsReady(readyState === ReadyState.OPEN); }, [readyState]); useEffect(() => { if (lastJsonMessage !== null) { // console.log( // `🐄: ${JSON.stringify(lastJsonMessage, null, 2)}`, // ) try { const payload: WebSocketMessage = lastJsonMessage; const timestamp = payload.timestamp; switch (payload.type) { case WebSocketMessageType.ERROR: toast({ title: "Error", description: (payload.data as TelegramError).message, variant: "destructive", }); break; case WebSocketMessageType.FILE_DOWNLOAD: const { downloadedSize, totalCount } = payload.data as { totalSize: number; totalCount: number; downloadedSize: number; }; if (totalCount === 0) { setAccountDownloadSpeed({ speed: 0, lastDownloadedSize: 0, lastTimestamp: 0, }); return; } setAccountDownloadSpeed((prev) => { const state = { lastTimestamp: timestamp, lastDownloadedSize: downloadedSize, }; const timeDiff = (timestamp - prev.lastTimestamp) / 1000; if (prev.lastTimestamp === 0 || timeDiff <= 0 || downloadedSize <= prev.lastDownloadedSize) { return { ...state, speed: prev.speed, }; } const speed = (downloadedSize - prev.lastDownloadedSize) / timeDiff; return { speed, lastDownloadedSize: downloadedSize, lastTimestamp: timestamp, }; }); } } catch (error) { console.error("Failed to parse WebSocket message:", error); } } }, [lastJsonMessage, toast]); const sendWebSocketMessage = useCallback( (message: WebSocketMessage) => { if (isReady) { sendMessage(JSON.stringify(message)); } }, [isReady, sendMessage], ); return ( {children} ); }; export function useWebsocket() { const context = useContext(WebSocketContext); if (context === undefined) { throw new Error( "useTelegramWebSocket must be used within a WebSocketProvider", ); } return context; }