import type { WebsocketStatus } from "@/components/ConnectionStatus";
import type { StatusMessage } from "@/services/types";
import { getLastMessage, setLastMessage } from "@/services/store";
import { initializeConnection } from "@/services/websocket";
import { useEffect, useState } from "react";

export function useWebsocket(agent: string) {
    const [status, setStatus] = useState<WebsocketStatus>();
    const [message, setMessage] = useState<StatusMessage | null>(
        getLastMessage(),
    );

    useEffect(() => {
        initializeConnection(agent, {
            onMessage: (data) => {
                const parsed = JSON.parse(data);

                setStatus("connected");
                setMessage(parsed);
                setLastMessage(parsed);
            },
            onOpen: () => setStatus("connecting"),
            onClose: () => setStatus("disconnected"),
            onError: () => setStatus("error"),
        });
    }, [agent]);

    return { status, message };
}