/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

/* eslint-disable @typescript-eslint/no-explicit-any */
import { describe, it, expect, vi, beforeEach, afterEach, Mock } from 'vitest';
import { renderHook, act } from '@testing-library/react';
import {
  useReactToolScheduler,
  mapToDisplay,
} from './useReactToolScheduler.js';
import {
  Part,
  PartListUnion,
  PartUnion,
  FunctionResponse,
} from '@google/genai';
import {
  Config,
  ToolCallRequestInfo,
  Tool,
  ToolRegistry,
  ToolResult,
  ToolCallConfirmationDetails,
  ToolConfirmationOutcome,
  ToolCallResponseInfo,
  formatLlmContentForFunctionResponse, // Import from core
  ToolCall, // Import from core
  Status as ToolCallStatusType,
  ApprovalMode, // Import from core
} from '@gemini-code/core';
import {
  HistoryItemWithoutId,
  ToolCallStatus,
  HistoryItemToolGroup,
} from '../types.js';

// Mocks
vi.mock('@gemini-code/core', async () => {
  const actual = await vi.importActual('@gemini-code/core');
  return {
    ...actual,
    ToolRegistry: vi.fn(),
    Config: vi.fn(),
  };
});

const mockToolRegistry = {
  getTool: vi.fn(),
};

const mockConfig = {
  getToolRegistry: vi.fn(() => mockToolRegistry as unknown as ToolRegistry),
  getApprovalMode: vi.fn(() => ApprovalMode.DEFAULT),
};

const mockTool: Tool = {
  name: 'mockTool',
  displayName: 'Mock Tool',
  description: 'A mock tool for testing',
  isOutputMarkdown: false,
  canUpdateOutput: false,
  schema: {},
  validateToolParams: vi.fn(),
  execute: vi.fn(),
  shouldConfirmExecute: vi.fn(),
  getDescription: vi.fn((args) => `Description for ${JSON.stringify(args)}`),
};

const mockToolWithLiveOutput: Tool = {
  ...mockTool,
  name: 'mockToolWithLiveOutput',
  displayName: 'Mock Tool With Live Output',
  canUpdateOutput: true,
};

let mockOnUserConfirmForToolConfirmation: Mock;

const mockToolRequiresConfirmation: Tool = {
  ...mockTool,
  name: 'mockToolRequiresConfirmation',
  displayName: 'Mock Tool Requires Confirmation',
  shouldConfirmExecute: vi.fn(
    async (): Promise<ToolCallConfirmationDetails | false> => ({
      type: 'edit',
      title: 'Mock Tool Requires Confirmation',
      onConfirm: mockOnUserConfirmForToolConfirmation,
      fileName: 'mockToolRequiresConfirmation.ts',
      fileDiff: 'Mock tool requires confirmation',
    }),
  ),
};

describe('formatLlmContentForFunctionResponse', () => {
  it('should handle simple string llmContent', () => {
    const llmContent = 'Simple text output';
    const { functionResponseJson, additionalParts } =
      formatLlmContentForFunctionResponse(llmContent);
    expect(functionResponseJson).toEqual({ output: 'Simple text output' });
    expect(additionalParts).toEqual([]);
  });

  it('should handle llmContent as a single Part with text', () => {
    const llmContent: Part = { text: 'Text from Part object' };
    const { functionResponseJson, additionalParts } =
      formatLlmContentForFunctionResponse(llmContent);
    expect(functionResponseJson).toEqual({ output: 'Text from Part object' });
    expect(additionalParts).toEqual([]);
  });

  it('should handle llmContent as a PartListUnion array with a single text Part', () => {
    const llmContent: PartListUnion = [{ text: 'Text from array' }];
    const { functionResponseJson, additionalParts } =
      formatLlmContentForFunctionResponse(llmContent);
    expect(functionResponseJson).toEqual({ output: 'Text from array' });
    expect(additionalParts).toEqual([]);
  });

  it('should handle llmContent with inlineData', () => {
    const llmContent: Part = {
      inlineData: { mimeType: 'image/png', data: 'base64...' },
    };
    const { functionResponseJson, additionalParts } =
      formatLlmContentForFunctionResponse(llmContent);
    expect(functionResponseJson).toEqual({
      status: 'Binary content of type image/png was processed.',
    });
    expect(additionalParts).toEqual([llmContent]);
  });

  it('should handle llmContent with fileData', () => {
    const llmContent: Part = {
      fileData: { mimeType: 'application/pdf', fileUri: 'gs://...' },
    };
    const { functionResponseJson, additionalParts } =
      formatLlmContentForFunctionResponse(llmContent);
    expect(functionResponseJson).toEqual({
      status: 'Binary content of type application/pdf was processed.',
    });
    expect(additionalParts).toEqual([llmContent]);
  });

  it('should handle llmContent as an array of multiple Parts (text and inlineData)', () => {
    const llmContent: PartListUnion = [
      { text: 'Some textual description' },
      { inlineData: { mimeType: 'image/jpeg', data: 'base64data...' } },
      { text: 'Another text part' },
    ];
    const { functionResponseJson, additionalParts } =
      formatLlmContentForFunctionResponse(llmContent);
    expect(functionResponseJson).toEqual({
      status: 'Tool execution succeeded.',
    });
    expect(additionalParts).toEqual(llmContent);
  });

  it('should handle llmContent as an array with a single inlineData Part', () => {
    const llmContent: PartListUnion = [
      { inlineData: { mimeType: 'image/gif', data: 'gifdata...' } },
    ];
    const { functionResponseJson, additionalParts } =
      formatLlmContentForFunctionResponse(llmContent);
    expect(functionResponseJson).toEqual({
      status: 'Binary content of type image/gif was processed.',
    });
    expect(additionalParts).toEqual(llmContent);
  });

  it('should handle llmContent as a generic Part (not text, inlineData, or fileData)', () => {
    const llmContent: Part = { functionCall: { name: 'test', args: {} } };
    const { functionResponseJson, additionalParts } =
      formatLlmContentForFunctionResponse(llmContent);
    expect(functionResponseJson).toEqual({
      status: 'Tool execution succeeded.',
    });
    expect(additionalParts).toEqual([llmContent]);
  });

  it('should handle empty string llmContent', () => {
    const llmContent = '';
    const { functionResponseJson, additionalParts } =
      formatLlmContentForFunctionResponse(llmContent);
    expect(functionResponseJson).toEqual({ output: '' });
    expect(additionalParts).toEqual([]);
  });

  it('should handle llmContent as an empty array', () => {
    const llmContent: PartListUnion = [];
    const { functionResponseJson, additionalParts } =
      formatLlmContentForFunctionResponse(llmContent);
    expect(functionResponseJson).toEqual({
      status: 'Tool execution succeeded.',
    });
    expect(additionalParts).toEqual([]);
  });

  it('should handle llmContent as a Part with undefined inlineData/fileData/text', () => {
    const llmContent: Part = {}; // An empty part object
    const { functionResponseJson, additionalParts } =
      formatLlmContentForFunctionResponse(llmContent);
    expect(functionResponseJson).toEqual({
      status: 'Tool execution succeeded.',
    });
    expect(additionalParts).toEqual([llmContent]);
  });
});

describe('useReactToolScheduler in YOLO Mode', () => {
  let onComplete: Mock;
  let setPendingHistoryItem: Mock;

  beforeEach(() => {
    onComplete = vi.fn();
    setPendingHistoryItem = vi.fn();
    mockToolRegistry.getTool.mockClear();
    (mockToolRequiresConfirmation.execute as Mock).mockClear();
    (mockToolRequiresConfirmation.shouldConfirmExecute as Mock).mockClear();

    // IMPORTANT: Enable YOLO mode for this test suite
    (mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.YOLO);

    vi.useFakeTimers();
  });

  afterEach(() => {
    vi.clearAllTimers();
    vi.useRealTimers();
    // IMPORTANT: Disable YOLO mode after this test suite
    (mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.DEFAULT);
  });

  const renderSchedulerInYoloMode = () =>
    renderHook(() =>
      useReactToolScheduler(
        onComplete,
        mockConfig as unknown as Config,
        setPendingHistoryItem,
      ),
    );

  it('should skip confirmation and execute tool directly when yoloMode is true', async () => {
    mockToolRegistry.getTool.mockReturnValue(mockToolRequiresConfirmation);
    const expectedOutput = 'YOLO Confirmed output';
    (mockToolRequiresConfirmation.execute as Mock).mockResolvedValue({
      llmContent: expectedOutput,
      returnDisplay: 'YOLO Formatted tool output',
    } as ToolResult);

    const { result } = renderSchedulerInYoloMode();
    const schedule = result.current[1];
    const request: ToolCallRequestInfo = {
      callId: 'yoloCall',
      name: 'mockToolRequiresConfirmation',
      args: { data: 'any data' },
    };

    act(() => {
      schedule(request);
    });

    await act(async () => {
      await vi.runAllTimersAsync(); // Process validation
    });
    await act(async () => {
      await vi.runAllTimersAsync(); // Process scheduling
    });
    await act(async () => {
      await vi.runAllTimersAsync(); // Process execution
    });

    // Check that shouldConfirmExecute was NOT called
    expect(
      mockToolRequiresConfirmation.shouldConfirmExecute,
    ).not.toHaveBeenCalled();

    // Check that execute WAS called
    expect(mockToolRequiresConfirmation.execute).toHaveBeenCalledWith(
      request.args,
      expect.any(AbortSignal),
      undefined,
    );

    // Check that onComplete was called with success
    expect(onComplete).toHaveBeenCalledWith([
      expect.objectContaining({
        status: 'success',
        request,
        response: expect.objectContaining({
          resultDisplay: 'YOLO Formatted tool output',
          responseParts: expect.arrayContaining([
            expect.objectContaining({
              functionResponse: expect.objectContaining({
                response: { output: expectedOutput },
              }),
            }),
          ]),
        }),
      }),
    ]);

    // Ensure no confirmation UI was triggered (setPendingHistoryItem should not have been called with confirmation details)
    const setPendingHistoryItemCalls = setPendingHistoryItem.mock.calls;
    const confirmationCall = setPendingHistoryItemCalls.find((call) => {
      const item = typeof call[0] === 'function' ? call[0]({}) : call[0];
      return item?.tools?.[0]?.confirmationDetails;
    });
    expect(confirmationCall).toBeUndefined();
  });
});

describe('useReactToolScheduler', () => {
  // TODO(ntaylormullen): The following tests are skipped due to difficulties in
  // reliably testing the asynchronous state updates and interactions with timers.
  // These tests involve complex sequences of events, including confirmations,
  // live output updates, and cancellations, which are challenging to assert
  // correctly with the current testing setup. Further investigation is needed
  // to find a robust way to test these scenarios.
  let onComplete: Mock;
  let setPendingHistoryItem: Mock;
  let capturedOnConfirmForTest:
    | ((outcome: ToolConfirmationOutcome) => void | Promise<void>)
    | undefined;

  beforeEach(() => {
    onComplete = vi.fn();
    capturedOnConfirmForTest = undefined;
    setPendingHistoryItem = vi.fn((updaterOrValue) => {
      let pendingItem: HistoryItemWithoutId | null = null;
      if (typeof updaterOrValue === 'function') {
        // Loosen the type for prevState to allow for more flexible updates in tests
        const prevState: Partial<HistoryItemToolGroup> = {
          type: 'tool_group', // Still default to tool_group for most cases
          tools: [],
        };

        pendingItem = updaterOrValue(prevState as any); // Allow any for more flexibility
      } else {
        pendingItem = updaterOrValue;
      }
      // Capture onConfirm if it exists, regardless of the exact type of pendingItem
      // This is a common pattern in these tests.
      if (
        (pendingItem as HistoryItemToolGroup)?.tools?.[0]?.confirmationDetails
          ?.onConfirm
      ) {
        capturedOnConfirmForTest = (pendingItem as HistoryItemToolGroup)
          .tools[0].confirmationDetails?.onConfirm;
      }
    });

    mockToolRegistry.getTool.mockClear();
    (mockTool.execute as Mock).mockClear();
    (mockTool.shouldConfirmExecute as Mock).mockClear();
    (mockToolWithLiveOutput.execute as Mock).mockClear();
    (mockToolWithLiveOutput.shouldConfirmExecute as Mock).mockClear();
    (mockToolRequiresConfirmation.execute as Mock).mockClear();
    (mockToolRequiresConfirmation.shouldConfirmExecute as Mock).mockClear();

    mockOnUserConfirmForToolConfirmation = vi.fn();
    (
      mockToolRequiresConfirmation.shouldConfirmExecute as Mock
    ).mockImplementation(
      async (): Promise<ToolCallConfirmationDetails | null> => ({
        onConfirm: mockOnUserConfirmForToolConfirmation,
        fileName: 'mockToolRequiresConfirmation.ts',
        fileDiff: 'Mock tool requires confirmation',
        type: 'edit',
        title: 'Mock Tool Requires Confirmation',
      }),
    );

    vi.useFakeTimers();
  });

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

  const renderScheduler = () =>
    renderHook(() =>
      useReactToolScheduler(
        onComplete,
        mockConfig as unknown as Config,
        setPendingHistoryItem,
      ),
    );

  it('initial state should be empty', () => {
    const { result } = renderScheduler();
    expect(result.current[0]).toEqual([]);
  });

  it('should schedule and execute a tool call successfully', async () => {
    mockToolRegistry.getTool.mockReturnValue(mockTool);
    (mockTool.execute as Mock).mockResolvedValue({
      llmContent: 'Tool output',
      returnDisplay: 'Formatted tool output',
    } as ToolResult);
    (mockTool.shouldConfirmExecute as Mock).mockResolvedValue(null);

    const { result } = renderScheduler();
    const schedule = result.current[1];
    const request: ToolCallRequestInfo = {
      callId: 'call1',
      name: 'mockTool',
      args: { param: 'value' },
    };

    act(() => {
      schedule(request);
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(mockTool.execute).toHaveBeenCalledWith(
      request.args,
      expect.any(AbortSignal),
      undefined,
    );
    expect(onComplete).toHaveBeenCalledWith([
      expect.objectContaining({
        status: 'success',
        request,
        response: expect.objectContaining({
          resultDisplay: 'Formatted tool output',
          responseParts: expect.arrayContaining([
            expect.objectContaining({
              functionResponse: expect.objectContaining({
                response: { output: 'Tool output' },
              }),
            }),
          ]),
        }),
      }),
    ]);
    expect(result.current[0]).toEqual([]);
  });

  it('should handle tool not found', async () => {
    mockToolRegistry.getTool.mockReturnValue(undefined);
    const { result } = renderScheduler();
    const schedule = result.current[1];
    const request: ToolCallRequestInfo = {
      callId: 'call1',
      name: 'nonExistentTool',
      args: {},
    };

    act(() => {
      schedule(request);
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(onComplete).toHaveBeenCalledWith([
      expect.objectContaining({
        status: 'error',
        request,
        response: expect.objectContaining({
          error: expect.objectContaining({
            message: 'Tool "nonExistentTool" not found in registry.',
          }),
        }),
      }),
    ]);
    expect(result.current[0]).toEqual([]);
  });

  it('should handle error during shouldConfirmExecute', async () => {
    mockToolRegistry.getTool.mockReturnValue(mockTool);
    const confirmError = new Error('Confirmation check failed');
    (mockTool.shouldConfirmExecute as Mock).mockRejectedValue(confirmError);

    const { result } = renderScheduler();
    const schedule = result.current[1];
    const request: ToolCallRequestInfo = {
      callId: 'call1',
      name: 'mockTool',
      args: {},
    };

    act(() => {
      schedule(request);
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(onComplete).toHaveBeenCalledWith([
      expect.objectContaining({
        status: 'error',
        request,
        response: expect.objectContaining({
          error: confirmError,
        }),
      }),
    ]);
    expect(result.current[0]).toEqual([]);
  });

  it('should handle error during execute', async () => {
    mockToolRegistry.getTool.mockReturnValue(mockTool);
    (mockTool.shouldConfirmExecute as Mock).mockResolvedValue(null);
    const execError = new Error('Execution failed');
    (mockTool.execute as Mock).mockRejectedValue(execError);

    const { result } = renderScheduler();
    const schedule = result.current[1];
    const request: ToolCallRequestInfo = {
      callId: 'call1',
      name: 'mockTool',
      args: {},
    };

    act(() => {
      schedule(request);
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(onComplete).toHaveBeenCalledWith([
      expect.objectContaining({
        status: 'error',
        request,
        response: expect.objectContaining({
          error: execError,
        }),
      }),
    ]);
    expect(result.current[0]).toEqual([]);
  });

  it.skip('should handle tool requiring confirmation - approved', async () => {
    mockToolRegistry.getTool.mockReturnValue(mockToolRequiresConfirmation);
    const expectedOutput = 'Confirmed output';
    (mockToolRequiresConfirmation.execute as Mock).mockResolvedValue({
      llmContent: expectedOutput,
      returnDisplay: 'Confirmed display',
    } as ToolResult);

    const { result } = renderScheduler();
    const schedule = result.current[1];
    const request: ToolCallRequestInfo = {
      callId: 'callConfirm',
      name: 'mockToolRequiresConfirmation',
      args: { data: 'sensitive' },
    };

    act(() => {
      schedule(request);
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(setPendingHistoryItem).toHaveBeenCalled();
    expect(capturedOnConfirmForTest).toBeDefined();

    await act(async () => {
      await capturedOnConfirmForTest?.(ToolConfirmationOutcome.ProceedOnce);
    });

    await act(async () => {
      await vi.runAllTimersAsync();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(mockOnUserConfirmForToolConfirmation).toHaveBeenCalledWith(
      ToolConfirmationOutcome.ProceedOnce,
    );
    expect(mockToolRequiresConfirmation.execute).toHaveBeenCalled();
    expect(onComplete).toHaveBeenCalledWith([
      expect.objectContaining({
        status: 'success',
        request,
        response: expect.objectContaining({
          resultDisplay: 'Confirmed display',
          responseParts: expect.arrayContaining([
            expect.objectContaining({
              functionResponse: expect.objectContaining({
                response: { output: expectedOutput },
              }),
            }),
          ]),
        }),
      }),
    ]);
  });

  it.skip('should handle tool requiring confirmation - cancelled by user', async () => {
    mockToolRegistry.getTool.mockReturnValue(mockToolRequiresConfirmation);
    const { result } = renderScheduler();
    const schedule = result.current[1];
    const request: ToolCallRequestInfo = {
      callId: 'callConfirmCancel',
      name: 'mockToolRequiresConfirmation',
      args: {},
    };

    act(() => {
      schedule(request);
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(setPendingHistoryItem).toHaveBeenCalled();
    expect(capturedOnConfirmForTest).toBeDefined();

    await act(async () => {
      await capturedOnConfirmForTest?.(ToolConfirmationOutcome.Cancel);
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(mockOnUserConfirmForToolConfirmation).toHaveBeenCalledWith(
      ToolConfirmationOutcome.Cancel,
    );
    expect(onComplete).toHaveBeenCalledWith([
      expect.objectContaining({
        status: 'cancelled',
        request,
        response: expect.objectContaining({
          responseParts: expect.arrayContaining([
            expect.objectContaining({
              functionResponse: expect.objectContaining({
                response: expect.objectContaining({
                  error: `User did not allow tool call ${request.name}. Reason: User cancelled.`,
                }),
              }),
            }),
          ]),
        }),
      }),
    ]);
  });

  it.skip('should handle live output updates', async () => {
    mockToolRegistry.getTool.mockReturnValue(mockToolWithLiveOutput);
    let liveUpdateFn: ((output: string) => void) | undefined;
    let resolveExecutePromise: (value: ToolResult) => void;
    const executePromise = new Promise<ToolResult>((resolve) => {
      resolveExecutePromise = resolve;
    });

    (mockToolWithLiveOutput.execute as Mock).mockImplementation(
      async (
        _args: any,
        _signal: any,
        updateFn: ((output: string) => void) | undefined,
      ) => {
        liveUpdateFn = updateFn;
        return executePromise;
      },
    );
    (mockToolWithLiveOutput.shouldConfirmExecute as Mock).mockResolvedValue(
      null,
    );

    const { result } = renderScheduler();
    const schedule = result.current[1];
    const request: ToolCallRequestInfo = {
      callId: 'liveCall',
      name: 'mockToolWithLiveOutput',
      args: {},
    };

    act(() => {
      schedule(request);
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(liveUpdateFn).toBeDefined();
    expect(setPendingHistoryItem).toHaveBeenCalled();

    await act(async () => {
      liveUpdateFn?.('Live output 1');
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    await act(async () => {
      liveUpdateFn?.('Live output 2');
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    act(() => {
      resolveExecutePromise({
        llmContent: 'Final output',
        returnDisplay: 'Final display',
      } as ToolResult);
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(onComplete).toHaveBeenCalledWith([
      expect.objectContaining({
        status: 'success',
        request,
        response: expect.objectContaining({
          resultDisplay: 'Final display',
          responseParts: expect.arrayContaining([
            expect.objectContaining({
              functionResponse: expect.objectContaining({
                response: { output: 'Final output' },
              }),
            }),
          ]),
        }),
      }),
    ]);
    expect(result.current[0]).toEqual([]);
  });

  it.skip('should cancel tool calls before execution (e.g. when status is scheduled)', async () => {
    mockToolRegistry.getTool.mockReturnValue(mockTool);
    (mockTool.shouldConfirmExecute as Mock).mockResolvedValue(null);
    (mockTool.execute as Mock).mockReturnValue(new Promise(() => {}));

    const { result } = renderScheduler();
    const schedule = result.current[1];
    const cancel = result.current[2];
    const request: ToolCallRequestInfo = {
      callId: 'cancelCall',
      name: 'mockTool',
      args: {},
    };

    act(() => {
      schedule(request);
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    act(() => {
      cancel();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(onComplete).toHaveBeenCalledWith([
      expect.objectContaining({
        status: 'cancelled',
        request,
        response: expect.objectContaining({
          responseParts: expect.arrayContaining([
            expect.objectContaining({
              functionResponse: expect.objectContaining({
                response: expect.objectContaining({
                  error:
                    '[Operation Cancelled] Reason: User cancelled before execution',
                }),
              }),
            }),
          ]),
        }),
      }),
    ]);
    expect(mockTool.execute).not.toHaveBeenCalled();
    expect(result.current[0]).toEqual([]);
  });

  it.skip('should cancel tool calls that are awaiting approval', async () => {
    mockToolRegistry.getTool.mockReturnValue(mockToolRequiresConfirmation);
    const { result } = renderScheduler();
    const schedule = result.current[1];
    const cancelFn = result.current[2];
    const request: ToolCallRequestInfo = {
      callId: 'cancelApprovalCall',
      name: 'mockToolRequiresConfirmation',
      args: {},
    };

    act(() => {
      schedule(request);
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    act(() => {
      cancelFn();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(onComplete).toHaveBeenCalledWith([
      expect.objectContaining({
        status: 'cancelled',
        request,
        response: expect.objectContaining({
          responseParts: expect.arrayContaining([
            expect.objectContaining({
              functionResponse: expect.objectContaining({
                response: expect.objectContaining({
                  error:
                    '[Operation Cancelled] Reason: User cancelled during approval',
                }),
              }),
            }),
          ]),
        }),
      }),
    ]);
    expect(result.current[0]).toEqual([]);
  });

  it('should schedule and execute multiple tool calls', async () => {
    const tool1 = {
      ...mockTool,
      name: 'tool1',
      displayName: 'Tool 1',
      execute: vi.fn().mockResolvedValue({
        llmContent: 'Output 1',
        returnDisplay: 'Display 1',
      } as ToolResult),
      shouldConfirmExecute: vi.fn().mockResolvedValue(null),
    };
    const tool2 = {
      ...mockTool,
      name: 'tool2',
      displayName: 'Tool 2',
      execute: vi.fn().mockResolvedValue({
        llmContent: 'Output 2',
        returnDisplay: 'Display 2',
      } as ToolResult),
      shouldConfirmExecute: vi.fn().mockResolvedValue(null),
    };

    mockToolRegistry.getTool.mockImplementation((name) => {
      if (name === 'tool1') return tool1;
      if (name === 'tool2') return tool2;
      return undefined;
    });

    const { result } = renderScheduler();
    const schedule = result.current[1];
    const requests: ToolCallRequestInfo[] = [
      { callId: 'multi1', name: 'tool1', args: { p: 1 } },
      { callId: 'multi2', name: 'tool2', args: { p: 2 } },
    ];

    act(() => {
      schedule(requests);
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(onComplete).toHaveBeenCalledTimes(1);
    const completedCalls = onComplete.mock.calls[0][0] as ToolCall[];
    expect(completedCalls.length).toBe(2);

    const call1Result = completedCalls.find(
      (c) => c.request.callId === 'multi1',
    );
    const call2Result = completedCalls.find(
      (c) => c.request.callId === 'multi2',
    );

    expect(call1Result).toMatchObject({
      status: 'success',
      request: requests[0],
      response: expect.objectContaining({
        resultDisplay: 'Display 1',
        responseParts: expect.arrayContaining([
          expect.objectContaining({
            functionResponse: expect.objectContaining({
              response: { output: 'Output 1' },
            }),
          }),
        ]),
      }),
    });
    expect(call2Result).toMatchObject({
      status: 'success',
      request: requests[1],
      response: expect.objectContaining({
        resultDisplay: 'Display 2',
        responseParts: expect.arrayContaining([
          expect.objectContaining({
            functionResponse: expect.objectContaining({
              response: { output: 'Output 2' },
            }),
          }),
        ]),
      }),
    });
    expect(result.current[0]).toEqual([]);
  });

  it.skip('should throw error if scheduling while already running', async () => {
    mockToolRegistry.getTool.mockReturnValue(mockTool);
    const longExecutePromise = new Promise<ToolResult>((resolve) =>
      setTimeout(
        () => resolve({ llmContent: 'done', returnDisplay: 'done display' }),
        50,
      ),
    );
    (mockTool.execute as Mock).mockReturnValue(longExecutePromise);
    (mockTool.shouldConfirmExecute as Mock).mockResolvedValue(null);

    const { result } = renderScheduler();
    const schedule = result.current[1];
    const request1: ToolCallRequestInfo = {
      callId: 'run1',
      name: 'mockTool',
      args: {},
    };
    const request2: ToolCallRequestInfo = {
      callId: 'run2',
      name: 'mockTool',
      args: {},
    };

    act(() => {
      schedule(request1);
    });
    await act(async () => {
      await vi.runAllTimersAsync();
    });

    expect(() => schedule(request2)).toThrow(
      'Cannot schedule tool calls while other tool calls are running',
    );

    await act(async () => {
      await vi.advanceTimersByTimeAsync(50);
      await vi.runAllTimersAsync();
      await act(async () => {
        await vi.runAllTimersAsync();
      });
    });
    expect(onComplete).toHaveBeenCalledWith([
      expect.objectContaining({
        status: 'success',
        request: request1,
        response: expect.objectContaining({ resultDisplay: 'done display' }),
      }),
    ]);
    expect(result.current[0]).toEqual([]);
  });
});

describe('mapToDisplay', () => {
  const baseRequest: ToolCallRequestInfo = {
    callId: 'testCallId',
    name: 'testTool',
    args: { foo: 'bar' },
  };

  const baseTool: Tool = {
    name: 'testTool',
    displayName: 'Test Tool Display',
    description: 'Test Description',
    isOutputMarkdown: false,
    canUpdateOutput: false,
    schema: {},
    validateToolParams: vi.fn(),
    execute: vi.fn(),
    shouldConfirmExecute: vi.fn(),
    getDescription: vi.fn((args) => `Desc: ${JSON.stringify(args)}`),
  };

  const baseResponse: ToolCallResponseInfo = {
    callId: 'testCallId',
    responseParts: [
      {
        functionResponse: {
          name: 'testTool',
          id: 'testCallId',
          response: { output: 'Test output' },
        } as FunctionResponse,
      } as PartUnion,
    ],
    resultDisplay: 'Test display output',
    error: undefined,
  };

  // Define a more specific type for extraProps for these tests
  // This helps ensure that tool and confirmationDetails are only accessed when they are expected to exist.
  type MapToDisplayExtraProps =
    | {
        tool?: Tool;
        liveOutput?: string;
        response?: ToolCallResponseInfo;
        confirmationDetails?: ToolCallConfirmationDetails;
      }
    | {
        tool: Tool;
        response?: ToolCallResponseInfo;
        confirmationDetails?: ToolCallConfirmationDetails;
      }
    | {
        response: ToolCallResponseInfo;
        tool?: undefined;
        confirmationDetails?: ToolCallConfirmationDetails;
      }
    | {
        confirmationDetails: ToolCallConfirmationDetails;
        tool?: Tool;
        response?: ToolCallResponseInfo;
      };

  const testCases: Array<{
    name: string;
    status: ToolCallStatusType;
    extraProps?: MapToDisplayExtraProps;
    expectedStatus: ToolCallStatus;
    expectedResultDisplay?: string;
    expectedName?: string;
    expectedDescription?: string;
  }> = [
    {
      name: 'validating',
      status: 'validating',
      extraProps: { tool: baseTool },
      expectedStatus: ToolCallStatus.Executing,
      expectedName: baseTool.displayName,
      expectedDescription: baseTool.getDescription(baseRequest.args),
    },
    {
      name: 'awaiting_approval',
      status: 'awaiting_approval',
      extraProps: {
        tool: baseTool,
        confirmationDetails: {
          onConfirm: vi.fn(),
          type: 'edit',
          title: 'Test Tool Display',
          serverName: 'testTool',
          toolName: 'testTool',
          toolDisplayName: 'Test Tool Display',
          fileName: 'test.ts',
          fileDiff: 'Test diff',
        } as ToolCallConfirmationDetails,
      },
      expectedStatus: ToolCallStatus.Confirming,
      expectedName: baseTool.displayName,
      expectedDescription: baseTool.getDescription(baseRequest.args),
    },
    {
      name: 'scheduled',
      status: 'scheduled',
      extraProps: { tool: baseTool },
      expectedStatus: ToolCallStatus.Pending,
      expectedName: baseTool.displayName,
      expectedDescription: baseTool.getDescription(baseRequest.args),
    },
    {
      name: 'executing no live output',
      status: 'executing',
      extraProps: { tool: baseTool },
      expectedStatus: ToolCallStatus.Executing,
      expectedName: baseTool.displayName,
      expectedDescription: baseTool.getDescription(baseRequest.args),
    },
    {
      name: 'executing with live output',
      status: 'executing',
      extraProps: { tool: baseTool, liveOutput: 'Live test output' },
      expectedStatus: ToolCallStatus.Executing,
      expectedResultDisplay: 'Live test output',
      expectedName: baseTool.displayName,
      expectedDescription: baseTool.getDescription(baseRequest.args),
    },
    {
      name: 'success',
      status: 'success',
      extraProps: { tool: baseTool, response: baseResponse },
      expectedStatus: ToolCallStatus.Success,
      expectedResultDisplay: baseResponse.resultDisplay as any,
      expectedName: baseTool.displayName,
      expectedDescription: baseTool.getDescription(baseRequest.args),
    },
    {
      name: 'error tool not found',
      status: 'error',
      extraProps: {
        response: {
          ...baseResponse,
          error: new Error('Test error tool not found'),
          resultDisplay: 'Error display tool not found',
        },
      },
      expectedStatus: ToolCallStatus.Error,
      expectedResultDisplay: 'Error display tool not found',
      expectedName: baseRequest.name,
      expectedDescription: '',
    },
    {
      name: 'error tool execution failed',
      status: 'error',
      extraProps: {
        tool: baseTool,
        response: {
          ...baseResponse,
          error: new Error('Tool execution failed'),
          resultDisplay: 'Execution failed display',
        },
      },
      expectedStatus: ToolCallStatus.Error,
      expectedResultDisplay: 'Execution failed display',
      expectedName: baseTool.displayName, // Changed from baseTool.name
      expectedDescription: '',
    },
    {
      name: 'cancelled',
      status: 'cancelled',
      extraProps: {
        tool: baseTool,
        response: {
          ...baseResponse,
          resultDisplay: 'Cancelled display',
        },
      },
      expectedStatus: ToolCallStatus.Canceled,
      expectedResultDisplay: 'Cancelled display',
      expectedName: baseTool.displayName,
      expectedDescription: baseTool.getDescription(baseRequest.args),
    },
  ];

  testCases.forEach(
    ({
      name: testName,
      status,
      extraProps,
      expectedStatus,
      expectedResultDisplay,
      expectedName,
      expectedDescription,
    }) => {
      it(`should map ToolCall with status '${status}' (${testName}) correctly`, () => {
        const toolCall: ToolCall = {
          request: baseRequest,
          status,
          ...(extraProps || {}),
        } as ToolCall;

        const display = mapToDisplay(toolCall);
        expect(display.type).toBe('tool_group');
        expect(display.tools.length).toBe(1);
        const toolDisplay = display.tools[0];

        expect(toolDisplay.callId).toBe(baseRequest.callId);
        expect(toolDisplay.status).toBe(expectedStatus);
        expect(toolDisplay.resultDisplay).toBe(expectedResultDisplay);

        expect(toolDisplay.name).toBe(expectedName);

        if (status === 'error' && !extraProps?.tool) {
          expect(toolDisplay.description).toBe('');
        } else {
          expect(toolDisplay.description).toBe(
            expectedDescription ?? baseTool.getDescription(baseRequest.args),
          );
        }

        expect(toolDisplay.renderOutputAsMarkdown).toBe(
          extraProps?.tool?.isOutputMarkdown ?? false,
        );
        if (status === 'awaiting_approval') {
          expect(toolDisplay.confirmationDetails).toBe(
            extraProps!.confirmationDetails,
          );
        } else {
          expect(toolDisplay.confirmationDetails).toBeUndefined();
        }
      });
    },
  );

  it('should map an array of ToolCalls correctly', () => {
    const toolCall1: ToolCall = {
      request: { ...baseRequest, callId: 'call1' },
      status: 'success',
      tool: baseTool,
      response: { ...baseResponse, callId: 'call1' },
    } as ToolCall;
    const toolCall2: ToolCall = {
      request: { ...baseRequest, callId: 'call2' },
      status: 'executing',
      tool: { ...baseTool, isOutputMarkdown: true },
      liveOutput: 'markdown output',
    } as ToolCall;

    const display = mapToDisplay([toolCall1, toolCall2]);
    expect(display.tools.length).toBe(2);
    expect(display.tools[0].callId).toBe('call1');
    expect(display.tools[0].status).toBe(ToolCallStatus.Success);
    expect(display.tools[0].renderOutputAsMarkdown).toBe(false);
    expect(display.tools[1].callId).toBe('call2');
    expect(display.tools[1].status).toBe(ToolCallStatus.Executing);
    expect(display.tools[1].resultDisplay).toBe('markdown output');
    expect(display.tools[1].renderOutputAsMarkdown).toBe(true);
  });
});