"use client";

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

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

type UseAgentChatSessionOptions = {
  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({
    storageSessionId: state.storageSessionId ?? null,
    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 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 = ({
  onToolCall,
  onBeforeSend,
  getModel,
}: UseAgentChatSessionOptions) => {
  const storageSessionIdRef = useRef<string | undefined>(undefined);
  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 isSessionTitleManuallyEditedRef = useRef(false);
  const cancelPromiseRef = useRef<Promise<void> | null>(null);
  const titleUpdateNonceRef = useRef(0);
  const lastPersistedStateKeyRef = useRef(
    createPersistedStateKey({
      storageSessionId: undefined,
      title: undefined,
      isTitleManuallyEdited: false,
      messages: [],
      sessionId: undefined,
      branchGroups: [],
    }),
  );

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

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

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

    const hydrate = async () => {
      try {
        const [loadedState, sessions] = await Promise.all([
          loadActiveChatState(),
          listChatSessions(),
        ]);
        if (cancelled) return;

        storageSessionIdRef.current = loadedState.storageSessionId;
        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);
      } catch (error) {
        console.error("[GlobalChatbox] Failed to hydrate chat state:", error);
      } finally {
        if (!cancelled) {
          setIsHydrating(false);
        }
      }
    };

    void hydrate();

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

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

    const currentHydrationNonce = hydrationNonceRef.current;
    const persistTimer = window.setTimeout(() => {
      const state: LoadedChatState = {
        storageSessionId: storageSessionIdRef.current,
        title: sessionTitle,
        isTitleManuallyEdited: isSessionTitleManuallyEdited,
        messages,
        sessionId,
        branchGroups,
      };
      const currentStateKey = createPersistedStateKey(state);
      if (currentStateKey === lastPersistedStateKeyRef.current) {
        return;
      }

      void saveActiveChatState(state)
        .then((storageSessionId) => {
          if (hydrationNonceRef.current !== currentHydrationNonce) return;
          storageSessionIdRef.current = storageSessionId;
          lastPersistedStateKeyRef.current = createPersistedStateKey({
            ...state,
            storageSessionId,
          });
          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, messages, 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 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];

      setIsStreaming(true);
      setMessages(cloneMessages(nextMessages));
      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) => {
            if ("sessionId" in event && event.sessionId && event.sessionId !== sessionIdRef.current) {
              sessionIdRef.current = event.sessionId;
              setSessionId(event.sessionId);
            }

            if (event.type === "token") {
              setMessages((prev) =>
                prev.map((message) =>
                  message.id === nextAssistantMessage.id
                    ? {
                        ...message,
                        content: message.content + event.content,
                        isError: false,
                      }
                    : message,
                ),
              );
            } else if (event.type === "progress") {
              setMessages((prev) =>
                prev.map((message) =>
                  message.id === nextAssistantMessage.id
                    ? { ...message, progress: upsertProgress(message.progress, event) }
                    : message,
                ),
              );
            } else if (event.type === "tool_call") {
              onToolCall(event, {
                assistantMessageId: nextAssistantMessage.id,
                appendArtifact,
              });
            } else if (event.type === "session_title") {
              const nextTitle = event.title.trim();
              if (nextTitle && !isSessionTitleManuallyEditedRef.current) {
                setSessionTitle(nextTitle);
                const currentStorageSessionId = storageSessionIdRef.current;
                if (currentStorageSessionId) {
                  const currentNonce = ++titleUpdateNonceRef.current;
                  void updateChatSessionTitle(currentStorageSessionId, 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 !== nextAssistantMessage.id) 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 === nextAssistantMessage.id
                    ? {
                        ...message,
                        content: message.content || `⚠️ **错误：** ${event.message}`,
                        isError: true,
                        progress: completeRunningProgress(message.progress),
                      }
                    : message,
                ),
              );
              setIsStreaming(false);
            }
          },
        });
      } 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);
      }
    },
    [appendArtifact, getModel, isHydrating, isStreaming, messages, onBeforeSend, onToolCall],
  );

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

    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;
  }, []);

  const reset = useCallback(() => {
    const controller = abortRef.current;
    controller?.abort();
    const activeSessionId = sessionIdRef.current;
    if (activeSessionId) {
      const cancelPromise = abortAgentChat(activeSessionId).catch((error) => {
        console.error("[GlobalChatbox] Failed to abort agent session during reset:", error);
      });
      const trackedCancelPromise = cancelPromise.finally(() => {
        if (cancelPromiseRef.current === trackedCancelPromise) {
          cancelPromiseRef.current = null;
        }
      });
      cancelPromiseRef.current = trackedCancelPromise;
    }
    setMessages([]);
    setSessionTitle(undefined);
    setIsSessionTitleManuallyEdited(false);
    setBranchGroups([]);
    setBranchTransition(null);
    setSessionId(undefined);
    sessionIdRef.current = undefined;
    storageSessionIdRef.current = undefined;
    lastPersistedStateKeyRef.current = createPersistedStateKey({
      storageSessionId: undefined,
      title: undefined,
      isTitleManuallyEdited: false,
      messages: [],
      sessionId: undefined,
      branchGroups: [],
    });
    titleUpdateNonceRef.current += 1;
    setIsStreaming(false);
  }, []);

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

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

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

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

        hydrationNonceRef.current += 1;
        titleUpdateNonceRef.current += 1;
        storageSessionIdRef.current = nextState.storageSessionId;
        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);
      } catch (error) {
        console.error("[GlobalChatbox] Failed to switch chat session:", error);
      } finally {
        setIsHydrating(false);
      }
    },
    [isHydrating, isStreaming],
  );

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

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

        if (storageSessionIdRef.current !== targetStorageSessionId) {
          return;
        }

        if (!nextActiveSessionId) {
          hydrationNonceRef.current += 1;
          titleUpdateNonceRef.current += 1;
          storageSessionIdRef.current = undefined;
          sessionIdRef.current = undefined;
          lastPersistedStateKeyRef.current = createPersistedStateKey({
            storageSessionId: undefined,
            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;
        storageSessionIdRef.current = nextState.storageSessionId;
        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);
      } finally {
        setIsHydrating(false);
      }
    },
    [isHydrating, isStreaming],
  );

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

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

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

        if (storageSessionIdRef.current === targetStorageSessionId) {
          setSessionTitle(normalizedTitle);
          setIsSessionTitleManuallyEdited(true);
          lastPersistedStateKeyRef.current = createPersistedStateKey({
            storageSessionId: targetStorageSessionId,
            title: normalizedTitle,
            isTitleManuallyEdited: true,
            messages,
            sessionId: sessionIdRef.current,
            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,
    activeStorageSessionId: storageSessionIdRef.current,
    branchGroups,
    branchTransition,
    isHydrating,
    isStreaming,
    sessionTitle,
    sessionId,
    sendPrompt,
    regenerate,
    editAndResubmit,
    cycleBranch,
    abort,
    createSession,
    reset,
    renameSession,
    removeSession,
    switchSession,
  };
};