Agent/frontend/react-template/src/components/FlowChart/hooks/useFlowChartData.ts

import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { useWebSocket } from "../../../hooks/useWebSocket";
import type { Goal } from "../../../types/goal";
import type { Message } from "../../../types/message";

// WebSocket 数据解析与状态聚合（goals + msgGroups）
const isRecord = (value: unknown): value is Record<string, unknown> =>
  !!value && typeof value === "object" && !Array.isArray(value);

const isGoalLike = (value: unknown): value is Partial<Goal> & { id: string } =>
  isRecord(value) && typeof value.id === "string";

const isMessage = (value: unknown): value is Message =>
  isRecord(value) &&
  (typeof value.id === "string" || typeof (value as { message_id?: string }).message_id === "string");

const buildSubGoals = (flatGoals: Goal[]): Goal[] => {
  const nodeMap = new Map<string, Goal>();
  flatGoals.forEach((goal) => {
    nodeMap.set(goal.id, { ...goal, sub_goals: [] });
  });

  flatGoals.forEach((goal) => {
    const parentId = typeof goal.parent_id === "string" && goal.parent_id ? goal.parent_id : undefined;
    if (!parentId) return;
    const parent = nodeMap.get(parentId);
    const child = nodeMap.get(goal.id);
    if (!parent || !child) return;
    if (!Array.isArray(parent.sub_goals)) parent.sub_goals = [];
    parent.sub_goals.push(child);
  });

  return flatGoals.map((goal) => {
    const node = nodeMap.get(goal.id);
    if (!node) return goal;
    if (Array.isArray(node.sub_goals) && node.sub_goals.length === 0) {
      delete node.sub_goals;
    }
    return node;
  });
};

// FlowChart 专用数据 Hook：处理实时事件并聚合消息组
export const useFlowChartData = (traceId: string | null, initialGoals: Goal[], refreshTrigger?: number) => {
  const [goals, setGoals] = useState<Goal[]>(initialGoals);
  const [messages, setMessages] = useState<Message[]>([]);
  const [msgGroups, setMsgGroups] = useState<Record<string, Message[]>>({});
  const [sinceEventId, setSinceEventId] = useState(0);
  const currentEventIdRef = useRef(0);
  const restReloadingRef = useRef(false);
  const [reloading, setReloading] = useState(false);

  const messageSortKey = useCallback((message: Message): number => {
    const mid =
      typeof message.message_id === "string"
        ? message.message_id
        : typeof message.id === "string"
          ? message.id
          : undefined;
    if (!mid) return 0;
    if (!mid.includes("-")) return 0;
    const suffix = mid.slice(mid.lastIndexOf("-") + 1);
    const num = Number.parseInt(suffix, 10);
    return Number.isFinite(num) ? num : 0;
  }, []);

  const updateMessageGroups = useCallback(
    (message: Message) => {
      const groupKey = typeof message.goal_id === "string" && message.goal_id ? message.goal_id : "START";
      if (groupKey === "START") {
        setGoals((prev) => {
          if (prev.some((g) => g.id === "START")) return prev;
          const startGoal: Goal = {
            id: "START",
            description: "START",
            status: "completed",
            created_at: "",
          };
          return [startGoal, ...prev];
        });
      }
      setMsgGroups((prev) => {
        const existing = prev[groupKey] ? [...prev[groupKey]] : [];
        existing.push(message);
        existing.sort((a, b) => messageSortKey(a) - messageSortKey(b));
        return { ...prev, [groupKey]: existing };
      });
    },
    [messageSortKey],
  );

  useEffect(() => {
    setGoals(initialGoals);
  }, [initialGoals]);

  useEffect(() => {
    setMessages([]);
    setMsgGroups({});
    setSinceEventId(0);
    currentEventIdRef.current = 0;
    restReloadingRef.current = false;
  }, [traceId]);

  const reloadViaRest = useCallback(async () => {
    if (!traceId) return;
    if (restReloadingRef.current) return;
    restReloadingRef.current = true;
    setReloading(true);
    let nextSinceEventId: number | null = null;
    try {
      const [traceRes, messagesRes] = await Promise.all([
        fetch(`http://localhost:8000/api/traces/${traceId}`),
        fetch(`http://localhost:8000/api/traces/${traceId}/messages?mode=all`),
      ]);

      if (traceRes.ok) {
        const json = (await traceRes.json()) as unknown;
        const root = isRecord(json) ? json : {};
        const trace = isRecord(root.trace) ? root.trace : undefined;
        const goalTree = isRecord(root.goal_tree) ? root.goal_tree : undefined;
        const goalList = goalTree && Array.isArray(goalTree.goals) ? (goalTree.goals as Goal[]) : [];

        const lastEventId = trace && typeof trace.last_event_id === "number" ? trace.last_event_id : undefined;
        if (typeof lastEventId === "number") {
          currentEventIdRef.current = Math.max(currentEventIdRef.current, lastEventId);
          setSinceEventId(lastEventId);
          nextSinceEventId = lastEventId;
        }

        if (goalList.length > 0) {
          setGoals((prev) => {
            const mergedFlat = goalList.map((ng) => {
              const existing = prev.find((p) => p.id === ng.id);
              if (!existing) return ng;
              const merged: Goal = { ...existing, ...ng };
              if (existing.sub_trace_ids && !merged.sub_trace_ids) {
                merged.sub_trace_ids = existing.sub_trace_ids;
              }
              if (existing.agent_call_mode && !merged.agent_call_mode) {
                merged.agent_call_mode = existing.agent_call_mode;
              }
              return merged;
            });
            return buildSubGoals(mergedFlat);
          });
        }
      }

      if (messagesRes.ok) {
        const json = (await messagesRes.json()) as unknown;
        const root = isRecord(json) ? json : {};
        const list = Array.isArray(root.messages) ? (root.messages as Message[]) : [];
        const filtered = list.filter((message) => (message as { status?: string }).status !== "abandoned");
        const nextMessages = [...filtered].sort((a, b) => messageSortKey(a) - messageSortKey(b));
        setMessages(nextMessages);
        const grouped: Record<string, Message[]> = {};
        nextMessages.forEach((message) => {
          const groupKey = typeof message.goal_id === "string" && message.goal_id ? message.goal_id : "START";
          if (!grouped[groupKey]) grouped[groupKey] = [];
          grouped[groupKey].push(message);
        });
        Object.keys(grouped).forEach((key) => {
          grouped[key].sort((a, b) => messageSortKey(a) - messageSortKey(b));
        });
        setMsgGroups(grouped);

        if (grouped.START && grouped.START.length > 0) {
          setGoals((prev) => {
            if (prev.some((g) => g.id === "START")) return prev;
            const startGoal: Goal = {
              id: "START",
              description: "START",
              status: "completed",
              created_at: "",
            };
            return [startGoal, ...prev];
          });
        }
      }
    } finally {
      restReloadingRef.current = false;
      setReloading(false);
    }
    return nextSinceEventId;
  }, [messageSortKey, traceId]);

  useEffect(() => {
    if (!traceId) return;
    if (!refreshTrigger) return;
    void reloadViaRest();
  }, [refreshTrigger, reloadViaRest, traceId]);

  const handleWebSocketMessage = useCallback(
    (payload: unknown) => {
      const raw = isRecord(payload) ? payload : {};
      const event = (typeof raw.event === "string" && raw.event) || (typeof raw.type === "string" && raw.type) || "";
      const data = isRecord(raw.data) ? raw.data : raw;
      console.log("%c [ data ]-182", "font-size:13px; background:pink; color:#bf2c9f;", data);

      const eventId = typeof raw.event_id === "number" ? raw.event_id : undefined;
      if (typeof eventId === "number") {
        currentEventIdRef.current = Math.max(currentEventIdRef.current, eventId);
      }

      if (event === "error") {
        const message =
          (typeof data.message === "string" ? data.message : undefined) ||
          (typeof raw.message === "string" ? raw.message : undefined) ||
          "";
        if (message.includes("Too many missed events")) {
          void reloadViaRest().then((nextSince) => {
            if (typeof nextSince === "number") return;
            const fallbackSince = currentEventIdRef.current;
            if (fallbackSince > 0) setSinceEventId(fallbackSince);
          });
        }
        return;
      }

      if (event === "connected") {
        const currentEventId =
          (typeof data.current_event_id === "number" ? data.current_event_id : undefined) ||
          (typeof raw.current_event_id === "number" ? raw.current_event_id : undefined);
        if (typeof currentEventId === "number") {
          currentEventIdRef.current = Math.max(currentEventIdRef.current, currentEventId);
        }

        const trace = isRecord(data.trace) ? data.trace : undefined;
        const rawTrace = isRecord(raw.trace) ? raw.trace : undefined;
        const goalTree =
          (isRecord(data.goal_tree) ? data.goal_tree : undefined) ||
          (trace && isRecord(trace.goal_tree) ? trace.goal_tree : undefined) ||
          (isRecord(raw.goal_tree) ? raw.goal_tree : undefined) ||
          (rawTrace && isRecord(rawTrace.goal_tree) ? rawTrace.goal_tree : undefined) ||
          {};
        const goalList = isRecord(goalTree) ? goalTree.goals : undefined;
        const nextGoals = Array.isArray(goalList) ? (goalList as Goal[]) : [];
        setGoals((prev) => {
          const mergedFlat = nextGoals.map((ng) => {
            const existing = prev.find((p) => p.id === ng.id);
            if (!existing) return ng;
            const merged: Goal = { ...existing, ...ng };
            if (existing.sub_trace_ids && !merged.sub_trace_ids) {
              merged.sub_trace_ids = existing.sub_trace_ids;
            }
            if (existing.agent_call_mode && !merged.agent_call_mode) {
              merged.agent_call_mode = existing.agent_call_mode;
            }
            return merged;
          });
          return buildSubGoals(mergedFlat);
        });
        return;
      }

      if (event === "goal_added") {
        const goal = isGoalLike(data.goal) ? data.goal : isGoalLike(raw.goal) ? raw.goal : null;
        if (!goal) return;
        setGoals((prev: Goal[]) => {
          const next = [...prev];
          const idx = next.findIndex((g) => g.id === goal.id);
          if (idx >= 0) {
            const existing = next[idx];
            const merged = { ...existing, ...goal };
            // 保留 sub_trace_ids，如果 WebSocket 数据中缺失但本地已有
            if (existing.sub_trace_ids && !merged.sub_trace_ids) {
              merged.sub_trace_ids = existing.sub_trace_ids;
            }
            if (existing.agent_call_mode && !merged.agent_call_mode) {
              merged.agent_call_mode = existing.agent_call_mode;
            }
            next[idx] = merged;
            return buildSubGoals(next);
          }
          next.push(goal as Goal);
          return buildSubGoals(next);
        });
        return;
      }

      if (event === "goal_updated") {
        const goalId =
          (typeof data.goal_id === "string" ? data.goal_id : undefined) ||
          (isRecord(data.goal) && typeof data.goal.id === "string" ? data.goal.id : undefined) ||
          (typeof raw.goal_id === "string" ? raw.goal_id : undefined);
        const updates = isRecord(data.updates) ? data.updates : isRecord(raw.updates) ? raw.updates : {};
        if (!goalId) return;
        setGoals((prev: Goal[]) =>
          prev.map((g: Goal) => {
            if (g.id !== goalId) return g;
            const next: Goal = { ...g };
            if ("status" in updates) {
              const status = updates.status;
              if (typeof status === "string") {
                next.status = status as Goal["status"];
              }
            }
            if ("summary" in updates) {
              const summary = updates.summary;
              if (typeof summary === "string") {
                next.summary = summary;
              }
            }
            return next;
          }),
        );
        return;
      }

      if (event === "message_added") {
        const message = isMessage(data.message) ? data.message : isMessage(raw.message) ? raw.message : null;
        if (message) {
          setMessages((prev) => {
            const next = [...prev, message];
            next.sort((a, b) => messageSortKey(a) - messageSortKey(b));
            return next;
          });
          updateMessageGroups(message);
        }
      }
    },
    [messageSortKey, reloadViaRest, updateMessageGroups],
  );

  // 主 Trace 连接
  const wsOptions = useMemo(
    () => ({ onMessage: handleWebSocketMessage, sinceEventId }),
    [handleWebSocketMessage, sinceEventId],
  );
  const { connected } = useWebSocket(traceId, wsOptions);

  return { goals, messages, msgGroups, connected, reloading };
};