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

import {
  describe,
  it,
  expect,
  vi,
  beforeEach,
  afterEach,
  type Mock,
  type Mocked,
} from 'vitest';
import { GeminiAgent, Session } from './zedIntegration.js';
import * as acp from '@agentclientprotocol/sdk';
import {
  AuthType,
  ToolConfirmationOutcome,
  StreamEventType,
  isWithinRoot,
  ReadManyFilesTool,
  type GeminiChat,
  type Config,
  type MessageBus,
  LlmRole,
} from '@google/gemini-cli-core';
import {
  SettingScope,
  type LoadedSettings,
  loadSettings,
} from '../config/settings.js';
import { loadCliConfig, type CliArgs } from '../config/config.js';
import * as fs from 'node:fs/promises';
import * as path from 'node:path';
import { ApprovalMode } from '@google/gemini-cli-core/src/policy/types.js';

vi.mock('../config/config.js', () => ({
  loadCliConfig: vi.fn(),
}));

vi.mock('../config/settings.js', async (importOriginal) => {
  const actual = await importOriginal<typeof import('../config/settings.js')>();
  return {
    ...actual,
    loadSettings: vi.fn(),
  };
});

vi.mock('node:crypto', () => ({
  randomUUID: () => 'test-session-id',
}));

vi.mock('node:fs/promises');
vi.mock('node:path', async (importOriginal) => {
  const actual = await importOriginal<typeof import('node:path')>();
  return {
    ...actual,
    resolve: vi.fn(),
  };
});

// Mock ReadManyFilesTool
vi.mock(
  '@google/gemini-cli-core',
  async (
    importOriginal: () => Promise<typeof import('@google/gemini-cli-core')>,
  ) => {
    const actual = await importOriginal();
    return {
      ...actual,
      ReadManyFilesTool: vi.fn().mockImplementation(() => ({
        name: 'read_many_files',
        kind: 'read',
        build: vi.fn().mockReturnValue({
          getDescription: () => 'Read files',
          toolLocations: () => [],
          execute: vi.fn().mockResolvedValue({
            llmContent: ['--- file.txt ---\n\nFile content\n\n'],
          }),
        }),
      })),
      logToolCall: vi.fn(),
      isWithinRoot: vi.fn().mockReturnValue(true),
      LlmRole: {
        MAIN: 'main',
        SUBAGENT: 'subagent',
        UTILITY_TOOL: 'utility_tool',
        UTILITY_COMPRESSOR: 'utility_compressor',
        UTILITY_SUMMARIZER: 'utility_summarizer',
        UTILITY_ROUTER: 'utility_router',
        UTILITY_LOOP_DETECTOR: 'utility_loop_detector',
        UTILITY_NEXT_SPEAKER: 'utility_next_speaker',
        UTILITY_EDIT_CORRECTOR: 'utility_edit_corrector',
        UTILITY_AUTOCOMPLETE: 'utility_autocomplete',
        UTILITY_FAST_ACK_HELPER: 'utility_fast_ack_helper',
      },
      CoreToolCallStatus: {
        Validating: 'validating',
        Scheduled: 'scheduled',
        Error: 'error',
        Success: 'success',
        Executing: 'executing',
        Cancelled: 'cancelled',
        AwaitingApproval: 'awaiting_approval',
      },
    };
  },
);

// Helper to create mock streams
// eslint-disable-next-line @typescript-eslint/no-explicit-any
async function* createMockStream(items: any[]) {
  for (const item of items) {
    yield item;
  }
}

describe('GeminiAgent', () => {
  let mockConfig: Mocked<Awaited<ReturnType<typeof loadCliConfig>>>;
  let mockSettings: Mocked<LoadedSettings>;
  let mockArgv: CliArgs;
  let mockConnection: Mocked<acp.AgentSideConnection>;
  let agent: GeminiAgent;

  beforeEach(() => {
    mockConfig = {
      refreshAuth: vi.fn(),
      initialize: vi.fn(),
      getFileSystemService: vi.fn(),
      setFileSystemService: vi.fn(),
      getContentGeneratorConfig: vi.fn(),
      getActiveModel: vi.fn().mockReturnValue('gemini-pro'),
      getModel: vi.fn().mockReturnValue('gemini-pro'),
      getGeminiClient: vi.fn().mockReturnValue({
        startChat: vi.fn().mockResolvedValue({}),
      }),
      getMessageBus: vi.fn().mockReturnValue({
        publish: vi.fn(),
        subscribe: vi.fn(),
        unsubscribe: vi.fn(),
      }),
      getApprovalMode: vi.fn().mockReturnValue('default'),
      isPlanEnabled: vi.fn().mockReturnValue(false),
    } as unknown as Mocked<Awaited<ReturnType<typeof loadCliConfig>>>;
    mockSettings = {
      merged: {
        security: { auth: { selectedType: 'login_with_google' } },
        mcpServers: {},
      },
      setValue: vi.fn(),
    } as unknown as Mocked<LoadedSettings>;
    mockArgv = {} as unknown as CliArgs;
    mockConnection = {
      sessionUpdate: vi.fn(),
    } as unknown as Mocked<acp.AgentSideConnection>;

    (loadCliConfig as unknown as Mock).mockResolvedValue(mockConfig);
    (loadSettings as unknown as Mock).mockImplementation(() => ({
      merged: {
        security: { auth: { selectedType: AuthType.LOGIN_WITH_GOOGLE } },
        mcpServers: {},
      },
      setValue: vi.fn(),
    }));

    agent = new GeminiAgent(mockConfig, mockSettings, mockArgv, mockConnection);
  });

  it('should initialize correctly', async () => {
    const response = await agent.initialize({
      clientCapabilities: { fs: { readTextFile: true, writeTextFile: true } },
      protocolVersion: 1,
    });

    expect(response.protocolVersion).toBe(acp.PROTOCOL_VERSION);
    expect(response.authMethods).toHaveLength(3);
    const geminiAuth = response.authMethods?.find(
      (m) => m.id === AuthType.USE_GEMINI,
    );
    expect(geminiAuth?._meta).toEqual({
      'api-key': {
        provider: 'google',
      },
    });
    expect(response.agentCapabilities?.loadSession).toBe(true);
  });

  it('should authenticate correctly', async () => {
    await agent.authenticate({
      methodId: AuthType.LOGIN_WITH_GOOGLE,
    });

    expect(mockConfig.refreshAuth).toHaveBeenCalledWith(
      AuthType.LOGIN_WITH_GOOGLE,
      undefined,
    );
    expect(mockSettings.setValue).toHaveBeenCalledWith(
      SettingScope.User,
      'security.auth.selectedType',
      AuthType.LOGIN_WITH_GOOGLE,
    );
  });

  it('should authenticate correctly with api-key in _meta', async () => {
    await agent.authenticate({
      methodId: AuthType.USE_GEMINI,
      _meta: {
        'api-key': 'test-api-key',
      },
    } as unknown as acp.AuthenticateRequest);

    expect(mockConfig.refreshAuth).toHaveBeenCalledWith(
      AuthType.USE_GEMINI,
      'test-api-key',
    );
    expect(mockSettings.setValue).toHaveBeenCalledWith(
      SettingScope.User,
      'security.auth.selectedType',
      AuthType.USE_GEMINI,
    );
  });

  it('should create a new session', async () => {
    mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({
      apiKey: 'test-key',
    });
    const response = await agent.newSession({
      cwd: '/tmp',
      mcpServers: [],
    });

    expect(response.sessionId).toBe('test-session-id');
    expect(loadCliConfig).toHaveBeenCalled();
    expect(mockConfig.initialize).toHaveBeenCalled();
    expect(mockConfig.getGeminiClient).toHaveBeenCalled();
  });

  it('should return modes without plan mode when plan is disabled', async () => {
    mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({
      apiKey: 'test-key',
    });
    mockConfig.isPlanEnabled = vi.fn().mockReturnValue(false);
    mockConfig.getApprovalMode = vi.fn().mockReturnValue('default');

    const response = await agent.newSession({
      cwd: '/tmp',
      mcpServers: [],
    });

    expect(response.modes).toEqual({
      availableModes: [
        { id: 'default', name: 'Default', description: 'Prompts for approval' },
        {
          id: 'autoEdit',
          name: 'Auto Edit',
          description: 'Auto-approves edit tools',
        },
        { id: 'yolo', name: 'YOLO', description: 'Auto-approves all tools' },
      ],
      currentModeId: 'default',
    });
  });

  it('should return modes with plan mode when plan is enabled', async () => {
    mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({
      apiKey: 'test-key',
    });
    mockConfig.isPlanEnabled = vi.fn().mockReturnValue(true);
    mockConfig.getApprovalMode = vi.fn().mockReturnValue('plan');

    const response = await agent.newSession({
      cwd: '/tmp',
      mcpServers: [],
    });

    expect(response.modes).toEqual({
      availableModes: [
        { id: 'default', name: 'Default', description: 'Prompts for approval' },
        {
          id: 'autoEdit',
          name: 'Auto Edit',
          description: 'Auto-approves edit tools',
        },
        { id: 'yolo', name: 'YOLO', description: 'Auto-approves all tools' },
        { id: 'plan', name: 'Plan', description: 'Read-only mode' },
      ],
      currentModeId: 'plan',
    });
  });

  it('should fail session creation if Gemini API key is missing', async () => {
    (loadSettings as unknown as Mock).mockImplementation(() => ({
      merged: {
        security: { auth: { selectedType: AuthType.USE_GEMINI } },
        mcpServers: {},
      },
      setValue: vi.fn(),
    }));
    mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({
      apiKey: undefined,
    });

    await expect(
      agent.newSession({
        cwd: '/tmp',
        mcpServers: [],
      }),
    ).rejects.toMatchObject({
      message: 'Gemini API key is missing or not configured.',
    });
  });

  it('should create a new session with mcp servers', async () => {
    const mcpServers = [
      {
        name: 'test-server',
        command: 'node',
        args: ['server.js'],
        env: [{ name: 'KEY', value: 'VALUE' }],
      },
    ];

    await agent.newSession({
      cwd: '/tmp',
      mcpServers,
    });

    expect(loadCliConfig).toHaveBeenCalledWith(
      expect.objectContaining({
        mcpServers: expect.objectContaining({
          'test-server': expect.objectContaining({
            command: 'node',
            args: ['server.js'],
            env: { KEY: 'VALUE' },
          }),
        }),
      }),
      'test-session-id',
      mockArgv,
      { cwd: '/tmp' },
    );
  });

  it('should handle authentication failure gracefully', async () => {
    mockConfig.refreshAuth.mockRejectedValue(new Error('Auth failed'));
    const debugSpy = vi.spyOn(console, 'error').mockImplementation(() => {});

    // Should throw RequestError with custom message
    await expect(
      agent.newSession({
        cwd: '/tmp',
        mcpServers: [],
      }),
    ).rejects.toMatchObject({
      message: 'Auth failed',
    });

    debugSpy.mockRestore();
  });

  it('should initialize file system service if client supports it', async () => {
    agent = new GeminiAgent(mockConfig, mockSettings, mockArgv, mockConnection);
    await agent.initialize({
      clientCapabilities: { fs: { readTextFile: true, writeTextFile: true } },
      protocolVersion: 1,
    });

    await agent.newSession({
      cwd: '/tmp',
      mcpServers: [],
    });

    expect(mockConfig.setFileSystemService).toHaveBeenCalled();
  });

  it('should cancel a session', async () => {
    await agent.newSession({ cwd: '/tmp', mcpServers: [] });
    // Mock the session's cancelPendingPrompt
    const session = (
      agent as unknown as { sessions: Map<string, Session> }
    ).sessions.get('test-session-id');
    if (!session) throw new Error('Session not found');
    session.cancelPendingPrompt = vi.fn();

    await agent.cancel({ sessionId: 'test-session-id' });

    expect(session.cancelPendingPrompt).toHaveBeenCalled();
  });

  it('should throw error when cancelling non-existent session', async () => {
    await expect(agent.cancel({ sessionId: 'unknown' })).rejects.toThrow(
      'Session not found',
    );
  });

  it('should delegate prompt to session', async () => {
    await agent.newSession({ cwd: '/tmp', mcpServers: [] });
    const session = (
      agent as unknown as { sessions: Map<string, Session> }
    ).sessions.get('test-session-id');
    if (!session) throw new Error('Session not found');
    session.prompt = vi.fn().mockResolvedValue({ stopReason: 'end_turn' });

    const result = await agent.prompt({
      sessionId: 'test-session-id',
      prompt: [],
    });

    expect(session.prompt).toHaveBeenCalled();
    expect(result).toEqual({ stopReason: 'end_turn' });
  });

  it('should delegate setMode to session', async () => {
    await agent.newSession({ cwd: '/tmp', mcpServers: [] });
    const session = (
      agent as unknown as { sessions: Map<string, Session> }
    ).sessions.get('test-session-id');
    if (!session) throw new Error('Session not found');
    session.setMode = vi.fn().mockReturnValue({});

    const result = await agent.setSessionMode({
      sessionId: 'test-session-id',
      modeId: 'plan',
    });

    expect(session.setMode).toHaveBeenCalledWith('plan');
    expect(result).toEqual({});
  });

  it('should throw error when setting mode on non-existent session', async () => {
    await expect(
      agent.setSessionMode({
        sessionId: 'unknown',
        modeId: 'plan',
      }),
    ).rejects.toThrow('Session not found: unknown');
  });
});

describe('Session', () => {
  let mockChat: Mocked<GeminiChat>;
  let mockConfig: Mocked<Config>;
  let mockConnection: Mocked<acp.AgentSideConnection>;
  let session: Session;
  let mockToolRegistry: { getTool: Mock };
  let mockTool: { kind: string; build: Mock };
  let mockMessageBus: Mocked<MessageBus>;

  beforeEach(() => {
    mockChat = {
      sendMessageStream: vi.fn(),
      addHistory: vi.fn(),
      recordCompletedToolCalls: vi.fn(),
    } as unknown as Mocked<GeminiChat>;
    mockTool = {
      kind: 'read',
      build: vi.fn().mockReturnValue({
        getDescription: () => 'Test Tool',
        toolLocations: () => [],
        shouldConfirmExecute: vi.fn().mockResolvedValue(null),
        execute: vi.fn().mockResolvedValue({ llmContent: 'Tool Result' }),
      }),
    };
    mockToolRegistry = {
      getTool: vi.fn().mockReturnValue(mockTool),
    };
    mockMessageBus = {
      publish: vi.fn(),
      subscribe: vi.fn(),
      unsubscribe: vi.fn(),
    } as unknown as Mocked<MessageBus>;
    mockConfig = {
      getModel: vi.fn().mockReturnValue('gemini-pro'),
      getActiveModel: vi.fn().mockReturnValue('gemini-pro'),
      getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry),
      getFileService: vi.fn().mockReturnValue({
        shouldIgnoreFile: vi.fn().mockReturnValue(false),
      }),
      getFileFilteringOptions: vi.fn().mockReturnValue({}),
      getTargetDir: vi.fn().mockReturnValue('/tmp'),
      getEnableRecursiveFileSearch: vi.fn().mockReturnValue(false),
      getDebugMode: vi.fn().mockReturnValue(false),
      getMessageBus: vi.fn().mockReturnValue(mockMessageBus),
      setApprovalMode: vi.fn(),
      isPlanEnabled: vi.fn().mockReturnValue(false),
    } as unknown as Mocked<Config>;
    mockConnection = {
      sessionUpdate: vi.fn(),
      requestPermission: vi.fn(),
      sendNotification: vi.fn(),
    } as unknown as Mocked<acp.AgentSideConnection>;

    session = new Session('session-1', mockChat, mockConfig, mockConnection);
  });

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

  it('should handle prompt with text response', async () => {
    const stream = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: {
          candidates: [{ content: { parts: [{ text: 'Hello' }] } }],
        },
      },
    ]);
    mockChat.sendMessageStream.mockResolvedValue(stream);

    const result = await session.prompt({
      sessionId: 'session-1',
      prompt: [{ type: 'text', text: 'Hi' }],
    });

    expect(mockChat.sendMessageStream).toHaveBeenCalled();
    expect(mockConnection.sessionUpdate).toHaveBeenCalledWith({
      sessionId: 'session-1',
      update: {
        sessionUpdate: 'agent_message_chunk',
        content: { type: 'text', text: 'Hello' },
      },
    });
    expect(result).toEqual({ stopReason: 'end_turn' });
  });

  it('should handle tool calls', async () => {
    const stream1 = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: {
          functionCalls: [{ name: 'test_tool', args: { foo: 'bar' } }],
        },
      },
    ]);
    const stream2 = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: {
          candidates: [{ content: { parts: [{ text: 'Result' }] } }],
        },
      },
    ]);

    mockChat.sendMessageStream
      .mockResolvedValueOnce(stream1)
      .mockResolvedValueOnce(stream2);

    const result = await session.prompt({
      sessionId: 'session-1',
      prompt: [{ type: 'text', text: 'Call tool' }],
    });

    expect(mockToolRegistry.getTool).toHaveBeenCalledWith('test_tool');
    expect(mockTool.build).toHaveBeenCalledWith({ foo: 'bar' });
    expect(mockConnection.sessionUpdate).toHaveBeenCalledWith(
      expect.objectContaining({
        update: expect.objectContaining({
          sessionUpdate: 'tool_call',
          status: 'in_progress',
          kind: 'read',
        }),
      }),
    );
    expect(mockConnection.sessionUpdate).toHaveBeenCalledWith(
      expect.objectContaining({
        update: expect.objectContaining({
          sessionUpdate: 'tool_call_update',
          status: 'completed',
        }),
      }),
    );
    expect(result).toEqual({ stopReason: 'end_turn' });
  });

  it('should handle tool call permission request', async () => {
    const confirmationDetails = {
      type: 'info',
      onConfirm: vi.fn(),
    };
    mockTool.build.mockReturnValue({
      getDescription: () => 'Test Tool',
      toolLocations: () => [],
      shouldConfirmExecute: vi.fn().mockResolvedValue(confirmationDetails),
      execute: vi.fn().mockResolvedValue({ llmContent: 'Tool Result' }),
    });

    mockConnection.requestPermission.mockResolvedValue({
      outcome: {
        outcome: 'selected',
        optionId: ToolConfirmationOutcome.ProceedOnce,
      },
    });

    const stream1 = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: {
          functionCalls: [{ name: 'test_tool', args: {} }],
        },
      },
    ]);
    const stream2 = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: { candidates: [] },
      },
    ]);

    mockChat.sendMessageStream
      .mockResolvedValueOnce(stream1)
      .mockResolvedValueOnce(stream2);

    await session.prompt({
      sessionId: 'session-1',
      prompt: [{ type: 'text', text: 'Call tool' }],
    });

    expect(mockConnection.requestPermission).toHaveBeenCalled();
    expect(confirmationDetails.onConfirm).toHaveBeenCalledWith(
      ToolConfirmationOutcome.ProceedOnce,
    );
  });

  it('should handle tool call cancellation by user', async () => {
    const confirmationDetails = {
      type: 'info',
      onConfirm: vi.fn(),
    };
    mockTool.build.mockReturnValue({
      getDescription: () => 'Test Tool',
      toolLocations: () => [],
      shouldConfirmExecute: vi.fn().mockResolvedValue(confirmationDetails),
      execute: vi.fn().mockResolvedValue({ llmContent: 'Tool Result' }),
    });

    mockConnection.requestPermission.mockResolvedValue({
      outcome: { outcome: 'cancelled' },
    });

    const stream1 = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: {
          functionCalls: [{ name: 'test_tool', args: {} }],
        },
      },
    ]);
    const stream2 = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: { candidates: [] },
      },
    ]);

    mockChat.sendMessageStream
      .mockResolvedValueOnce(stream1)
      .mockResolvedValueOnce(stream2);

    await session.prompt({
      sessionId: 'session-1',
      prompt: [{ type: 'text', text: 'Call tool' }],
    });

    // When cancelled, it sends an error response to the model
    // We can verify that the second call to sendMessageStream contains the error
    expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2);
    const secondCallArgs = mockChat.sendMessageStream.mock.calls[1];
    const parts = secondCallArgs[1]; // parts
    expect(parts).toEqual(
      expect.arrayContaining([
        expect.objectContaining({
          functionResponse: expect.objectContaining({
            response: {
              error: expect.stringContaining('canceled by the user'),
            },
          }),
        }),
      ]),
    );
  });

  it('should include _meta.kind in diff tool calls', async () => {
    // Test 'add' (no original content)
    const addConfirmation = {
      type: 'edit',
      fileName: 'new.txt',
      originalContent: null,
      newContent: 'New content',
      onConfirm: vi.fn(),
    };

    // Test 'modify' (original and new content)
    const modifyConfirmation = {
      type: 'edit',
      fileName: 'existing.txt',
      originalContent: 'Old content',
      newContent: 'New content',
      onConfirm: vi.fn(),
    };

    // Test 'delete' (original content, no new content)
    const deleteConfirmation = {
      type: 'edit',
      fileName: 'deleted.txt',
      originalContent: 'Old content',
      newContent: '',
      onConfirm: vi.fn(),
    };

    const mockBuild = vi.fn();
    mockTool.build = mockBuild;

    // Helper to simulate tool call and check permission request
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    const checkDiffKind = async (confirmation: any, expectedKind: string) => {
      mockBuild.mockReturnValueOnce({
        getDescription: () => 'Test Tool',
        toolLocations: () => [],
        shouldConfirmExecute: vi.fn().mockResolvedValue(confirmation),
        execute: vi.fn().mockResolvedValue({ llmContent: 'Result' }),
      });

      mockConnection.requestPermission.mockResolvedValueOnce({
        outcome: {
          outcome: 'selected',
          optionId: ToolConfirmationOutcome.ProceedOnce,
        },
      });

      const stream = createMockStream([
        {
          type: StreamEventType.CHUNK,
          value: {
            functionCalls: [{ name: 'test_tool', args: {} }],
          },
        },
      ]);
      const emptyStream = createMockStream([]);

      mockChat.sendMessageStream
        .mockResolvedValueOnce(stream)
        .mockResolvedValueOnce(emptyStream);

      await session.prompt({
        sessionId: 'session-1',
        prompt: [{ type: 'text', text: 'Call tool' }],
      });

      expect(mockConnection.requestPermission).toHaveBeenCalledWith(
        expect.objectContaining({
          toolCall: expect.objectContaining({
            content: expect.arrayContaining([
              expect.objectContaining({
                type: 'diff',
                _meta: { kind: expectedKind },
              }),
            ]),
          }),
        }),
      );
    };

    await checkDiffKind(addConfirmation, 'add');
    await checkDiffKind(modifyConfirmation, 'modify');
    await checkDiffKind(deleteConfirmation, 'delete');
  });

  it('should handle @path resolution', async () => {
    (path.resolve as unknown as Mock).mockReturnValue('/tmp/file.txt');
    (fs.stat as unknown as Mock).mockResolvedValue({
      isDirectory: () => false,
    });
    (isWithinRoot as unknown as Mock).mockReturnValue(true);

    const stream = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: { candidates: [] },
      },
    ]);
    mockChat.sendMessageStream.mockResolvedValue(stream);

    await session.prompt({
      sessionId: 'session-1',
      prompt: [
        { type: 'text', text: 'Read' },
        {
          type: 'resource_link',
          uri: 'file://file.txt',
          mimeType: 'text/plain',
          name: 'file.txt',
        },
      ],
    });

    expect(path.resolve).toHaveBeenCalled();
    expect(fs.stat).toHaveBeenCalled();

    // Verify ReadManyFilesTool was used (implicitly by checking if sendMessageStream was called with resolved content)
    // Since we mocked ReadManyFilesTool to return specific content, we can check the args passed to sendMessageStream
    expect(mockChat.sendMessageStream).toHaveBeenCalledWith(
      expect.anything(),
      expect.arrayContaining([
        expect.objectContaining({
          text: expect.stringContaining('Content from @file.txt'),
        }),
      ]),
      expect.anything(),
      expect.any(AbortSignal),
      LlmRole.MAIN,
    );
  });

  it('should handle cancellation during prompt', async () => {
    let streamController: ReadableStreamDefaultController<unknown>;
    const stream = new ReadableStream({
      start(controller) {
        streamController = controller;
      },
    });

    let streamStarted: (value: unknown) => void;
    const streamStartedPromise = new Promise((resolve) => {
      streamStarted = resolve;
    });

    // Adapt web stream to async iterable
    async function* asyncStream() {
      process.stdout.write('TEST: asyncStream started\n');
      streamStarted(true);
      const reader = stream.getReader();
      try {
        while (true) {
          process.stdout.write('TEST: waiting for read\n');
          const { done, value } = await reader.read();
          process.stdout.write(`TEST: read returned done=${done}\n`);
          if (done) break;
          yield value;
        }
      } finally {
        process.stdout.write('TEST: releasing lock\n');
        reader.releaseLock();
      }
    }

    mockChat.sendMessageStream.mockResolvedValue(asyncStream());

    process.stdout.write('TEST: calling prompt\n');
    const promptPromise = session.prompt({
      sessionId: 'session-1',
      prompt: [{ type: 'text', text: 'Hi' }],
    });

    process.stdout.write('TEST: waiting for streamStarted\n');
    await streamStartedPromise;
    process.stdout.write('TEST: streamStarted\n');
    await session.cancelPendingPrompt();
    process.stdout.write('TEST: cancelled\n');

    // Close the stream to allow prompt loop to continue and check aborted signal
    streamController!.close();
    process.stdout.write('TEST: stream closed\n');

    const result = await promptPromise;
    process.stdout.write(`TEST: result received ${JSON.stringify(result)}\n`);
    expect(result).toEqual({ stopReason: 'cancelled' });
  });

  it('should handle rate limit error', async () => {
    const error = new Error('Rate limit');
    (error as unknown as { status: number }).status = 429;
    mockChat.sendMessageStream.mockRejectedValue(error);

    await expect(
      session.prompt({
        sessionId: 'session-1',
        prompt: [{ type: 'text', text: 'Hi' }],
      }),
    ).rejects.toMatchObject({
      code: 429,
      message: 'Rate limit exceeded. Try again later.',
    });
  });

  it('should handle tool execution error', async () => {
    mockTool.build.mockReturnValue({
      getDescription: () => 'Test Tool',
      toolLocations: () => [],
      shouldConfirmExecute: vi.fn().mockResolvedValue(null),
      execute: vi.fn().mockRejectedValue(new Error('Tool failed')),
    });

    const stream1 = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: {
          functionCalls: [{ name: 'test_tool', args: {} }],
        },
      },
    ]);
    const stream2 = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: { candidates: [] },
      },
    ]);

    mockChat.sendMessageStream
      .mockResolvedValueOnce(stream1)
      .mockResolvedValueOnce(stream2);

    await session.prompt({
      sessionId: 'session-1',
      prompt: [{ type: 'text', text: 'Call tool' }],
    });

    expect(mockConnection.sessionUpdate).toHaveBeenCalledWith(
      expect.objectContaining({
        update: expect.objectContaining({
          sessionUpdate: 'tool_call_update',
          status: 'failed',
          content: expect.arrayContaining([
            expect.objectContaining({
              content: expect.objectContaining({ text: 'Tool failed' }),
            }),
          ]),
        }),
      }),
    );
  });

  it('should handle missing tool', async () => {
    mockToolRegistry.getTool.mockReturnValue(undefined);

    const stream1 = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: {
          functionCalls: [{ name: 'unknown_tool', args: {} }],
        },
      },
    ]);
    const stream2 = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: { candidates: [] },
      },
    ]);

    mockChat.sendMessageStream
      .mockResolvedValueOnce(stream1)
      .mockResolvedValueOnce(stream2);

    await session.prompt({
      sessionId: 'session-1',
      prompt: [{ type: 'text', text: 'Call tool' }],
    });

    // Should send error response to model
    expect(mockChat.sendMessageStream).toHaveBeenCalledTimes(2);
    const secondCallArgs = mockChat.sendMessageStream.mock.calls[1];
    const parts = secondCallArgs[1];
    expect(parts).toEqual(
      expect.arrayContaining([
        expect.objectContaining({
          functionResponse: expect.objectContaining({
            response: {
              error: expect.stringContaining('not found in registry'),
            },
          }),
        }),
      ]),
    );
  });

  it('should ignore files based on configuration', async () => {
    (
      mockConfig.getFileService().shouldIgnoreFile as unknown as Mock
    ).mockReturnValue(true);
    const stream = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: { candidates: [] },
      },
    ]);
    mockChat.sendMessageStream.mockResolvedValue(stream);

    await session.prompt({
      sessionId: 'session-1',
      prompt: [
        {
          type: 'resource_link',
          uri: 'file://ignored.txt',
          mimeType: 'text/plain',
          name: 'ignored.txt',
        },
      ],
    });

    // Should not read file
    expect(mockToolRegistry.getTool).not.toHaveBeenCalledWith(
      'read_many_files',
    );
  });

  it('should handle directory resolution with glob', async () => {
    (path.resolve as unknown as Mock).mockReturnValue('/tmp/dir');
    (fs.stat as unknown as Mock).mockResolvedValue({
      isDirectory: () => true,
    });
    (isWithinRoot as unknown as Mock).mockReturnValue(true);

    const stream = createMockStream([
      {
        type: StreamEventType.CHUNK,
        value: { candidates: [] },
      },
    ]);
    mockChat.sendMessageStream.mockResolvedValue(stream);

    await session.prompt({
      sessionId: 'session-1',
      prompt: [
        {
          type: 'resource_link',
          uri: 'file://dir',
          mimeType: 'text/plain',
          name: 'dir',
        },
      ],
    });

    // Should use glob
    // ReadManyFilesTool is instantiated directly, so we check if the mock instance's build method was called
    const MockReadManyFilesTool = ReadManyFilesTool as unknown as Mock;
    const mockInstance =
      MockReadManyFilesTool.mock.results[
        MockReadManyFilesTool.mock.results.length - 1
      ].value;
    expect(mockInstance.build).toHaveBeenCalled();
  });

  it('should set mode on config', () => {
    session.setMode(ApprovalMode.AUTO_EDIT);
    expect(mockConfig.setApprovalMode).toHaveBeenCalledWith(
      ApprovalMode.AUTO_EDIT,
    );
  });

  it('should throw error for invalid mode', () => {
    expect(() => session.setMode('invalid-mode')).toThrow(
      'Invalid or unavailable mode: invalid-mode',
    );
  });
});