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

import { useCallback, useEffect, useMemo, 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[]) => {
  const [goals, setGoals] = useState<Goal[]>(initialGoals);
  const [messages, setMessages] = useState<Message[]>([]);
  const [msgGroups, setMsgGroups] = useState<Record<string, Message[]>>({});

  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);
    setMessages([]);
    setMsgGroups({});
  }, [initialGoals, 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;

      if (event === "connected") {
        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, updateMessageGroups],
  );

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

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