frontend/src/hooks/useLangGraph.ts

"use client";

import { useState, useCallback, useRef } from "react";
import { Client } from "@langchain/langgraph-sdk";

interface UseLangGraphOptions {
  apiUrl?: string;
  assistantId?: string;
}

export interface LangGraphUIEvent {
  id: string;
  name: string;
  props: Record<string, unknown>;
  status: "loading" | "streaming" | "complete" | "error";
}

export function useLangGraph<T extends LangGraphUIEvent = LangGraphUIEvent>({
  apiUrl = "http://localhost:2024",
  assistantId = "report_agent",
}: UseLangGraphOptions = {}) {
  const [uiEvents, setUiEvents] = useState<Record<string, T>>({});
  const [isProcessing, setIsProcessing] = useState(false);
  const clientRef = useRef(new Client({ apiUrl }));

  const submit = useCallback(async (input: Record<string, unknown> = {}) => {
    setIsProcessing(true);
    try {
      // 1. Create a thread
      const thread = await clientRef.current.threads.create();
      
      // 2. Stream steps
      const stream = clientRef.current.runs.stream(
        thread.thread_id,
        assistantId,
        { input }
      );

      for await (const chunk of stream) {
        // Handle UI events pushed from the server
        // Native LangGraph sends UI messages in 'values' or as custom events if configured.
        // In our report_builder.py, we use push_ui_message.
        
        const data = chunk.data as Record<string, any>;
        if (chunk.event === "values" && data?.ui) {
          const uiMessages = data.ui as any[];
          // Reduce ALL ui messages by id so merged updates (merge=True)
          // replace the previous version of the same component, and
          // multiple components can coexist.
          setUiEvents((prev) => {
            const next = { ...prev };
            for (const msg of uiMessages) {
              next[msg.id] = {
                id: msg.id,
                name: msg.name,
                props: msg.props,
                status: msg.props?.status || "complete",
              } as T;
            }
            return next;
          });
        }
      }
    } catch (error) {
      console.error("LangGraph Sidecar Error:", error);
    } finally {
      setIsProcessing(false);
    }
  }, [assistantId]);

  const clear = useCallback(() => {
    setUiEvents({});
  }, []);

  return {
    submit,
    clear,
    uiEvents: Object.values(uiEvents),
    isProcessing
  };
}