"use client";

import { useCallback, useEffect, useRef, useState } from "react";

import {
  abortAgentChat,
  forkAgentChat,
  replyAgentPermission,
  resumeAgentChatStream,
  streamAgentChat,
} from "@/lib/chatStream";
import type { AgentModel, PermissionReply, StreamEvent } from "@/lib/chatStream";
import type {
  AgentArtifact,
  AgentPermissionRequest,
  BranchGroup,
  BranchTransition,
  ChatProgress,
  ChatSessionSummary,
  LoadedChatState,
  Message,
} from "../GlobalChatbox.types";
import {
  cloneBranchGroups,
  cloneMessages,
  createId,
} from "../GlobalChatbox.utils";
import {
  createEmptyChatState,
  deleteChatSession,
  listChatSessions,
  loadChatSessionById,
  saveActiveChatState,
  updateChatSessionTitle,
} from "../chatStorage";

type UseAgentChatSessionOptions = {
  projectId?: string | null;
  onToolCall: (
    event: StreamEvent & { type: "tool_call" },
    options: {
      assistantMessageId: string;
      appendArtifact: (messageId: string, artifact: AgentArtifact) => void;
    },
  ) => void;
  onBeforeSend?: () => void;
  getModel?: () => AgentModel;
};

type PromptRunOptions = {
  prompt: string;
  sessionIdOverride?: string;
  preparedMessages?: Message[];
  userMessage?: Message;
  assistantMessage?: Message;
};

const createPersistedStateKey = (state: LoadedChatState) =>
  JSON.stringify({
    title: state.title ?? null,
    isTitleManuallyEdited: state.isTitleManuallyEdited ?? false,
    sessionId: state.sessionId ?? null,
    messages: state.messages,
    branchGroups: state.branchGroups,
  });

const upsertProgress = (
  progress: ChatProgress[] | undefined,
  event: StreamEvent & { type: "progress" },
) => {
  const next = [...(progress ?? [])];
  const index = next.findIndex((item) => item.id === event.id);
  const existing = index >= 0 ? next[index] : undefined;
  const now = Date.now();
  const startedAt = event.startedAt ?? existing?.startedAt;
  const isRunning = event.status === "running";
  const endedAt = isRunning ? undefined : event.endedAt ?? existing?.endedAt ?? now;
  const elapsedMs = isRunning
    ? event.elapsedMs ??
      existing?.elapsedMs ??
      (startedAt !== undefined ? Math.max(0, now - startedAt) : undefined)
    : undefined;
  const elapsedSnapshotAt = isRunning
    ? event.elapsedMs !== undefined
      ? now
      : existing?.elapsedSnapshotAt ?? now
    : undefined;
  const durationMs = !isRunning
    ? event.durationMs ??
      existing?.durationMs ??
      (startedAt !== undefined && endedAt !== undefined
        ? Math.max(0, endedAt - startedAt)
        : undefined)
    : undefined;
  const nextItem: ChatProgress = {
    id: event.id,
    phase: event.phase,
    status: event.status,
    title: event.title,
    detail: event.detail,
    startedAt,
    endedAt,
    elapsedMs,
    elapsedSnapshotAt,
    durationMs,
  };
  if (index >= 0) {
    next[index] = nextItem;
  } else {
    next.push(nextItem);
  }
  return next;
};

const completeRunningProgress = (progress: ChatProgress[] | undefined) =>
  progress?.map((item) => {
    if (item.status !== "running") {
      return item;
    }
    const endedAt = Date.now();
    return {
      ...item,
      status: "completed" as const,
      endedAt,
      elapsedMs: undefined,
      elapsedSnapshotAt: undefined,
      durationMs:
        item.durationMs ??
        (item.startedAt !== undefined
          ? Math.max(0, endedAt - item.startedAt)
          : item.elapsedMs),
    };
  });

const upsertPermission = (
  permissions: AgentPermissionRequest[] | undefined,
  event: StreamEvent & { type: "permission_request" },
) => {
  const next = [...(permissions ?? [])];
  const index = next.findIndex((item) => item.requestId === event.requestId);
  const nextItem: AgentPermissionRequest = {
    requestId: event.requestId,
    sessionId: event.sessionId,
    permission: event.permission,
    patterns: event.patterns,
    metadata: event.metadata,
    always: event.always,
    tool: event.tool,
    createdAt: event.createdAt,
    status: "pending",
  };
  if (index >= 0) {
    next[index] = {
      ...next[index],
      ...nextItem,
      status: next[index].status === "submitting" ? "submitting" : nextItem.status,
    };
  } else {
    next.push(nextItem);
  }
  return next;
};

const toPermissionStatus = (reply: PermissionReply): AgentPermissionRequest["status"] => {
  if (reply === "always") return "approved_always";
  if (reply === "once") return "approved_once";
  return "rejected";
};

const finalizeAssistantMessageAfterAbort = (message: Message): Message => {
  const completedProgress = completeRunningProgress(message.progress);
  const hasVisibleOutput =
    message.content.trim().length > 0 ||
    Boolean(message.artifacts?.length) ||
    Boolean(completedProgress?.length);

  if (!hasVisibleOutput) {
    return message;
  }

  return {
    ...message,
    content: message.content || "⚠️ **请求已中断**",
    isError: true,
    progress: completedProgress,
  };
};

const createUserMessage = (content: string, branchRootId?: string): Message => {
  const id = createId();
  return {
    id,
    role: "user",
    content,
    branchRootId: branchRootId ?? id,
  };
};

const createAssistantMessage = (): Message => ({
  id: createId(),
  role: "assistant",
  content: "",
});

const messagesEqual = (left: Message[], right: Message[]) =>
  JSON.stringify(left) === JSON.stringify(right);

export const useAgentChatSession = ({
  projectId,
  onToolCall,
  onBeforeSend,
  getModel,
}: UseAgentChatSessionOptions) => {
  const hydrationCompletedRef = useRef(false);
  const hydrationNonceRef = useRef(0);

  const [messages, setMessages] = useState<Message[]>([]);
  const [sessionTitle, setSessionTitle] = useState<string | undefined>(undefined);
  const [isSessionTitleManuallyEdited, setIsSessionTitleManuallyEdited] = useState(false);
  const [sessionId, setSessionId] = useState<string | undefined>(undefined);
  const [branchGroups, setBranchGroups] = useState<BranchGroup[]>([]);
  const [chatSessions, setChatSessions] = useState<ChatSessionSummary[]>([]);
  const [branchTransition, setBranchTransition] = useState<BranchTransition | null>(null);
  const [isStreaming, setIsStreaming] = useState(false);
  const [isHydrating, setIsHydrating] = useState(true);
  const abortRef = useRef<AbortController | null>(null);
  const sessionIdRef = useRef<string | undefined>(undefined);
  const messagesRef = useRef<Message[]>([]);
  const resumeStreamingSessionRef = useRef<((sessionId: string) => void) | null>(null);
  const isSessionTitleManuallyEditedRef = useRef(false);
  const cancelPromiseRef = useRef<Promise<void> | null>(null);
  const titleUpdateNonceRef = useRef(0);
  const lastPersistedStateKeyRef = useRef(
    createPersistedStateKey({
      sessionId: undefined,
      title: undefined,
      isTitleManuallyEdited: false,
      messages: [],
      branchGroups: [],
    }),
  );

  useEffect(() => {
    sessionIdRef.current = sessionId;
  }, [sessionId]);

  useEffect(() => {
    messagesRef.current = messages;
  }, [messages]);

  useEffect(() => {
    isSessionTitleManuallyEditedRef.current = isSessionTitleManuallyEdited;
  }, [isSessionTitleManuallyEdited]);

  useEffect(() => {
    let cancelled = false;

    const hydrate = async () => {
      setIsHydrating(true);
      hydrationCompletedRef.current = false;

      if (!projectId) {
        sessionIdRef.current = undefined;
        lastPersistedStateKeyRef.current = createPersistedStateKey({
          title: undefined,
          isTitleManuallyEdited: false,
          messages: [],
          sessionId: undefined,
          branchGroups: [],
        });
        hydrationCompletedRef.current = true;
        hydrationNonceRef.current += 1;
        titleUpdateNonceRef.current += 1;
        setBranchTransition(null);
        setMessages([]);
        setSessionTitle(undefined);
        setIsSessionTitleManuallyEdited(false);
        setSessionId(undefined);
        setBranchGroups([]);
        setChatSessions([]);
        setIsHydrating(false);
        return;
      }

      try {
        const sessions = await listChatSessions();
        const streamingSession = sessions.find((session) => session.isStreaming);
        const loadedState = streamingSession
          ? await loadChatSessionById(streamingSession.id)
          : createEmptyChatState();
        if (cancelled) return;

        sessionIdRef.current = loadedState.sessionId;
        lastPersistedStateKeyRef.current = createPersistedStateKey(loadedState);
        hydrationCompletedRef.current = true;
        hydrationNonceRef.current += 1;
        titleUpdateNonceRef.current += 1;

        setMessages(loadedState.messages);
        setSessionTitle(loadedState.title);
        setIsSessionTitleManuallyEdited(loadedState.isTitleManuallyEdited ?? false);
        setSessionId(loadedState.sessionId);
        setBranchGroups(loadedState.branchGroups);
        setChatSessions(sessions);
        if (
          loadedState.sessionId &&
          (loadedState.isStreaming || streamingSession?.isStreaming)
        ) {
          resumeStreamingSessionRef.current?.(loadedState.sessionId);
        }
      } catch (error) {
        console.error("[GlobalChatbox] Failed to hydrate chat state:", error);
      } finally {
        if (!cancelled) {
          setIsHydrating(false);
        }
      }
    };

    void hydrate();

    return () => {
      cancelled = true;
    };
  }, [projectId]);

  useEffect(() => {
    if (!projectId || isHydrating || !hydrationCompletedRef.current) return;

    const currentHydrationNonce = hydrationNonceRef.current;
    const persistTimer = window.setTimeout(() => {
      if (isStreaming) {
        return;
      }

      const state: LoadedChatState = {
        title: sessionTitle,
        isTitleManuallyEdited: isSessionTitleManuallyEdited,
        messages,
        sessionId,
        branchGroups,
      };

      const currentStateKey = createPersistedStateKey(state);
      if (currentStateKey === lastPersistedStateKeyRef.current) {
        return;
      }

      void saveActiveChatState(state)
        .then((sessionId) => {
          if (hydrationNonceRef.current !== currentHydrationNonce) return;
          sessionIdRef.current = sessionId;
          lastPersistedStateKeyRef.current = createPersistedStateKey({
            ...state,
            sessionId,
          });
          return listChatSessions();
        })
        .then((sessions) => {
          if (!sessions || hydrationNonceRef.current !== currentHydrationNonce) return;
          setChatSessions(sessions);
        })
        .catch((error) => {
          console.error("[GlobalChatbox] Failed to persist chat state:", error);
        });
    }, 150);

    return () => {
      window.clearTimeout(persistTimer);
    };
  }, [branchGroups, isHydrating, isSessionTitleManuallyEdited, isStreaming, messages, projectId, sessionId, sessionTitle]);

  useEffect(() => {
    setBranchGroups((prev) => {
      let changed = false;
      const next = prev.map((group) => {
        const rootMessage = messages[group.parentCount];
        if (
          !rootMessage ||
          rootMessage.role !== "user" ||
          (rootMessage.branchRootId ?? rootMessage.id) !== group.rootMessageId
        ) {
          return group;
        }

        const activeBranch = group.branches[group.activeIndex];
        if (!activeBranch) {
          return group;
        }

        const nextSuffix = cloneMessages(messages.slice(group.parentCount));
        if (
          activeBranch.sessionId === sessionId &&
          messagesEqual(activeBranch.messages, nextSuffix)
        ) {
          return group;
        }

        changed = true;
        const branches = group.branches.map((branch, index) =>
          index === group.activeIndex
            ? { ...branch, sessionId, messages: nextSuffix }
            : branch,
        );
        return { ...group, branches };
      });

      return changed ? next : prev;
    });
  }, [messages, sessionId]);

  const appendArtifact = useCallback((messageId: string, artifact: AgentArtifact) => {
    setMessages((prev) =>
      prev.map((message) =>
        message.id === messageId
          ? {
              ...message,
              artifacts: [...(message.artifacts ?? []), artifact],
            }
          : message,
      ),
    );
  }, []);

  const getLastAssistantMessageId = useCallback((fallback?: string) => {
    const assistant = [...messagesRef.current]
      .reverse()
      .find((message) => message.role === "assistant");
    return assistant?.id ?? fallback;
  }, []);

  const applyStreamEvent = useCallback(
    (
      event: StreamEvent,
      options?: {
        assistantMessageId?: string;
      },
    ) => {
      if ("sessionId" in event && event.sessionId && event.sessionId !== sessionIdRef.current) {
        sessionIdRef.current = event.sessionId;
        setSessionId(event.sessionId);
      }

      if (event.type === "state") {
        const nextMessages = cloneMessages(event.messages as Message[]);
        messagesRef.current = nextMessages;
        setMessages(nextMessages);
        setIsStreaming(event.isStreaming);
        return;
      }

      const assistantMessageId = getLastAssistantMessageId(options?.assistantMessageId);
      if (!assistantMessageId) {
        return;
      }

      if (event.type === "token") {
        setMessages((prev) =>
          prev.map((message) =>
            message.id === assistantMessageId
              ? {
                  ...message,
                  content: message.content + event.content,
                  isError: false,
                }
              : message,
          ),
        );
      } else if (event.type === "progress") {
        setMessages((prev) =>
          prev.map((message) =>
            message.id === assistantMessageId
              ? { ...message, progress: upsertProgress(message.progress, event) }
              : message,
          ),
        );
      } else if (event.type === "tool_call") {
        onToolCall(event, {
          assistantMessageId,
          appendArtifact,
        });
      } else if (event.type === "permission_request") {
        setMessages((prev) =>
          prev.map((message) =>
            message.id === assistantMessageId
              ? {
                  ...message,
                  permissions: upsertPermission(message.permissions, event),
                }
              : message,
          ),
        );
      } else if (event.type === "permission_response") {
        setMessages((prev) =>
          prev.map((message) => {
            if (message.id !== assistantMessageId || !message.permissions?.length) {
              return message;
            }
            return {
              ...message,
              permissions: message.permissions.map((permission) =>
                permission.requestId === event.requestId
                  ? {
                      ...permission,
                      status: toPermissionStatus(event.reply),
                      repliedAt: Date.now(),
                      error: undefined,
                    }
                  : permission,
              ),
            };
          }),
        );
      } else if (event.type === "session_title") {
        const nextTitle = event.title.trim();
        if (nextTitle && !isSessionTitleManuallyEditedRef.current) {
          setSessionTitle(nextTitle);
          const currentSessionId = sessionIdRef.current;
          if (currentSessionId) {
            const currentNonce = ++titleUpdateNonceRef.current;
            void updateChatSessionTitle(currentSessionId, nextTitle, {
              isTitleManuallyEdited: false,
            })
              .then(() => listChatSessions())
              .then((sessions) => {
                if (titleUpdateNonceRef.current !== currentNonce) return;
                setChatSessions(sessions);
              })
              .catch((error) => {
                console.error("[GlobalChatbox] Failed to persist session title:", error);
              });
          }
        }
      } else if (event.type === "done") {
        setMessages((prev) =>
          prev.map((message) => {
            if (message.id !== assistantMessageId) return message;
            const completedProgress = completeRunningProgress(message.progress);
            if (
              message.content.trim().length === 0 &&
              !(message.artifacts?.length)
            ) {
              return {
                ...message,
                content:
                  "Agent 已完成处理，但没有生成文本回答。请查看过程记录，或换个更具体的问题重试。",
                progress: completedProgress,
              };
            }
            return { ...message, progress: completedProgress };
          }),
        );
        setIsStreaming(false);
      } else if (event.type === "error") {
        setMessages((prev) =>
          prev.map((message) =>
            message.id === assistantMessageId
              ? {
                  ...message,
                  content: message.content || `⚠️ **错误：** ${event.message}`,
                  isError: true,
                  progress: completeRunningProgress(message.progress),
                }
              : message,
          ),
        );
        setIsStreaming(false);
      }
    },
    [appendArtifact, getLastAssistantMessageId, onToolCall],
  );

  const resumeStreamingSession = useCallback(
    (nextSessionId: string) => {
      const controller = new AbortController();
      abortRef.current?.abort();
      abortRef.current = controller;
      setIsStreaming(true);

      void resumeAgentChatStream({
        sessionId: nextSessionId,
        signal: controller.signal,
        onEvent: (event) => applyStreamEvent(event),
      })
        .catch((error) => {
          if (!controller.signal.aborted) {
            console.error("[GlobalChatbox] Failed to resume chat stream:", error);
            setIsStreaming(false);
          }
        })
        .finally(() => {
          if (abortRef.current === controller) {
            abortRef.current = null;
          }
        });
    },
    [applyStreamEvent],
  );
  resumeStreamingSessionRef.current = resumeStreamingSession;

  const runPrompt = useCallback(
    async ({
      prompt: rawPrompt,
      sessionIdOverride,
      preparedMessages,
      userMessage,
      assistantMessage,
    }: PromptRunOptions) => {
      const prompt = rawPrompt.trim();
      if (!prompt || isStreaming || isHydrating) return;

      await cancelPromiseRef.current?.catch(() => undefined);
      onBeforeSend?.();
      setBranchTransition(null);

      const nextUserMessage = userMessage ?? createUserMessage(prompt);
      const nextAssistantMessage = assistantMessage ?? createAssistantMessage();
      const nextMessages =
        preparedMessages ??
        [...messages, nextUserMessage, nextAssistantMessage];

      const clonedNextMessages = cloneMessages(nextMessages);
      setIsStreaming(true);
      messagesRef.current = clonedNextMessages;
      setMessages(clonedNextMessages);
      if (sessionIdOverride !== undefined) {
        sessionIdRef.current = sessionIdOverride;
        setSessionId(sessionIdOverride);
      }

      const controller = new AbortController();
      abortRef.current = controller;

      try {
        await streamAgentChat({
          message: prompt,
          sessionId: sessionIdOverride ?? sessionIdRef.current,
          model: getModel?.(),
          signal: controller.signal,
          onEvent: (event) =>
            applyStreamEvent(event, {
              assistantMessageId: nextAssistantMessage.id,
            }),
        });
      } catch (error) {
        if (controller.signal.aborted) {
          setMessages((prev) =>
            prev
              .map((message) =>
                message.id === nextAssistantMessage.id
                  ? {
                      ...message,
                      content: message.content || "⚠️ **请求已中断**",
                      isError: true,
                    }
                  : message,
              )
              .filter(
                (message) =>
                  !(
                    message.id === nextAssistantMessage.id &&
                    message.role === "assistant" &&
                    message.content.trim().length === 0 &&
                    !(message.artifacts?.length) &&
                    !(message.progress?.length)
                  ),
              ),
          );
          return;
        }
        setMessages((prev) =>
          prev.map((message) =>
            message.id === nextAssistantMessage.id
              ? {
                  ...message,
                  content: `⚠️ **错误：** ${String(error)}`,
                  isError: true,
                  progress: completeRunningProgress(message.progress),
                }
              : message,
          ),
        );
        setIsStreaming(false);
      } finally {
        abortRef.current = null;
        setIsStreaming(false);
      }
    },
    [applyStreamEvent, getModel, isHydrating, isStreaming, messages, onBeforeSend],
  );

  const abort = useCallback(() => {
    const controller = abortRef.current;
    controller?.abort();
    setIsStreaming(false);
    const assistantMessageId = getLastAssistantMessageId();

    if (assistantMessageId) {
      setMessages((prev) =>
        prev.map((message) =>
          message.id === assistantMessageId
            ? finalizeAssistantMessageAfterAbort(message)
            : message,
        ),
      );
    }

    const cancelPromise = abortAgentChat(sessionIdRef.current).catch((error) => {
      console.error("[GlobalChatbox] Failed to abort agent session:", error);
    });
    const trackedCancelPromise = cancelPromise.finally(() => {
      if (cancelPromiseRef.current === trackedCancelPromise) {
        cancelPromiseRef.current = null;
      }
    });
    cancelPromiseRef.current = trackedCancelPromise;
  }, [getLastAssistantMessageId]);

  const replyPermission = useCallback(
    async (requestId: string, reply: PermissionReply) => {
      const target = messagesRef.current
        .flatMap((message) => message.permissions ?? [])
        .find((permission) => permission.requestId === requestId);
      if (!target || target.status === "submitting") {
        return;
      }

      setMessages((prev) =>
        prev.map((message) =>
          !message.permissions?.some((permission) => permission.requestId === requestId)
            ? message
            : {
                ...message,
                permissions: message.permissions.map((permission) =>
                  permission.requestId === requestId
                    ? { ...permission, status: "submitting", error: undefined }
                    : permission,
                ),
              },
        ),
      );

      try {
        await replyAgentPermission(target.sessionId, requestId, reply);
        setMessages((prev) =>
          prev.map((message) =>
            !message.permissions?.some((permission) => permission.requestId === requestId)
              ? message
              : {
                  ...message,
                  permissions: message.permissions.map((permission) =>
                    permission.requestId === requestId
                      ? {
                          ...permission,
                          status: toPermissionStatus(reply),
                          repliedAt: Date.now(),
                          error: undefined,
                        }
                      : permission,
                  ),
                },
          ),
        );
      } catch (error) {
        setMessages((prev) =>
          prev.map((message) =>
            !message.permissions?.some((permission) => permission.requestId === requestId)
              ? message
              : {
                  ...message,
                  permissions: message.permissions.map((permission) =>
                    permission.requestId === requestId
                      ? {
                          ...permission,
                          status: "error",
                          error: error instanceof Error ? error.message : String(error),
                        }
                      : permission,
                  ),
                },
          ),
        );
      }
    },
    [],
  );

  const createSession = useCallback(() => {
    if (isHydrating || isStreaming) return;

    const controller = abortRef.current;
    controller?.abort();
    setBranchTransition(null);
    hydrationNonceRef.current += 1;
    titleUpdateNonceRef.current += 1;
    sessionIdRef.current = undefined;
    lastPersistedStateKeyRef.current = createPersistedStateKey({
      title: "新对话",
      isTitleManuallyEdited: false,
      messages: [],
      sessionId: undefined,
      branchGroups: [],
    });
    setMessages([]);
    setSessionTitle("新对话");
    setIsSessionTitleManuallyEdited(false);
    setSessionId(undefined);
    setBranchGroups([]);
    setIsStreaming(false);
  }, [isHydrating, isStreaming]);

  const switchSession = useCallback(
    async (nextSessionId: string) => {
      if (isHydrating || isStreaming || sessionIdRef.current === nextSessionId) {
        return;
      }

      setIsHydrating(true);
      try {
        const [nextState, sessions] = await Promise.all([
          loadChatSessionById(nextSessionId),
          listChatSessions(),
        ]);

        hydrationNonceRef.current += 1;
        titleUpdateNonceRef.current += 1;
        sessionIdRef.current = nextState.sessionId;
        lastPersistedStateKeyRef.current = createPersistedStateKey(nextState);
        setBranchTransition(null);
        setMessages(nextState.messages);
        setSessionTitle(nextState.title);
        setIsSessionTitleManuallyEdited(nextState.isTitleManuallyEdited ?? false);
        setSessionId(nextState.sessionId);
        setBranchGroups(nextState.branchGroups);
        setChatSessions(sessions);
        if (nextState.sessionId && nextState.isStreaming) {
          resumeStreamingSession(nextState.sessionId);
        } else {
          setIsStreaming(false);
        }
      } catch (error) {
        console.error("[GlobalChatbox] Failed to switch chat session:", error);
      } finally {
        setIsHydrating(false);
      }
    },
    [isHydrating, isStreaming, resumeStreamingSession],
  );

  const removeSession = useCallback(
    async (targetSessionId: string) => {
      if (isHydrating || isStreaming) return;

      setChatSessions((prev) =>
        prev.filter((session) => session.id !== targetSessionId),
      );

      try {
        const nextActiveSessionId = await deleteChatSession(
          targetSessionId,
        );
        const sessions = await listChatSessions();
        setChatSessions(sessions);

        if (sessionIdRef.current !== targetSessionId) {
          return;
        }

        if (!nextActiveSessionId) {
          hydrationNonceRef.current += 1;
          titleUpdateNonceRef.current += 1;
          sessionIdRef.current = undefined;
          lastPersistedStateKeyRef.current = createPersistedStateKey({
            title: undefined,
            isTitleManuallyEdited: false,
            messages: [],
            sessionId: undefined,
            branchGroups: [],
          });
          setBranchTransition(null);
          setMessages([]);
          setSessionTitle(undefined);
          setIsSessionTitleManuallyEdited(false);
          setSessionId(undefined);
          setBranchGroups([]);
          return;
        }

        setIsHydrating(true);
        const [nextState, sessionsAfterDelete] = await Promise.all([
          loadChatSessionById(nextActiveSessionId),
          listChatSessions(),
        ]);
        hydrationNonceRef.current += 1;
        titleUpdateNonceRef.current += 1;
        sessionIdRef.current = nextState.sessionId;
        lastPersistedStateKeyRef.current = createPersistedStateKey(nextState);
        setBranchTransition(null);
        setMessages(nextState.messages);
        setSessionTitle(nextState.title);
        setIsSessionTitleManuallyEdited(nextState.isTitleManuallyEdited ?? false);
        setSessionId(nextState.sessionId);
        setBranchGroups(nextState.branchGroups);
        setChatSessions(sessionsAfterDelete);
      } catch (error) {
        console.error("[GlobalChatbox] Failed to delete chat session:", error);
        try {
          setChatSessions(await listChatSessions());
        } catch (refreshError) {
          console.error("[GlobalChatbox] Failed to refresh chat sessions:", refreshError);
        }
      } finally {
        setIsHydrating(false);
      }
    },
    [isHydrating, isStreaming],
  );

  const sendPrompt = useCallback(
    async (rawPrompt: string) => {
      await runPrompt({ prompt: rawPrompt });
    },
    [runPrompt],
  );

  const renameSession = useCallback(
    async (targetSessionId: string, nextTitle: string) => {
      const normalizedTitle = nextTitle.trim();
      if (!normalizedTitle || isHydrating) return;

      try {
        await updateChatSessionTitle(targetSessionId, normalizedTitle, {
          isTitleManuallyEdited: true,
        });
        const sessions = await listChatSessions();
        setChatSessions(sessions);

        if (sessionIdRef.current === targetSessionId) {
          setSessionTitle(normalizedTitle);
          setIsSessionTitleManuallyEdited(true);
          lastPersistedStateKeyRef.current = createPersistedStateKey({
            sessionId: targetSessionId,
            title: normalizedTitle,
            isTitleManuallyEdited: true,
            messages,
            branchGroups,
          });
        }
      } catch (error) {
        console.error("[GlobalChatbox] Failed to rename chat session:", error);
      }
    },
    [branchGroups, isHydrating, messages],
  );

  const regenerate = useCallback(async () => {
    if (isHydrating || isStreaming || messages.length === 0) return;

    let lastUserIndex = messages.length - 1;
    while (lastUserIndex >= 0 && messages[lastUserIndex].role !== "user") {
      lastUserIndex--;
    }

    if (lastUserIndex < 0) return;

      const lastUser = messages[lastUserIndex];
      const lastUserContent = lastUser.content;
      const nextMessages = cloneMessages(messages.slice(0, lastUserIndex));
      const nextUserMessage = createUserMessage(
        lastUserContent,
        lastUser.branchRootId ?? lastUser.id,
      );
      const nextAssistantMessage = createAssistantMessage();

      setMessages(nextMessages);
      await runPrompt({
        prompt: lastUserContent,
        preparedMessages: [
          ...nextMessages,
          nextUserMessage,
          nextAssistantMessage,
        ],
        userMessage: nextUserMessage,
        assistantMessage: nextAssistantMessage,
      });
    }, [isHydrating, isStreaming, messages, runPrompt]);

  const editAndResubmit = useCallback(
    async (messageId: string, newContent: string) => {
      if (isHydrating || isStreaming) return;

      const trimmedContent = newContent.trim();
      if (!trimmedContent) return;

      const messageIndex = messages.findIndex((m) => m.id === messageId);
      if (messageIndex < 0 || messages[messageIndex].role !== "user") return;

      const originalMessage = messages[messageIndex];
      if (trimmedContent === originalMessage.content.trim()) return;

      const rootMessageId = originalMessage.branchRootId ?? originalMessage.id;
      const currentSessionId = sessionIdRef.current;
      const keepMessageCount = messageIndex;
      const prefix = cloneMessages(messages.slice(0, messageIndex));
      const originalSuffix = cloneMessages(messages.slice(messageIndex));
      const forkedSessionId = await forkAgentChat(currentSessionId, keepMessageCount);

      const nextUserMessage = createUserMessage(trimmedContent, rootMessageId);
      const nextAssistantMessage = createAssistantMessage();
      const nextSuffix = [nextUserMessage, nextAssistantMessage];

      setBranchGroups((prev) => {
        const next = cloneBranchGroups(prev);
        const groupIndex = next.findIndex(
          (group) =>
            group.rootMessageId === rootMessageId && group.parentCount === messageIndex,
        );

        if (groupIndex >= 0) {
          const group = next[groupIndex];
          group.branches[group.activeIndex] = {
            ...group.branches[group.activeIndex],
            sessionId: currentSessionId,
            messages: originalSuffix,
          };
          group.branches.push({
            id: createId(),
            label: `分支 ${group.branches.length + 1}`,
            sessionId: forkedSessionId,
            messages: cloneMessages(nextSuffix),
          });
          group.activeIndex = group.branches.length - 1;
        } else {
          next.push({
            id: rootMessageId,
            rootMessageId,
            parentCount: messageIndex,
            activeIndex: 1,
            branches: [
              {
                id: createId(),
                label: "分支 1",
                sessionId: currentSessionId,
                messages: originalSuffix,
              },
              {
                id: createId(),
                label: "分支 2",
                sessionId: forkedSessionId,
                messages: cloneMessages(nextSuffix),
              },
            ],
          });
        }

        return next;
      });

      sessionIdRef.current = forkedSessionId;
      setSessionId(forkedSessionId);
      await runPrompt({
        prompt: trimmedContent,
        sessionIdOverride: forkedSessionId,
        preparedMessages: [...prefix, ...nextSuffix],
        userMessage: nextUserMessage,
        assistantMessage: nextAssistantMessage,
      });
    },
    [isHydrating, isStreaming, messages, runPrompt],
  );

  const cycleBranch = useCallback(
    (rootMessageId: string, direction: -1 | 1) => {
      if (isHydrating || isStreaming) return;

      setBranchGroups((prev) => {
        const next = cloneBranchGroups(prev);
        const group = next.find((item) => item.rootMessageId === rootMessageId);
        if (!group || group.branches.length < 2) {
          return prev;
        }

        const nextIndex =
          (group.activeIndex + direction + group.branches.length) % group.branches.length;
        const selectedBranch = group.branches[nextIndex];
        group.activeIndex = nextIndex;

      const nextMessages = [
          ...cloneMessages(messages.slice(0, group.parentCount)),
          ...cloneMessages(selectedBranch.messages),
        ];
        setBranchTransition({
          rootMessageId,
          parentCount: group.parentCount,
          activeBranchId: selectedBranch.id,
          nonce: Date.now(),
        });
        sessionIdRef.current = selectedBranch.sessionId;
        setSessionId(selectedBranch.sessionId);
        setMessages(nextMessages);

        return next;
      });
    },
    [isHydrating, isStreaming, messages],
  );

  return {
    messages,
    chatSessions,
    activeSessionId: sessionIdRef.current,
    branchGroups,
    branchTransition,
    isHydrating,
    isStreaming,
    sessionTitle,
    sessionId,
    sendPrompt,
    regenerate,
    editAndResubmit,
    cycleBranch,
    abort,
    replyPermission,
    createSession,
    renameSession,
    removeSession,
    switchSession,
  };
};