path: root/makima/frontend/src/hooks/useFileSubscription.ts

import { useState, useCallback, useRef, useEffect } from "react";
import { FILE_SUBSCRIBE_ENDPOINT } from "../lib/api";

export interface FileUpdateEvent {
  fileId: string;
  version: number;
  updatedFields: string[];
  updatedBy: "user" | "llm" | "system";
}

interface UseFileSubscriptionOptions {
  fileId: string | null;
  onUpdate?: (event: FileUpdateEvent) => void;
  onError?: (error: string) => void;
}

export function useFileSubscription(options: UseFileSubscriptionOptions) {
  const { fileId, onUpdate, onError } = options;
  const [connected, setConnected] = useState(false);
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimeoutRef = useRef<number | null>(null);
  const subscribedFileRef = useRef<string | null>(null);

  // Store callbacks in refs to avoid re-connecting when callbacks change
  const callbacksRef = useRef({ onUpdate, onError });
  useEffect(() => {
    callbacksRef.current = { onUpdate, onError };
  }, [onUpdate, onError]);

  const connect = useCallback(() => {
    if (wsRef.current?.readyState === WebSocket.OPEN) return;

    try {
      const ws = new WebSocket(FILE_SUBSCRIBE_ENDPOINT);
      wsRef.current = ws;

      ws.onopen = () => {
        setConnected(true);
        // Re-subscribe if we had a subscription
        if (subscribedFileRef.current) {
          ws.send(
            JSON.stringify({
              type: "subscribe",
              fileId: subscribedFileRef.current,
            })
          );
        }
      };

      ws.onmessage = (event) => {
        try {
          const message = JSON.parse(event.data);

          if (message.type === "fileUpdated") {
            callbacksRef.current.onUpdate?.({
              fileId: message.fileId,
              version: message.version,
              updatedFields: message.updatedFields,
              updatedBy: message.updatedBy,
            });
          } else if (message.type === "error") {
            callbacksRef.current.onError?.(message.message);
          }
        } catch (e) {
          console.error("Failed to parse file subscription message:", e);
        }
      };

      ws.onerror = () => {
        callbacksRef.current.onError?.("WebSocket connection error");
      };

      ws.onclose = () => {
        setConnected(false);
        wsRef.current = null;

        // Attempt reconnection after 3 seconds if we still have a subscription
        if (subscribedFileRef.current) {
          reconnectTimeoutRef.current = window.setTimeout(() => {
            connect();
          }, 3000);
        }
      };
    } catch (e) {
      callbacksRef.current.onError?.(
        e instanceof Error ? e.message : "Failed to connect"
      );
    }
  }, []);

  const subscribe = useCallback(
    (id: string) => {
      subscribedFileRef.current = id;

      if (wsRef.current?.readyState === WebSocket.OPEN) {
        wsRef.current.send(
          JSON.stringify({
            type: "subscribe",
            fileId: id,
          })
        );
      } else {
        connect();
      }
    },
    [connect]
  );

  const unsubscribe = useCallback(() => {
    if (
      subscribedFileRef.current &&
      wsRef.current?.readyState === WebSocket.OPEN
    ) {
      wsRef.current.send(
        JSON.stringify({
          type: "unsubscribe",
          fileId: subscribedFileRef.current,
        })
      );
    }
    subscribedFileRef.current = null;
  }, []);

  // Auto-subscribe when fileId changes
  useEffect(() => {
    if (fileId) {
      subscribe(fileId);
    } else {
      unsubscribe();
    }

    return () => {
      unsubscribe();
    };
  }, [fileId, subscribe, unsubscribe]);

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
      }
      if (wsRef.current) {
        wsRef.current.close();
      }
    };
  }, []);

  return {
    connected,
    subscribe,
    unsubscribe,
  };
}