import { AgentManagementIdentifier } from '@lobechat/builtin-tool-agent-management';
import { act } from '@testing-library/react';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';

import { INPUT_LOADING_OPERATION_TYPES } from '@/store/chat/slices/operation/types';

import { type ConversationContext, type ConversationHooks } from '../../../types';
import { createStore } from '../../index';

// Mock useChatStore
const mockCancelOperations = vi.fn();
const mockCancelOperation = vi.fn();
const mockCancelSendMessageInServer = vi.fn();
const mockRegenerateUserMessage = vi.fn();
const mockRegenerateAssistantMessage = vi.fn();
const mockContinueGenerationMessage = vi.fn();
const mockDeleteMessage = vi.fn();
const mockSwitchMessageBranch = vi.fn();
const mockStartOperation = vi.fn(() => ({ operationId: 'test-op-id' }));
const mockCompleteOperation = vi.fn();
const mockFailOperation = vi.fn();
const mockExecuteClientAgent = vi.fn();
const mockIsGatewayModeEnabled = vi.fn(() => false);
const mockExecuteGatewayAgent = vi.fn();

vi.mock('@/store/chat', () => ({
  useChatStore: {
    getState: vi.fn(() => ({
      messagesMap: {
        'session-1-': [
          { id: 'msg-1', role: 'user', content: 'Hello' },
          { id: 'msg-2', role: 'assistant', content: 'Hi there', parentId: 'msg-1' },
        ],
      },
      operations: {},
      operationsByMessage: {},

      cancelOperations: mockCancelOperations,
      cancelOperation: mockCancelOperation,
      cancelSendMessageInServer: mockCancelSendMessageInServer,
      regenerateUserMessage: mockRegenerateUserMessage,
      regenerateAssistantMessage: mockRegenerateAssistantMessage,
      continueGenerationMessage: mockContinueGenerationMessage,
      deleteMessage: mockDeleteMessage,
      switchMessageBranch: mockSwitchMessageBranch,
      startOperation: mockStartOperation,
      completeOperation: mockCompleteOperation,
      failOperation: mockFailOperation,
      executeClientAgent: mockExecuteClientAgent,
      isGatewayModeEnabled: mockIsGatewayModeEnabled,
      executeGatewayAgent: mockExecuteGatewayAgent,
    })),
    setState: vi.fn(),
  },
}));

describe('Generation Actions', () => {
  beforeEach(() => {
    vi.clearAllMocks();
  });

  afterEach(() => {
    vi.clearAllTimers();
  });

  describe('stopGenerating', () => {
    it('should cancel all running execAgentRuntime operations', () => {
      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: 'topic-1',
        threadId: null,
      };

      const store = createStore({ context });

      act(() => {
        store.getState().stopGenerating();
      });

      expect(mockCancelOperations).toHaveBeenCalledWith(
        {
          type: INPUT_LOADING_OPERATION_TYPES,
          status: 'running',
          agentId: 'session-1',
          topicId: 'topic-1',
        },
        expect.any(String),
      );
    });

    it('should call onGenerationStop hook', () => {
      const onGenerationStop = vi.fn();
      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: null,
        threadId: null,
      };
      const hooks: ConversationHooks = { onGenerationStop };

      const store = createStore({ context, hooks });

      act(() => {
        store.getState().stopGenerating();
      });

      expect(onGenerationStop).toHaveBeenCalledTimes(1);
    });
  });

  describe('cancelOperation', () => {
    it('should cancel specific operation via ChatStore', () => {
      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: null,
        threadId: null,
      };

      const store = createStore({ context });

      act(() => {
        store.getState().cancelOperation('op-1', 'User cancelled');
      });

      expect(mockCancelOperation).toHaveBeenCalledWith('op-1', 'User cancelled');
    });

    it('should call onOperationCancelled hook', () => {
      const onOperationCancelled = vi.fn();
      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: null,
        threadId: null,
      };
      const hooks: ConversationHooks = { onOperationCancelled };

      const store = createStore({ context, hooks });

      act(() => {
        store.getState().cancelOperation('op-1');
      });

      expect(onOperationCancelled).toHaveBeenCalledWith('op-1');
    });
  });

  describe('clearOperations', () => {
    it('should be a no-op since operations are managed by ChatStore', () => {
      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: null,
        threadId: null,
      };

      const store = createStore({ context });

      // clearOperations should be callable without error
      act(() => {
        store.getState().clearOperations();
      });

      // No assertions needed - it's a no-op
    });
  });

  describe('continueGeneration', () => {
    it('should continue generation from assistantGroup message with last child as blockId', async () => {
      // Reset mock to ensure all required functions are available
      vi.mocked(await import('@/store/chat').then((m) => m.useChatStore.getState)).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},
        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        executeClientAgent: mockExecuteClientAgent,
        isGatewayModeEnabled: mockIsGatewayModeEnabled,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: 'topic-1',
        threadId: null,
        groupId: 'group-1',
      };

      const store = createStore({ context });

      // Set displayMessages with assistantGroup message containing children
      act(() => {
        store.setState({
          displayMessages: [
            {
              id: 'group-msg-1',
              role: 'assistantGroup',
              content: '',
              children: [
                { id: 'child-1', content: 'First response' },
                { id: 'child-2', content: 'Second response' },
              ],
            },
          ],
        } as any);
      });

      await act(async () => {
        await store.getState().continueGeneration('group-msg-1');
      });

      // Should create operation with groupMessageId
      expect(mockStartOperation).toHaveBeenCalledWith({
        context: { ...context, messageId: 'group-msg-1' },
        type: 'continue',
      });

      // Should call executeClientAgent with last child id as parentMessageId
      expect(mockExecuteClientAgent).toHaveBeenCalledWith(
        expect.objectContaining({
          context,
          parentMessageId: 'child-2', // last child's id
          parentMessageType: 'assistantGroup',
          parentOperationId: 'test-op-id',
        }),
      );
    });

    it('should not continue if message is not assistantGroup', async () => {
      // Reset mock to ensure all required functions are available
      vi.mocked(await import('@/store/chat').then((m) => m.useChatStore.getState)).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},
        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        executeClientAgent: mockExecuteClientAgent,
        isGatewayModeEnabled: mockIsGatewayModeEnabled,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: null,
        threadId: null,
      };

      const store = createStore({ context });

      // Set displayMessages with regular assistant message (not assistantGroup)
      act(() => {
        store.setState({
          displayMessages: [{ id: 'msg-1', role: 'assistant', content: 'Hello' }],
        } as any);
      });

      await act(async () => {
        await store.getState().continueGeneration('msg-1');
      });

      // Should not create operation if message is not assistantGroup
      expect(mockStartOperation).not.toHaveBeenCalled();
      expect(mockExecuteClientAgent).not.toHaveBeenCalled();
    });

    it('should not continue if assistantGroup has no children', async () => {
      // Reset mock to ensure all required functions are available
      vi.mocked(await import('@/store/chat').then((m) => m.useChatStore.getState)).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},
        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        executeClientAgent: mockExecuteClientAgent,
        isGatewayModeEnabled: mockIsGatewayModeEnabled,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: null,
        threadId: null,
      };

      const store = createStore({ context });

      // Set displayMessages with assistantGroup but no children
      act(() => {
        store.setState({
          displayMessages: [
            { id: 'group-msg-1', role: 'assistantGroup', content: '', children: [] },
          ],
        } as any);
      });

      await act(async () => {
        await store.getState().continueGeneration('group-msg-1');
      });

      // Should not create operation if no children
      expect(mockStartOperation).not.toHaveBeenCalled();
      expect(mockExecuteClientAgent).not.toHaveBeenCalled();
    });

    it('should call onBeforeContinue hook and respect false return', async () => {
      // Reset mock to ensure all required functions are available
      vi.mocked(await import('@/store/chat').then((m) => m.useChatStore.getState)).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},
        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        executeClientAgent: mockExecuteClientAgent,
        isGatewayModeEnabled: mockIsGatewayModeEnabled,
      } as any);

      const onBeforeContinue = vi.fn().mockResolvedValue(false);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: null,
        threadId: null,
      };
      const hooks: ConversationHooks = { onBeforeContinue };

      const store = createStore({ context, hooks });

      // Set displayMessages with assistantGroup
      act(() => {
        store.setState({
          displayMessages: [
            {
              id: 'group-msg-1',
              role: 'assistantGroup',
              content: '',
              children: [{ id: 'child-1', content: 'Response' }],
            },
          ],
        } as any);
      });

      await act(async () => {
        await store.getState().continueGeneration('group-msg-1');
      });

      expect(onBeforeContinue).toHaveBeenCalledWith('group-msg-1');
      // Should not call executeClientAgent if hook returns false
      expect(mockExecuteClientAgent).not.toHaveBeenCalled();
    });

    it('should call onContinueComplete hook after continuation', async () => {
      // Reset mock to ensure all required functions are available
      vi.mocked(await import('@/store/chat').then((m) => m.useChatStore.getState)).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},
        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        executeClientAgent: mockExecuteClientAgent.mockResolvedValue(undefined),
        isGatewayModeEnabled: mockIsGatewayModeEnabled,
      } as any);

      const onContinueComplete = vi.fn();

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: null,
        threadId: null,
      };
      const hooks: ConversationHooks = { onContinueComplete };

      const store = createStore({ context, hooks });

      // Set displayMessages with assistantGroup
      act(() => {
        store.setState({
          displayMessages: [
            {
              id: 'group-msg-1',
              role: 'assistantGroup',
              content: '',
              children: [{ id: 'child-1', content: 'Response' }],
            },
          ],
        } as any);
      });

      await act(async () => {
        await store.getState().continueGeneration('group-msg-1');
      });

      expect(onContinueComplete).toHaveBeenCalledWith('group-msg-1');
    });

    it('should not continue if message is not found', async () => {
      // Reset mock to ensure all required functions are available
      vi.mocked(await import('@/store/chat').then((m) => m.useChatStore.getState)).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},
        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        executeClientAgent: mockExecuteClientAgent,
        isGatewayModeEnabled: mockIsGatewayModeEnabled,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: null,
        threadId: null,
      };

      const store = createStore({ context });

      // Set empty displayMessages
      act(() => {
        store.setState({
          displayMessages: [],
        } as any);
      });

      await act(async () => {
        await store.getState().continueGeneration('msg-not-exist');
      });

      // Should not create operation if message not found
      expect(mockStartOperation).not.toHaveBeenCalled();
      expect(mockExecuteClientAgent).not.toHaveBeenCalled();
    });
  });

  describe('delAndRegenerateMessage', () => {
    it('should create operation with context and pass operationId to deleteMessage', async () => {
      // Re-setup mock to ensure all required functions are available
      const { useChatStore } = await import('@/store/chat');
      vi.mocked(useChatStore.getState).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},

        cancelOperations: mockCancelOperations,
        cancelOperation: mockCancelOperation,
        deleteMessage: mockDeleteMessage,
        switchMessageBranch: mockSwitchMessageBranch,
        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        executeClientAgent: mockExecuteClientAgent,
        isGatewayModeEnabled: mockIsGatewayModeEnabled,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: 'topic-1',
        threadId: null,
        groupId: 'group-1',
      };

      const store = createStore({ context });

      // Set displayMessages after store creation
      act(() => {
        store.setState({
          displayMessages: [
            { id: 'msg-1', role: 'user', content: 'Hello' },
            { id: 'msg-2', role: 'assistant', content: 'Hi there', parentId: 'msg-1' },
          ],
        } as any);
      });

      await act(async () => {
        await store.getState().delAndRegenerateMessage('msg-2');
      });

      // Should create operation with context
      expect(mockStartOperation).toHaveBeenCalledWith({
        context: { ...context, messageId: 'msg-2' },
        type: 'regenerate',
      });

      // Should pass operationId to deleteMessage for correct context isolation
      expect(mockDeleteMessage).toHaveBeenCalledWith('msg-2', { operationId: 'test-op-id' });

      // Should complete operation
      expect(mockCompleteOperation).toHaveBeenCalledWith('test-op-id');
    });

    it('should delete message BEFORE regeneration to prevent message not found issue ()', async () => {
      // This test verifies the fix:
      // When "delete and regenerate" is called, if regeneration happens first,
      // it switches to a new branch, causing the original message to no longer
      // appear in displayMessages. Then deleteMessage cannot find the message
      // and fails silently.
      //
      // The fix: delete first, then regenerate.

      const callOrder: string[] = [];

      // Re-setup mock to track call order
      const { useChatStore } = await import('@/store/chat');
      vi.mocked(useChatStore.getState).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},

        cancelOperations: mockCancelOperations,
        cancelOperation: mockCancelOperation,
        deleteMessage: vi.fn().mockImplementation(() => {
          callOrder.push('deleteMessage');
          return Promise.resolve();
        }),
        switchMessageBranch: vi.fn().mockImplementation(() => {
          callOrder.push('switchMessageBranch');
          return Promise.resolve();
        }),
        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        executeClientAgent: vi.fn().mockImplementation(() => {
          callOrder.push('executeClientAgent');
          return Promise.resolve();
        }),
        isGatewayModeEnabled: mockIsGatewayModeEnabled,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: 'topic-1',
        threadId: null,
        groupId: 'group-1',
      };

      const store = createStore({ context });

      // Set displayMessages and dbMessages
      act(() => {
        store.setState({
          displayMessages: [
            { id: 'msg-1', role: 'user', content: 'Hello' },
            { id: 'msg-2', role: 'assistant', content: 'Hi there', parentId: 'msg-1' },
          ],
          dbMessages: [
            { id: 'msg-1', role: 'user', content: 'Hello' },
            { id: 'msg-2', role: 'assistant', content: 'Hi there', parentId: 'msg-1' },
          ],
        } as any);
      });

      await act(async () => {
        await store.getState().delAndRegenerateMessage('msg-2');
      });

      // CRITICAL: deleteMessage must be called BEFORE switchMessageBranch and executeClientAgent
      // If regeneration (which calls switchMessageBranch) happens first, the message
      // won't be found in displayMessages and deletion will fail silently.
      expect(callOrder[0]).toBe('deleteMessage');
      expect(callOrder).toContain('switchMessageBranch');
      expect(callOrder).toContain('executeClientAgent');

      // Verify deleteMessage is called before any regeneration-related calls
      const deleteIndex = callOrder.indexOf('deleteMessage');
      const switchIndex = callOrder.indexOf('switchMessageBranch');
      const execIndex = callOrder.indexOf('executeClientAgent');

      expect(deleteIndex).toBeLessThan(switchIndex);
      expect(deleteIndex).toBeLessThan(execIndex);
    });

    it('should not proceed if assistant message has no parentId', async () => {
      const { useChatStore } = await import('@/store/chat');
      vi.mocked(useChatStore.getState).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},

        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        deleteMessage: mockDeleteMessage,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: null,
        threadId: null,
      };

      const store = createStore({ context });

      // Set displayMessages with assistant message that has no parentId
      act(() => {
        store.setState({
          displayMessages: [
            { id: 'msg-1', role: 'assistant', content: 'Hi there' }, // no parentId
          ],
        } as any);
      });

      await act(async () => {
        await store.getState().delAndRegenerateMessage('msg-1');
      });

      // Should not proceed - no operation created, no delete called
      expect(mockStartOperation).not.toHaveBeenCalled();
      expect(mockDeleteMessage).not.toHaveBeenCalled();
    });
  });

  describe('delAndResendThreadMessage', () => {
    it('should create operation with context and pass operationId to deleteMessage', async () => {
      // Re-setup mock to ensure startOperation is available
      const { useChatStore } = await import('@/store/chat');
      vi.mocked(useChatStore.getState).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},

        cancelOperations: mockCancelOperations,
        cancelOperation: mockCancelOperation,
        deleteMessage: mockDeleteMessage,
        switchMessageBranch: mockSwitchMessageBranch,
        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        executeClientAgent: mockExecuteClientAgent,
        isGatewayModeEnabled: mockIsGatewayModeEnabled,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: 'topic-1',
        threadId: 'thread-1',
        groupId: 'group-1',
      };

      const store = createStore({ context });

      // Set displayMessages after store creation
      act(() => {
        store.setState({
          displayMessages: [{ id: 'msg-1', role: 'user', content: 'Hello' }],
        } as any);
      });

      await act(async () => {
        await store.getState().delAndResendThreadMessage('msg-1');
      });

      // Should create operation with context including threadId
      expect(mockStartOperation).toHaveBeenCalledWith({
        context: { ...context, messageId: 'msg-1' },
        type: 'regenerate',
      });

      // Should pass operationId to deleteMessage
      expect(mockDeleteMessage).toHaveBeenCalledWith('msg-1', { operationId: 'test-op-id' });

      // Should complete operation
      expect(mockCompleteOperation).toHaveBeenCalledWith('test-op-id');
    });
  });

  describe('regenerateUserMessage', () => {
    it('should pass operationId to switchMessageBranch for correct context', async () => {
      // Re-setup mock with all required properties
      const { useChatStore } = await import('@/store/chat');
      vi.mocked(useChatStore.getState).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},

        cancelOperations: mockCancelOperations,
        cancelOperation: mockCancelOperation,
        deleteMessage: mockDeleteMessage,
        switchMessageBranch: mockSwitchMessageBranch,
        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        executeClientAgent: mockExecuteClientAgent,
        isGatewayModeEnabled: mockIsGatewayModeEnabled,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: 'topic-1',
        threadId: null,
        groupId: 'group-1',
      };

      const store = createStore({ context });

      // Set displayMessages and dbMessages after store creation
      // dbMessages is used to calculate children count for branch index
      act(() => {
        store.setState({
          displayMessages: [{ id: 'msg-1', role: 'user', content: 'Hello', branch: { count: 2 } }],
          dbMessages: [
            { id: 'msg-1', role: 'user', content: 'Hello' },
            { id: 'child-1', role: 'assistant', content: 'Response 1', parentId: 'msg-1' },
            { id: 'child-2', role: 'assistant', content: 'Response 2', parentId: 'msg-1' },
          ],
        } as any);
      });

      await act(async () => {
        await store.getState().regenerateUserMessage('msg-1');
      });

      // Should create operation with context
      expect(mockStartOperation).toHaveBeenCalledWith({
        context: { ...context, messageId: 'msg-1' },
        type: 'regenerate',
      });

      // Should pass operationId to switchMessageBranch
      // nextBranchIndex = childrenCount = 2 (two assistant messages with parentId: 'msg-1')
      expect(mockSwitchMessageBranch).toHaveBeenCalledWith('msg-1', 2, {
        operationId: 'test-op-id',
      });
    });

    it('should pass context to executeClientAgent', async () => {
      // Re-setup mock with all required properties
      const { useChatStore } = await import('@/store/chat');
      vi.mocked(useChatStore.getState).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},

        cancelOperations: mockCancelOperations,
        cancelOperation: mockCancelOperation,
        deleteMessage: mockDeleteMessage,
        switchMessageBranch: mockSwitchMessageBranch,
        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        executeClientAgent: mockExecuteClientAgent,
        isGatewayModeEnabled: mockIsGatewayModeEnabled,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: 'topic-1',
        threadId: null,
        groupId: 'group-1',
      };

      const store = createStore({ context });

      // Set displayMessages after store creation
      act(() => {
        store.setState({
          displayMessages: [{ id: 'msg-1', role: 'user', content: 'Hello' }],
        } as any);
      });

      await act(async () => {
        await store.getState().regenerateUserMessage('msg-1');
      });

      // Should pass full context to executeClientAgent
      expect(mockExecuteClientAgent).toHaveBeenCalledWith(
        expect.objectContaining({
          context,
          parentMessageId: 'msg-1',
          parentMessageType: 'user',
          parentOperationId: 'test-op-id',
        }),
      );
    });

    it('should restore mention-based initialContext when regenerating a user message', async () => {
      const { useChatStore } = await import('@/store/chat');
      vi.mocked(useChatStore.getState).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},

        cancelOperations: mockCancelOperations,
        cancelOperation: mockCancelOperation,
        deleteMessage: mockDeleteMessage,
        switchMessageBranch: mockSwitchMessageBranch,
        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        executeClientAgent: mockExecuteClientAgent,
        isGatewayModeEnabled: mockIsGatewayModeEnabled,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: 'topic-1',
        threadId: null,
      };

      const store = createStore({ context });

      act(() => {
        store.setState({
          displayMessages: [
            {
              id: 'msg-1',
              role: 'user',
              content: '<mention name="Agent A" id="agent-a" /> hello',
              editorData: {
                root: {
                  type: 'root',
                  children: [
                    {
                      type: 'paragraph',
                      children: [
                        {
                          type: 'mention',
                          label: 'Agent A',
                          metadata: { id: 'agent-a', type: 'agent' },
                        },
                        { type: 'text', text: ' hello' },
                      ],
                    },
                  ],
                },
              },
            },
          ],
        } as any);
      });

      await act(async () => {
        await store.getState().regenerateUserMessage('msg-1');
      });

      expect(mockExecuteClientAgent).toHaveBeenCalledWith(
        expect.objectContaining({
          initialContext: {
            initialContext: {
              mentionedAgents: [{ id: 'agent-a', name: 'Agent A' }],
              selectedTools: [{ identifier: AgentManagementIdentifier, name: 'Agent Management' }],
            },
            phase: 'init',
          },
        }),
      );
    });

    it('should use executeGatewayAgent when gateway mode is enabled', async () => {
      const { useChatStore } = await import('@/store/chat');
      vi.mocked(useChatStore.getState).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},

        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        isGatewayModeEnabled: vi.fn(() => true),
        executeGatewayAgent: mockExecuteGatewayAgent,
        executeClientAgent: mockExecuteClientAgent,
        switchMessageBranch: mockSwitchMessageBranch,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: 'topic-1',
        threadId: null,
      };

      const store = createStore({ context });

      act(() => {
        store.setState({
          displayMessages: [{ id: 'msg-1', role: 'user', content: 'Hello world' }],
        } as any);
      });

      await act(async () => {
        await store.getState().regenerateUserMessage('msg-1');
      });

      // Should switch message branch before gateway call
      expect(mockSwitchMessageBranch).toHaveBeenCalledWith('msg-1', 0, {
        operationId: 'test-op-id',
      });

      // Should call executeGatewayAgent with parentMessageId, original content, and onComplete
      expect(mockExecuteGatewayAgent).toHaveBeenCalledWith(
        expect.objectContaining({
          context,
          message: 'Hello world',
          parentMessageId: 'msg-1',
          onComplete: expect.any(Function),
        }),
      );

      // Should NOT call client-mode executeClientAgent
      expect(mockExecuteClientAgent).not.toHaveBeenCalled();

      // regenerate operation stays running until onComplete is called
      expect(mockCompleteOperation).not.toHaveBeenCalled();

      // Simulate gateway session complete
      const onComplete = mockExecuteGatewayAgent.mock.calls[0][0].onComplete;
      onComplete();
      expect(mockCompleteOperation).toHaveBeenCalledWith('test-op-id');
    });

    it('should call onRegenerateComplete hook after gateway regeneration', async () => {
      const { useChatStore } = await import('@/store/chat');
      vi.mocked(useChatStore.getState).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},

        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        isGatewayModeEnabled: vi.fn(() => true),
        executeGatewayAgent: mockExecuteGatewayAgent,
        switchMessageBranch: mockSwitchMessageBranch,
      } as any);

      const onRegenerateComplete = vi.fn();
      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: 'topic-1',
        threadId: null,
      };

      const store = createStore({ context, hooks: { onRegenerateComplete } });

      act(() => {
        store.setState({
          displayMessages: [{ id: 'msg-1', role: 'user', content: 'Hello' }],
        } as any);
      });

      await act(async () => {
        await store.getState().regenerateUserMessage('msg-1');
      });

      // Hook is called via onComplete callback, not directly
      expect(onRegenerateComplete).not.toHaveBeenCalled();

      // Simulate gateway session complete
      const onComplete = mockExecuteGatewayAgent.mock.calls[0][0].onComplete;
      onComplete();
      expect(onRegenerateComplete).toHaveBeenCalledWith('msg-1');
    });

    it('should fall back to client mode when gateway is disabled', async () => {
      const { useChatStore } = await import('@/store/chat');
      vi.mocked(useChatStore.getState).mockReturnValue({
        messagesMap: {},
        operations: {},
        operationsByMessage: {},

        startOperation: mockStartOperation,
        completeOperation: mockCompleteOperation,
        failOperation: mockFailOperation,
        isGatewayModeEnabled: vi.fn(() => false),
        executeGatewayAgent: mockExecuteGatewayAgent,
        executeClientAgent: mockExecuteClientAgent,
        switchMessageBranch: mockSwitchMessageBranch,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: 'topic-1',
        threadId: null,
      };

      const store = createStore({ context });

      act(() => {
        store.setState({
          displayMessages: [{ id: 'msg-1', role: 'user', content: 'Hello' }],
        } as any);
      });

      await act(async () => {
        await store.getState().regenerateUserMessage('msg-1');
      });

      // Should NOT call executeGatewayAgent
      expect(mockExecuteGatewayAgent).not.toHaveBeenCalled();

      // Should call client-mode executeClientAgent
      expect(mockExecuteClientAgent).toHaveBeenCalled();
    });

    it('should not regenerate if message is already loading', async () => {
      // Mock operation system to indicate message is processing
      const { useChatStore } = await import('@/store/chat');
      vi.mocked(useChatStore.getState).mockReturnValue({
        messagesMap: {},
        operations: {
          'op-1': {
            id: 'op-1',
            type: 'sendMessage',
            status: 'running',
            context: { messageIds: ['msg-1'] },
          },
        },
        operationsByMessage: { 'msg-1': ['op-1'] },
        startOperation: mockStartOperation,
      } as any);

      const context: ConversationContext = {
        agentId: 'session-1',
        topicId: null,
        threadId: null,
      };

      const store = createStore({ context });

      // Set displayMessages after store creation
      act(() => {
        store.setState({
          displayMessages: [{ id: 'msg-1', role: 'user', content: 'Hello' }],
        } as any);
      });

      await act(async () => {
        await store.getState().regenerateUserMessage('msg-1');
      });

      // Should not create operation if already regenerating
      expect(mockStartOperation).not.toHaveBeenCalled();
    });
  });
});
