// @vitest-environment node

import { describe, expect, it } from 'vitest';
import {
  PiSubagentChildError,
  PiSubagentScheduler,
  parsePiSubagentDispatchRequest,
  type PiSubagentChild,
  type PiSubagentChildOpenInput,
} from '../../electron/coding-runtime/pi/subagent';
import { PiProcessBudget, PiWorkerPool } from '../../electron/coding-runtime/pi/worker-pool';

function parent(runId = 'run-a') {
  return {
    conversationId: 'conversation-a',
    workerGeneration: 1,
    runId,
    projectId: 'project-a',
  };
}

function task(agentId: string, toolProfile: 'read-only' | 'coding' = 'read-only') {
  return { agentId, task: `Inspect ${agentId}`, toolProfile };
}

function deferred() {
  let resolve!: () => void;
  const promise = new Promise<void>((done) => { resolve = done; });
  return { promise, resolve };
}

describe('Pi subagent scheduler', () => {
  it('validates one bounded dispatch mode and rejects a ninth task', () => {
    expect(parsePiSubagentDispatchRequest({ mode: 'single', tasks: [task('one')] })).toEqual({
      mode: 'single', tasks: [task('one')],
    });
    expect(parsePiSubagentDispatchRequest({
      mode: 'parallel', tasks: Array.from({ length: 8 }, (_, index) => task(`agent-${index}`)),
    }).tasks).toHaveLength(8);
    expect(() => parsePiSubagentDispatchRequest({
      mode: 'parallel', tasks: Array.from({ length: 9 }, (_, index) => task(`agent-${index}`)),
    })).toThrowError('Subagent dispatch accepts at most 8 tasks');
    expect(() => parsePiSubagentDispatchRequest({
      mode: 'single', tasks: [task('one'), task('two')],
    })).toThrowError('Single subagent dispatch requires exactly one task');
    expect(() => parsePiSubagentDispatchRequest({ mode: 'unknown', tasks: [task('one')] }))
      .toThrowError('Subagent dispatch mode is invalid');
  });

  it('shares one FIFO four-child cap across two parents and releases the process budget', async () => {
    const gate = deferred();
    const processBudget = new PiProcessBudget(8);
    let running = 0;
    let maxRunning = 0;
    const opened: PiSubagentChildOpenInput[] = [];
    const scheduler = new PiSubagentScheduler({
      processBudget,
      createId: (() => {
        let id = 0;
        return (kind) => `${kind}-${++id}`;
      })(),
      openChild: async (input) => {
        opened.push(structuredClone(input));
        return {
          id: input.taskId,
          async run() {
            running += 1;
            maxRunning = Math.max(maxRunning, running);
            await gate.promise;
            running -= 1;
            return {
              summary: `done ${input.agentId}`,
              usage: { inputTokens: 1, outputTokens: 2 },
            };
          },
          async stop() {},
        } satisfies PiSubagentChild;
      },
    });
    const request = {
      mode: 'parallel' as const,
      tasks: Array.from({ length: 4 }, (_, index) => task(`agent-${index}`)),
    };
    const left = scheduler.dispatch({ ...parent('run-left'), request });
    const right = scheduler.dispatch({ ...parent('run-right'), request });
    await expect.poll(() => running).toBe(4);
    expect(maxRunning).toBe(4);
    expect(processBudget.activeCount).toBe(4);
    expect(scheduler.getDiagnostics()).toEqual({
      activeChildPermits: 4,
      waitingChildPermits: 4,
      activeDispatches: 2,
      activeParents: 2,
    });
    gate.resolve();
    const [leftResult, rightResult] = await Promise.all([left, right]);
    expect(leftResult.details.tasks).toHaveLength(4);
    expect(rightResult.details.tasks).toHaveLength(4);
    expect(opened).toHaveLength(8);
    expect(maxRunning).toBe(4);
    expect(processBudget.activeCount).toBe(0);
    expect(scheduler.getDiagnostics()).toEqual({
      activeChildPermits: 0,
      waitingChildPermits: 0,
      activeDispatches: 0,
      activeParents: 0,
    });
    await scheduler.close();
  });

  it('preserves parallel sibling results and stops a chain at its first failure', async () => {
    const scheduler = new PiSubagentScheduler({
      processBudget: new PiProcessBudget(8),
      createId: (() => {
        let id = 0;
        return (kind) => `${kind}-${++id}`;
      })(),
      openChild: async (input) => ({
        id: input.taskId,
        async run(prompt) {
          if (input.agentId === 'broken') {
            throw new PiSubagentChildError('SUBAGENT_CHILD_FAILED');
          }
          return {
            summary: `${input.agentId}:${prompt}`,
            usage: { inputTokens: 3, outputTokens: 5 },
          };
        },
        async stop() {},
      }),
    });
    const parallel = await scheduler.dispatch({
      ...parent('parallel'),
      request: {
        mode: 'parallel',
        tasks: [task('left'), task('broken'), task('right')],
      },
    });
    expect(parallel.details.tasks).toEqual(expect.arrayContaining([
      expect.objectContaining({ agentId: 'left', status: 'complete', summary: 'left:Inspect left' }),
      expect.objectContaining({ agentId: 'broken', status: 'error', errorCode: 'SUBAGENT_CHILD_FAILED' }),
      expect.objectContaining({ agentId: 'right', status: 'complete', usage: { inputTokens: 3, outputTokens: 5 } }),
    ]));

    const chain = await scheduler.dispatch({
      ...parent('chain'),
      request: {
        mode: 'chain',
        tasks: [
          { agentId: 'first', task: 'first', toolProfile: 'read-only' },
          { agentId: 'broken', task: 'review {previous}', toolProfile: 'read-only' },
          { agentId: 'never', task: 'never', toolProfile: 'coding' },
        ],
      },
    });
    expect(chain.details.tasks.map(({ status }) => status)).toEqual(['complete', 'error', 'skipped']);
    expect(chain.details.tasks[1]).toMatchObject({ errorCode: 'SUBAGENT_CHILD_FAILED' });
    await scheduler.close();
  });

  it('aborts every unfinished child with the parent and leaves no permit or process lease', async () => {
    const processBudget = new PiProcessBudget(8);
    const stopped: Array<{ taskId: string; reason: string }> = [];
    const scheduler = new PiSubagentScheduler({
      processBudget,
      openChild: async (input) => ({
        id: input.taskId,
        async run(_prompt, signal) {
          await new Promise<void>((_resolve, reject) => {
            const abort = () => reject(new PiSubagentChildError('SUBAGENT_ABORTED'));
            if (signal.aborted) abort();
            else signal.addEventListener('abort', abort, { once: true });
          });
          return { summary: 'unreachable' };
        },
        async stop(reason) { stopped.push({ taskId: input.taskId, reason }); },
      }),
    });
    const controller = new AbortController();
    const flight = scheduler.dispatch({
      ...parent('abort'),
      request: { mode: 'parallel', tasks: [task('left'), task('right', 'coding')] },
    }, { signal: controller.signal });
    await expect.poll(() => processBudget.activeCount).toBe(2);
    controller.abort();
    const result = await flight;
    expect(result.details.tasks.map(({ status }) => status)).toEqual(['aborted', 'aborted']);
    expect(stopped).toHaveLength(2);
    expect(stopped.map(({ reason }) => reason)).toEqual(['subagent_abort', 'subagent_abort']);
    expect(processBudget.activeCount).toBe(0);
    await scheduler.close();
  });

  it('marks the rest of an in-flight chain aborted when its parent aborts', async () => {
    const processBudget = new PiProcessBudget(8);
    const scheduler = new PiSubagentScheduler({
      processBudget,
      openChild: async (input) => ({
        id: input.taskId,
        async run(_prompt, signal) {
          await new Promise<void>((_resolve, reject) => {
            const abort = () => reject(new PiSubagentChildError('SUBAGENT_ABORTED'));
            if (signal.aborted) abort();
            else signal.addEventListener('abort', abort, { once: true });
          });
          return { summary: 'unreachable' };
        },
        async stop() {},
      }),
    });
    const controller = new AbortController();
    const flight = scheduler.dispatch({
      ...parent('abort-chain'),
      request: { mode: 'chain', tasks: [task('one'), task('two'), task('three')] },
    }, { signal: controller.signal });
    await expect.poll(() => processBudget.activeCount).toBe(1);
    controller.abort();
    const result = await flight;
    expect(result.details.tasks.map(({ status }) => status))
      .toEqual(['aborted', 'aborted', 'aborted']);
    expect(processBudget.activeCount).toBe(0);
    await scheduler.close();
  });

  it('reclaims an idle parent process instead of deadlocking on a full global budget', async () => {
    const processBudget = new PiProcessBudget(2);
    const firstParent = await processBudget.acquire();
    const idleParent = await processBudget.acquire();
    let waitingParentSettled = false;
    const waitingParent = processBudget.acquire().then((lease) => {
      waitingParentSettled = true;
      return lease;
    });
    let reclaimed = 0;
    let childRanBeforeWaitingParent = false;
    const scheduler = new PiSubagentScheduler({
      processBudget,
      reclaimProcessCapacity: async () => {
        reclaimed += 1;
        idleParent.release();
        return true;
      },
      openChild: async (input) => ({
        id: input.taskId,
        async run() {
          childRanBeforeWaitingParent = !waitingParentSettled;
          return { summary: 'done' };
        },
        async stop() {},
      }),
    });
    await expect(scheduler.dispatch({
      ...parent('full-budget'),
      request: { mode: 'single', tasks: [task('agent-a')] },
    })).resolves.toMatchObject({
      details: { tasks: [{ status: 'complete' }] },
    });
    expect(reclaimed).toBe(1);
    expect(childRanBeforeWaitingParent).toBe(true);
    const waitingParentLease = await waitingParent;
    expect(processBudget.activeCount).toBe(2);
    waitingParentLease.release();
    firstParent.release();
    expect(processBudget.activeCount).toBe(0);
    await scheduler.close();
  });

  it('reclaims a parent that becomes idle after child capacity is already reserved', async () => {
    const parentOpenGate = deferred();
    const processBudget = new PiProcessBudget(1);
    let parentStopped = false;
    let childRan = false;
    const pool = new PiWorkerPool({
      processBudget,
      maxIdle: 1,
      openWorker: async () => {
        await parentOpenGate.promise;
        return {
          worker: {
            id: 'delayed-parent',
            generation: 1,
            async request() {
              return { type: 'response' as const, id: 'parent', success: true as const };
            },
            async send() {},
            subscribe() { return () => undefined; },
            subscribeInvalidation() { return () => undefined; },
            async stop() {
              parentStopped = true;
              return { mode: 'stdin-close' as const, code: 0, signal: null };
            },
          },
          session: { piSessionId: 'parent-session', sessionKey: 'parent-key' },
        };
      },
    });
    const preparingParent = pool.prepare({
      conversationId: 'delayed-parent',
      projectId: 'project-a',
      agentId: 'agent-a',
      title: 'Delayed parent',
      model: {
        model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
        modelResolution: 'resolved',
      },
    });
    await expect.poll(() => processBudget.activeCount).toBe(1);

    const scheduler = new PiSubagentScheduler({
      processBudget,
      reclaimProcessCapacity: (signal) => pool.reclaimIdleWorker(signal),
      openChild: async (input) => ({
        id: input.taskId,
        async run() {
          childRan = true;
          return { summary: 'done' };
        },
        async stop() {},
      }),
    });
    const child = scheduler.dispatch({
      ...parent('delayed-idle'),
      request: { mode: 'single', tasks: [task('agent-a')] },
    });
    await expect.poll(() => processBudget.waitingCount).toBe(1);

    parentOpenGate.resolve();
    await preparingParent;
    await expect(child).resolves.toMatchObject({
      details: { tasks: [{ status: 'complete' }] },
    });
    expect(parentStopped).toBe(true);
    expect(childRan).toBe(true);
    expect(processBudget.activeCount).toBe(0);
    expect(processBudget.waitingCount).toBe(0);
    await scheduler.close();
    await pool.shutdown();
  });

  it('cancels a queued child reservation when the capacity reclaimer rejects', async () => {
    const processBudget = new PiProcessBudget(1);
    const parentLease = await processBudget.acquire();
    const scheduler = new PiSubagentScheduler({
      processBudget,
      reclaimProcessCapacity: async () => {
        throw new Error('reclaim failed');
      },
      openChild: async () => {
        throw new Error('child must not open');
      },
    });

    await expect(scheduler.dispatch({
      ...parent('reclaim-rejected'),
      request: { mode: 'single', tasks: [task('agent-a')] },
    })).resolves.toMatchObject({
      details: { tasks: [{ status: 'error', errorCode: 'SUBAGENT_CHILD_FAILED' }] },
    });
    expect(processBudget.activeCount).toBe(1);
    expect(processBudget.waitingCount).toBe(0);
    parentLease.release();
    expect(processBudget.activeCount).toBe(0);
    expect(processBudget.waitingCount).toBe(0);
    await scheduler.close();
  });

  it('cancels both capacity waits when the parent aborts before an idle worker exists', async () => {
    const processBudget = new PiProcessBudget(1);
    const parentLease = await processBudget.acquire();
    let reclaimerCancelled = false;
    const scheduler = new PiSubagentScheduler({
      processBudget,
      reclaimProcessCapacity: async (signal) => await new Promise<boolean>((_resolve, reject) => {
        const cancel = () => {
          reclaimerCancelled = true;
          reject(new Error('reclaim cancelled'));
        };
        if (signal?.aborted) cancel();
        else signal?.addEventListener('abort', cancel, { once: true });
      }),
      openChild: async () => {
        throw new Error('child must not open');
      },
    });
    const controller = new AbortController();
    const child = scheduler.dispatch({
      ...parent('reclaim-aborted'),
      request: { mode: 'single', tasks: [task('agent-a')] },
    }, { signal: controller.signal });
    await expect.poll(() => processBudget.waitingCount).toBe(1);

    controller.abort();
    await expect(child).resolves.toMatchObject({
      details: { tasks: [{ status: 'aborted', errorCode: 'SUBAGENT_ABORTED' }] },
    });
    expect(reclaimerCancelled).toBe(true);
    expect(processBudget.activeCount).toBe(1);
    expect(processBudget.waitingCount).toBe(0);
    parentLease.release();
    expect(processBudget.activeCount).toBe(0);
    await scheduler.close();
  });
});