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");

// 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(nextGoals);
        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) {
            next[idx] = { ...next[idx], ...goal };
            return next;
          }
          next.push(goal as Goal);
          return 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) return;
        setMessages((prev: Message[]) => [...prev, message]);
        updateMessageGroups(message);
      }
    },
    [updateMessageGroups],
  );

  const wsOptions = useMemo(() => ({ onMessage: handleWebSocketMessage }), [handleWebSocketMessage]);
  const { connected } = useWebSocket(traceId, wsOptions);

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