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

import { describe, it, expect, vi, beforeEach } from 'vitest';
import { SchedulerStateManager } from './state-manager.js';
import {
  CoreToolCallStatus,
  ROOT_SCHEDULER_ID,
  type ValidatingToolCall,
  type WaitingToolCall,
  type SuccessfulToolCall,
  type ErroredToolCall,
  type CancelledToolCall,
  type ExecutingToolCall,
  type ToolCallRequestInfo,
  type ToolCallResponseInfo,
} from './types.js';
import {
  ToolConfirmationOutcome,
  type AnyDeclarativeTool,
  type AnyToolInvocation,
  type FileDiff,
} from '../tools/tools.js';
import { MessageBusType } from '../confirmation-bus/types.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import { ApprovalMode } from '../policy/types.js';

describe('SchedulerStateManager', () => {
  const mockRequest: ToolCallRequestInfo = {
    callId: 'call-1',
    name: 'test-tool',
    args: { foo: 'bar' },
    isClientInitiated: false,
    prompt_id: 'prompt-1',
  };

  const mockTool = {
    name: 'test-tool',
    displayName: 'Test Tool',
  } as AnyDeclarativeTool;

  const mockInvocation = {
    shouldConfirmExecute: vi.fn(),
    execute: vi.fn(),
    getDescription: vi.fn(),
  } as unknown as AnyToolInvocation;

  const createValidatingCall = (
    id = 'call-1',
    mode: ApprovalMode = ApprovalMode.DEFAULT,
  ): ValidatingToolCall => ({
    status: CoreToolCallStatus.Validating,
    request: { ...mockRequest, callId: id },
    tool: mockTool,
    invocation: mockInvocation,
    startTime: Date.now(),
    approvalMode: mode,
  });

  const createMockResponse = (id: string): ToolCallResponseInfo => ({
    callId: id,
    responseParts: [],
    resultDisplay: 'Success',
    error: undefined,
    errorType: undefined,
  });

  let stateManager: SchedulerStateManager;
  let mockMessageBus: MessageBus;
  let onUpdate: (calls: unknown[]) => void;

  beforeEach(() => {
    onUpdate = vi.fn();
    mockMessageBus = {
      publish: vi.fn(),
      subscribe: vi.fn(),
      unsubscribe: vi.fn(),
    } as unknown as MessageBus;

    // Capture the update when published
    vi.mocked(mockMessageBus.publish).mockImplementation((msg) => {
      // Return a Promise to satisfy the void | Promise<void> signature if needed,
      // though typically mocks handle it.
      if (msg.type === MessageBusType.TOOL_CALLS_UPDATE) {
        onUpdate(msg.toolCalls);
      }
      return Promise.resolve();
    });

    stateManager = new SchedulerStateManager(mockMessageBus);
  });

  describe('Observer Callback', () => {
    it('should trigger onTerminalCall when finalizing a call', () => {
      const onTerminalCall = vi.fn();
      const manager = new SchedulerStateManager(
        mockMessageBus,
        ROOT_SCHEDULER_ID,
        onTerminalCall,
      );
      const call = createValidatingCall();
      manager.enqueue([call]);
      manager.dequeue();
      manager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Success,
        createMockResponse(call.request.callId),
      );
      manager.finalizeCall(call.request.callId);

      expect(onTerminalCall).toHaveBeenCalledTimes(1);
      expect(onTerminalCall).toHaveBeenCalledWith(
        expect.objectContaining({
          status: CoreToolCallStatus.Success,
          request: expect.objectContaining({ callId: call.request.callId }),
        }),
      );
    });

    it('should trigger onTerminalCall for every call in cancelAllQueued', () => {
      const onTerminalCall = vi.fn();
      const manager = new SchedulerStateManager(
        mockMessageBus,
        ROOT_SCHEDULER_ID,
        onTerminalCall,
      );
      manager.enqueue([createValidatingCall('1'), createValidatingCall('2')]);

      manager.cancelAllQueued('Test cancel');

      expect(onTerminalCall).toHaveBeenCalledTimes(2);
      expect(onTerminalCall).toHaveBeenCalledWith(
        expect.objectContaining({
          status: CoreToolCallStatus.Cancelled,
          request: expect.objectContaining({ callId: '1' }),
        }),
      );
      expect(onTerminalCall).toHaveBeenCalledWith(
        expect.objectContaining({
          status: CoreToolCallStatus.Cancelled,
          request: expect.objectContaining({ callId: '2' }),
        }),
      );
    });
  });

  describe('Initialization', () => {
    it('should start with empty state', () => {
      expect(stateManager.isActive).toBe(false);
      expect(stateManager.activeCallCount).toBe(0);
      expect(stateManager.queueLength).toBe(0);
      expect(stateManager.getSnapshot()).toEqual([]);
    });
  });

  describe('Lookup Operations', () => {
    it('should find tool calls in active calls', () => {
      const call = createValidatingCall('active-1');
      stateManager.enqueue([call]);
      stateManager.dequeue();
      expect(stateManager.getToolCall('active-1')).toEqual(call);
    });

    it('should find tool calls in the queue', () => {
      const call = createValidatingCall('queued-1');
      stateManager.enqueue([call]);
      expect(stateManager.getToolCall('queued-1')).toEqual(call);
    });

    it('should find tool calls in the completed batch', () => {
      const call = createValidatingCall('completed-1');
      stateManager.enqueue([call]);
      stateManager.dequeue();
      stateManager.updateStatus(
        'completed-1',
        CoreToolCallStatus.Success,
        createMockResponse('completed-1'),
      );
      stateManager.finalizeCall('completed-1');
      expect(stateManager.getToolCall('completed-1')).toBeDefined();
    });

    it('should return undefined for non-existent callIds', () => {
      expect(stateManager.getToolCall('void')).toBeUndefined();
    });
  });

  describe('Queue Management', () => {
    it('should enqueue calls and notify', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);

      expect(stateManager.queueLength).toBe(1);
      expect(onUpdate).toHaveBeenCalledWith([call]);
    });

    it('should dequeue calls and notify', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);

      const dequeued = stateManager.dequeue();

      expect(dequeued).toEqual(call);
      expect(stateManager.queueLength).toBe(0);
      expect(stateManager.activeCallCount).toBe(1);
      expect(onUpdate).toHaveBeenCalled();
    });

    it('should return undefined when dequeueing from empty queue', () => {
      const dequeued = stateManager.dequeue();
      expect(dequeued).toBeUndefined();
    });
  });

  describe('Status Transitions', () => {
    it('should transition validating to scheduled', () => {
      const call = createValidatingCall('call-1', ApprovalMode.PLAN);
      stateManager.enqueue([call]);
      stateManager.dequeue();

      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Scheduled,
      );

      const snapshot = stateManager.getSnapshot();
      expect(snapshot[0].status).toBe(CoreToolCallStatus.Scheduled);
      expect(snapshot[0].request.callId).toBe(call.request.callId);
      expect(snapshot[0].approvalMode).toBe(ApprovalMode.PLAN);
    });

    it('should transition scheduled to executing', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Scheduled,
      );

      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Executing,
      );

      expect(stateManager.firstActiveCall?.status).toBe(
        CoreToolCallStatus.Executing,
      );
    });

    it('should transition to success and move to completed batch', () => {
      const call = createValidatingCall('call-1', ApprovalMode.PLAN);
      stateManager.enqueue([call]);
      stateManager.dequeue();

      const response: ToolCallResponseInfo = {
        callId: call.request.callId,
        responseParts: [],
        resultDisplay: 'Success',
        error: undefined,
        errorType: undefined,
      };

      vi.mocked(onUpdate).mockClear();
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Success,
        response,
      );
      expect(onUpdate).toHaveBeenCalledTimes(1);

      vi.mocked(onUpdate).mockClear();
      stateManager.finalizeCall(call.request.callId);
      expect(onUpdate).toHaveBeenCalledTimes(1);

      expect(stateManager.isActive).toBe(false);
      expect(stateManager.completedBatch).toHaveLength(1);
      const completed = stateManager.completedBatch[0] as SuccessfulToolCall;
      expect(completed.status).toBe(CoreToolCallStatus.Success);
      expect(completed.response).toEqual(response);
      expect(completed.durationMs).toBeDefined();
      expect(completed.approvalMode).toBe(ApprovalMode.PLAN);
    });

    it('should transition to error and move to completed batch', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();

      const response: ToolCallResponseInfo = {
        callId: call.request.callId,
        responseParts: [],
        resultDisplay: 'Error',
        error: new Error('Failed'),
        errorType: undefined,
      };

      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Error,
        response,
      );
      stateManager.finalizeCall(call.request.callId);

      expect(stateManager.isActive).toBe(false);
      expect(stateManager.completedBatch).toHaveLength(1);
      const completed = stateManager.completedBatch[0] as ErroredToolCall;
      expect(completed.status).toBe(CoreToolCallStatus.Error);
      expect(completed.response).toEqual(response);
    });

    it('should transition to awaiting_approval with details', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();

      const details = {
        type: 'info' as const,
        title: 'Confirm',
        prompt: 'Proceed?',
        onConfirm: vi.fn(),
      };

      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.AwaitingApproval,
        details,
      );

      const active = stateManager.firstActiveCall as WaitingToolCall;
      expect(active.status).toBe(CoreToolCallStatus.AwaitingApproval);
      expect(active.confirmationDetails).toEqual(details);
    });

    it('should transition to awaiting_approval with event-driven format', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();

      const details = {
        type: 'info' as const,
        title: 'Confirm',
        prompt: 'Proceed?',
      };
      const eventDrivenData = {
        correlationId: 'corr-123',
        confirmationDetails: details,
      };

      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.AwaitingApproval,
        eventDrivenData,
      );

      const active = stateManager.firstActiveCall as WaitingToolCall;
      expect(active.status).toBe(CoreToolCallStatus.AwaitingApproval);
      expect(active.correlationId).toBe('corr-123');
      expect(active.confirmationDetails).toEqual(details);
    });

    it('should preserve diff and derive stats when cancelling an edit tool call', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();

      const details = {
        type: 'edit' as const,
        title: 'Edit',
        fileName: 'test.txt',
        filePath: '/path/to/test.txt',
        fileDiff: '@@ -1,1 +1,1 @@\n-old line\n+new line',
        originalContent: 'old line',
        newContent: 'new line',
        onConfirm: vi.fn(),
      };

      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.AwaitingApproval,
        details,
      );
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Cancelled,
        'User said no',
      );
      stateManager.finalizeCall(call.request.callId);

      const completed = stateManager.completedBatch[0] as CancelledToolCall;
      expect(completed.status).toBe(CoreToolCallStatus.Cancelled);
      const result = completed.response.resultDisplay as FileDiff;
      expect(result.fileDiff).toBe(details.fileDiff);
      expect(result.diffStat).toEqual(
        expect.objectContaining({
          model_added_lines: 1,
          model_removed_lines: 1,
        }),
      );
    });

    it('should ignore status updates for non-existent callIds', () => {
      stateManager.updateStatus('unknown', CoreToolCallStatus.Scheduled);
      expect(onUpdate).not.toHaveBeenCalled();
    });

    it('should ignore status updates for terminal calls', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Success,
        createMockResponse(call.request.callId),
      );
      stateManager.finalizeCall(call.request.callId);

      vi.mocked(onUpdate).mockClear();
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Scheduled,
      );
      expect(onUpdate).not.toHaveBeenCalled();
    });

    it('should only finalize terminal calls', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();

      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Executing,
      );
      stateManager.finalizeCall(call.request.callId);

      expect(stateManager.isActive).toBe(true);
      expect(stateManager.completedBatch).toHaveLength(0);

      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Success,
        createMockResponse(call.request.callId),
      );
      stateManager.finalizeCall(call.request.callId);

      expect(stateManager.isActive).toBe(false);
      expect(stateManager.completedBatch).toHaveLength(1);
    });

    it('should merge liveOutput and pid during executing updates', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();

      // Start executing
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Executing,
      );
      let active = stateManager.firstActiveCall as ExecutingToolCall;
      expect(active.status).toBe(CoreToolCallStatus.Executing);
      expect(active.liveOutput).toBeUndefined();

      // Update with live output
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Executing,
        {
          liveOutput: 'chunk 1',
        },
      );
      active = stateManager.firstActiveCall as ExecutingToolCall;
      expect(active.liveOutput).toBe('chunk 1');

      // Update with pid (should preserve liveOutput)
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Executing,
        {
          pid: 1234,
        },
      );
      active = stateManager.firstActiveCall as ExecutingToolCall;
      expect(active.liveOutput).toBe('chunk 1');
      expect(active.pid).toBe(1234);

      // Update live output again (should preserve pid)
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Executing,
        {
          liveOutput: 'chunk 2',
        },
      );
      active = stateManager.firstActiveCall as ExecutingToolCall;
      expect(active.liveOutput).toBe('chunk 2');
      expect(active.pid).toBe(1234);
    });

    it('should update progressMessage and progressPercent during executing updates', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();

      // Update with progress
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Executing,
        {
          progressMessage: 'Starting...',
          progressPercent: 10,
        },
      );
      let active = stateManager.firstActiveCall as ExecutingToolCall;
      expect(active.progressMessage).toBe('Starting...');
      expect(active.progressPercent).toBe(10);

      // Update progress further
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Executing,
        {
          progressMessage: 'Halfway!',
          progressPercent: 50,
        },
      );
      active = stateManager.firstActiveCall as ExecutingToolCall;
      expect(active.progressMessage).toBe('Halfway!');
      expect(active.progressPercent).toBe(50);
    });
  });

  describe('Argument Updates', () => {
    it('should update args and invocation', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();

      const newArgs = { foo: 'updated' };
      const newInvocation = { ...mockInvocation } as AnyToolInvocation;

      stateManager.updateArgs(call.request.callId, newArgs, newInvocation);

      const active = stateManager.firstActiveCall;
      if (active && 'invocation' in active) {
        expect(active.invocation).toEqual(newInvocation);
      } else {
        throw new Error('Active call should have invocation');
      }
    });

    it('should ignore arg updates for errored calls', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Error,
        createMockResponse(call.request.callId),
      );
      stateManager.finalizeCall(call.request.callId);

      stateManager.updateArgs(
        call.request.callId,
        { foo: 'new' },
        mockInvocation,
      );

      const completed = stateManager.completedBatch[0];
      expect(completed.request.args).toEqual(mockRequest.args);
    });
  });

  describe('Outcome Tracking', () => {
    it('should set outcome and notify', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();

      stateManager.setOutcome(
        call.request.callId,
        ToolConfirmationOutcome.ProceedAlways,
      );

      const active = stateManager.firstActiveCall;
      expect(active?.outcome).toBe(ToolConfirmationOutcome.ProceedAlways);
      expect(onUpdate).toHaveBeenCalled();
    });
  });

  describe('Batch Operations', () => {
    it('should cancel all queued calls', () => {
      stateManager.enqueue([
        createValidatingCall('1'),
        createValidatingCall('2'),
      ]);

      vi.mocked(onUpdate).mockClear();
      stateManager.cancelAllQueued('Batch cancel');

      expect(stateManager.queueLength).toBe(0);
      expect(stateManager.completedBatch).toHaveLength(2);
      expect(
        stateManager.completedBatch.every(
          (c) => c.status === CoreToolCallStatus.Cancelled,
        ),
      ).toBe(true);
      expect(onUpdate).toHaveBeenCalledTimes(1);
    });

    it('should use originalRequestName when cancelling queued calls', () => {
      const call = createValidatingCall('tail-1');
      call.request.originalRequestName = 'original-tool';
      stateManager.enqueue([call]);

      stateManager.cancelAllQueued('Batch cancel');

      const completed = stateManager.completedBatch[0] as CancelledToolCall;
      expect(completed.response.responseParts[0]?.functionResponse?.name).toBe(
        'original-tool',
      );
    });

    it('should not notify if cancelAllQueued is called on an empty queue', () => {
      vi.mocked(onUpdate).mockClear();
      stateManager.cancelAllQueued('Batch cancel');
      expect(onUpdate).not.toHaveBeenCalled();
    });

    it('should clear batch and notify', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Success,
        createMockResponse(call.request.callId),
      );
      stateManager.finalizeCall(call.request.callId);

      stateManager.clearBatch();

      expect(stateManager.completedBatch).toHaveLength(0);
      expect(onUpdate).toHaveBeenCalledWith([]);
    });

    it('should return a copy of the completed batch (defensive)', () => {
      const call = createValidatingCall();
      stateManager.enqueue([call]);
      stateManager.dequeue();
      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Success,
        createMockResponse(call.request.callId),
      );
      stateManager.finalizeCall(call.request.callId);

      const batch = stateManager.completedBatch;
      expect(batch).toHaveLength(1);

      // Mutate the returned array
      batch.pop();
      expect(batch).toHaveLength(0);

      // Verify internal state is unchanged
      expect(stateManager.completedBatch).toHaveLength(1);
    });
  });

  describe('Snapshot and Ordering', () => {
    it('should return snapshot in order: completed, active, queue', () => {
      // 1. Completed
      const call1 = createValidatingCall('1');
      stateManager.enqueue([call1]);
      stateManager.dequeue();
      stateManager.updateStatus(
        '1',
        CoreToolCallStatus.Success,
        createMockResponse('1'),
      );
      stateManager.finalizeCall('1');

      // 2. Active
      const call2 = createValidatingCall('2');
      stateManager.enqueue([call2]);
      stateManager.dequeue();

      // 3. Queue
      const call3 = createValidatingCall('3');
      stateManager.enqueue([call3]);

      const snapshot = stateManager.getSnapshot();
      expect(snapshot).toHaveLength(3);
      expect(snapshot[0].request.callId).toBe('1');
      expect(snapshot[1].request.callId).toBe('2');
      expect(snapshot[2].request.callId).toBe('3');
    });
  });

  describe('progress field preservation', () => {
    it('should preserve progress and progressTotal in toExecuting', () => {
      const call = createValidatingCall('progress-1');
      stateManager.enqueue([call]);
      stateManager.dequeue();

      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Executing,
        {
          progress: 5,
          progressTotal: 10,
          progressMessage: 'Working',
          progressPercent: 50,
        },
      );

      const active = stateManager.firstActiveCall as ExecutingToolCall;
      expect(active.status).toBe(CoreToolCallStatus.Executing);
      expect(active.progress).toBe(5);
      expect(active.progressTotal).toBe(10);
      expect(active.progressMessage).toBe('Working');
      expect(active.progressPercent).toBe(50);
    });

    it('should preserve progress fields after a liveOutput update', () => {
      const call = createValidatingCall('progress-2');
      stateManager.enqueue([call]);
      stateManager.dequeue();

      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Executing,
        {
          progress: 5,
          progressTotal: 10,
          progressMessage: 'Working',
          progressPercent: 50,
        },
      );

      stateManager.updateStatus(
        call.request.callId,
        CoreToolCallStatus.Executing,
        {
          liveOutput: 'some output',
        },
      );

      const active = stateManager.firstActiveCall as ExecutingToolCall;
      expect(active.status).toBe(CoreToolCallStatus.Executing);
      expect(active.liveOutput).toBe('some output');
      expect(active.progress).toBe(5);
      expect(active.progressTotal).toBe(10);
      expect(active.progressMessage).toBe('Working');
      expect(active.progressPercent).toBe(50);
    });
  });
});