import { useEffect, useRef, useState } from 'react';
export function useDownloadSocket(onEvent) {
const [connection, setConnection] = useState('connecting');
const wsRef = useRef(null);
const subscribedRef = useRef(new Set());
const onEventRef = useRef(onEvent);
onEventRef.current = onEvent;
useEffect(() => {
let cancelled = false;
let socket = null;
let reconnectDelay = 1000;
let reconnectTimer = null;
let heartbeat = null;
function send(message) {
if (socket && socket.readyState === WebSocket.OPEN) {
socket.send(JSON.stringify(message));
}
}
function connect() {
if (cancelled) return;
const proto = location.protocol === 'https:' ? 'wss' : 'ws';
socket = new WebSocket(`${proto}://${location.host}/ws`);
wsRef.current = socket;
socket.addEventListener('open', () => {
reconnectDelay = 1000;
setConnection('live');
subscribedRef.current.forEach((id) => send({ type: 'subscribe', session_id: id }));
heartbeat = setInterval(() => send({ type: 'ping' }), 25000);
});
socket.addEventListener('message', (event) => {
try {
const message = JSON.parse(event.data);
if (message.type !== 'pong') onEventRef.current(message);
} catch {
}
});
socket.addEventListener('close', () => {
if (cancelled) return;
setConnection('offline');
clearInterval(heartbeat);
reconnectTimer = setTimeout(connect, reconnectDelay);
reconnectDelay = Math.min(reconnectDelay * 1.6, 15000);
});
socket.addEventListener('error', () => socket.close());
}
connect();
return () => {
cancelled = true;
clearTimeout(reconnectTimer);
clearInterval(heartbeat);
socket?.close();
};
}, []);
function subscribe(sessionId) {
subscribedRef.current.add(sessionId);
if (wsRef.current?.readyState === WebSocket.OPEN) {
wsRef.current.send(JSON.stringify({ type: 'subscribe', session_id: sessionId }));
}
}
return { connection, subscribe };
}