Agent/frontend/react-template/src/components/MainContent/MainContent.tsx

import { useRef, useState, useEffect } from "react";
import type { FC, RefObject } from "react";
import { Select } from "@douyinfe/semi-ui";
import { FlowChart } from "../FlowChart/FlowChart";
import type { FlowChartRef } from "../FlowChart/FlowChart";
import { useFlowChartData } from "../FlowChart/hooks/useFlowChartData";
import { traceApi } from "../../api/traceApi";
import type { Goal } from "../../types/goal";
import type { Edge, Message } from "../../types/message";
import type { TraceListItem } from "../../types/trace";
import styles from "./MainContent.module.css";

interface MainContentProps {
  traceId: string | null;
  onNodeClick?: (node: Goal | Message, edge?: Edge) => void;
  onTraceChange?: (traceId: string, title?: string) => void;
  refreshTrigger?: number;
  messageRefreshTrigger?: number;
  flowChartRef?: RefObject<FlowChartRef>;
}

interface ConnectionStatusProps {
  isConnected: boolean;
}

const ConnectionStatus: FC<ConnectionStatusProps> = ({ isConnected }) => {
  return (
    <div
      className={`${styles.status} ${isConnected ? styles.connected : styles.disconnected}`}
      data-testid="connection-status"
    >
      <span className={styles.statusDot} />
      <span>{isConnected ? "WebSocket 已连接" : "WebSocket 未连接"}</span>
    </div>
  );
};

export const MainContent: FC<MainContentProps> = ({
  traceId,
  onNodeClick,
  onTraceChange,
  refreshTrigger,
  messageRefreshTrigger,
  flowChartRef: externalFlowChartRef,
}) => {
  const internalFlowChartRef = useRef<FlowChartRef>(null);
  const flowChartRef = externalFlowChartRef ?? internalFlowChartRef;
  const [isAllExpanded, setIsAllExpanded] = useState(true);
  const [traceList, setTraceList] = useState<TraceListItem[]>([]);
  const [cachedGoals, setCachedGoals] = useState<Goal[]>([]);
  const [cachedMsgGroups, setCachedMsgGroups] = useState<Record<string, Message[]>>({});
  const [cachedInvalidBranches, setCachedInvalidBranches] = useState<Message[][]>([]);
  const { goals, connected, msgGroups, reloading, invalidBranches, traceCompleted } = useFlowChartData(
    traceId,
    messageRefreshTrigger,
  );
  console.log("%c [ msgGroups ]-34", "font-size:13px; background:pink; color:#bf2c9f;", msgGroups);
  const displayGoals = goals.length > 0 ? goals : cachedGoals;
  const displayMsgGroups = Object.keys(msgGroups).length > 0 ? msgGroups : cachedMsgGroups;
  const displayInvalidBranches =
    invalidBranches && invalidBranches.length > 0 ? invalidBranches : cachedInvalidBranches;

  useEffect(() => {
    const fetchTraces = async () => {
      try {
        const data = await traceApi.fetchTraces({ limit: 100 });
        setTraceList(data.traces);
      } catch (error) {
        console.error("Failed to load traces:", error);
      }
    };
    fetchTraces();
  }, [refreshTrigger]);

  useEffect(() => {
    // 移除 reload 调用，因为 useFlowChartData 内部会监听 messageRefreshTrigger 并重新加载
  }, [messageRefreshTrigger]);

  useEffect(() => {
    if (goals.length > 0) {
      setCachedGoals(goals);
    }
  }, [goals]);

  useEffect(() => {
    if (Object.keys(msgGroups).length > 0) {
      setCachedMsgGroups(msgGroups);
    }
  }, [msgGroups]);

  useEffect(() => {
    if (invalidBranches && invalidBranches.length > 0) {
      setCachedInvalidBranches(invalidBranches);
    }
  }, [invalidBranches]);

  useEffect(() => {
    setCachedGoals([]);
    setCachedMsgGroups({});
    setCachedInvalidBranches([]);
  }, [traceId]);

  if (!traceId && !reloading) {
    return (
      <div className={styles.main}>
        <div className={styles.header}>
          <div className={styles.title}>暂无 Trace</div>
          <ConnectionStatus isConnected={connected} />
        </div>
        <div className={styles.content}>
          <div className={styles.empty}>暂无可展示的数据</div>
        </div>
      </div>
    );
  }

  return (
    <div className={styles.main}>
      <div className={styles.header}>
        <div style={{ display: "flex", alignItems: "center", gap: 16 }}>
          <ConnectionStatus isConnected={connected} />
          {traceCompleted && (
            <div className={styles.completedBadge}>
              <span style={{ fontSize: 16 }}>✓</span>
              <span>执行完成</span>
            </div>
          )}
        </div>
        <div className={styles.headerRight}>
          <Select
            value={traceId}
            onChange={(value: unknown) => {
              const trace = traceList.find((t) => t.trace_id === value);
              onTraceChange?.(value as string, trace?.task || trace?.trace_id);
            }}
            style={{ width: 400 }}
            placeholder="选择 Trace"
            optionList={traceList.map((t) => {
              const taskDesc = t.task && t.task.length > 50 ? `${t.task.slice(0, 50)}...` : t.task;
              return {
                label: taskDesc || t.trace_id,
                value: t.trace_id,
              };
            })}
          />
          {/* <div className={styles.status}>{connected ? "WebSocket 已连接" : "WebSocket 未连接"}</div> */}
          {/* <div className={styles.legend}>
            <div className={styles.legendItem}>
              <span
                className={styles.legendDot}
                style={{ background: "#00c853" }}
              />
              已完成
            </div>
            <div className={styles.legendItem}>
              <span
                className={styles.legendDot}
                style={{ background: "#f44336" }}
              />
              失败
            </div>
            <div className={styles.legendItem}>
              <span
                className={styles.legendDot}
                style={{ background: "#ff9800" }}
              />
              运行中
            </div>
            <div className={styles.legendItem}>
              <span
                className={styles.legendDot}
                style={{ background: "#4e79a7" }}
              />
              默认
            </div>
          </div> */}
          <div className={styles.buttons}>
            <button
              className={styles.btn}
              onClick={() => {
                if (isAllExpanded) {
                  flowChartRef.current?.collapseAll();
                } else {
                  flowChartRef.current?.expandAll();
                }
                setIsAllExpanded(!isAllExpanded);
              }}
            >
              {isAllExpanded ? "全折叠" : "全展开"}
            </button>
          </div>
        </div>
      </div>
      <div className={styles.content}>
        {reloading ? (
          <div className={styles.loading}>加载中...</div>
        ) : displayGoals.length === 0 ? (
          <div className={styles.empty}>暂无数据</div>
        ) : (
          <FlowChart
            ref={flowChartRef}
            goals={displayGoals}
            msgGroups={displayMsgGroups}
            invalidBranches={displayInvalidBranches}
            onNodeClick={onNodeClick}
            onSubTraceClick={(_parentGoal, entry) => {
              onTraceChange?.(entry.id, entry.mission || entry.id);
            }}
          />
        )}
      </div>
    </div>
  );
};