// @vitest-environment node

import { mkdtemp, readFile, rm } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import path from 'node:path';
import { afterEach, describe, expect, it } from 'vitest';
import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
import { PiProjectWriteLeaseCoordinator } from '../../electron/coding-runtime/pi/write-lease';
import {
  PiSubagentChildError,
  PiSubagentScheduler,
} from '../../electron/coding-runtime/pi/subagent';
import { PiProcessBudget } from '../../electron/coding-runtime/pi/worker-pool';

const roots: string[] = [];
const hosts: PiManagedExtensionHost[] = [];

afterEach(async () => {
  await Promise.all(hosts.splice(0).map((host) => host.close()));
  await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true })));
});

async function post(
  registration: Awaited<ReturnType<PiManagedExtensionHost['registerWorker']>>,
  body: Record<string, unknown>,
): Promise<Response> {
  return await fetch(registration.env.MAKELORE_PI_BRIDGE_URL as string, {
    method: 'POST',
    headers: {
      authorization: `Bearer ${registration.env.MAKELORE_PI_WORKER_TOKEN}`,
      'content-type': 'application/json',
    },
    body: JSON.stringify(body),
  });
}

describe('managed Pi extension bridge', () => {
  it('single-flights the shared managed extension for concurrent worker registration', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-concurrent-'));
    roots.push(root);
    const host = new PiManagedExtensionHost();
    hosts.push(host);

    const registrations = await Promise.all(Array.from({ length: 8 }, async (_, index) => (
      await host.registerWorker({
        conversationId: `conversation-${index + 1}`,
        generation: 1,
        projectId: `project-${index + 1}`,
        extensionsDir: root,
      })
    )));

    expect(new Set(registrations.map(({ extensionPath }) => extensionPath)).size).toBe(1);
    expect(host.getDiagnostics().registrations).toEqual({ parent: 8, child: 0 });
    await Promise.all(registrations.map(({ dispose }) => dispose()));
  });

  it('finishes a queued HTTP request while closing a holder and waiter', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-close-'));
    roots.push(root);
    const leases = new PiProjectWriteLeaseCoordinator();
    const host = new PiManagedExtensionHost(leases);
    hosts.push(host);
    const holder = await host.registerWorker({
      conversationId: 'conversation-holder', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    const waiter = await host.registerWorker({
      conversationId: 'conversation-waiter', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    await Promise.all([
      host.bindRun('conversation-holder', 1, 'run-holder'),
      host.bindRun('conversation-waiter', 1, 'run-waiter'),
    ]);
    const held = await post(holder, {
      action: 'lease.acquire', conversationId: 'conversation-holder', workerGeneration: 1,
      runId: 'run-holder', resourceId: 'held-tool',
    });
    expect(held.status).toBe(200);
    const waiting = post(waiter, {
      action: 'lease.acquire', conversationId: 'conversation-waiter', workerGeneration: 1,
      runId: 'run-waiter', resourceId: 'waiting-tool',
    });
    await expect.poll(() => leases.waitingCount('project-a')).toBe(1);

    await expect(Promise.race([
      host.close().then(() => 'closed'),
      new Promise<string>((resolve) => setTimeout(() => resolve('timeout'), 500)),
    ])).resolves.toBe('closed');
    await expect(waiting).resolves.toMatchObject({ status: 409 });
  });

  it('materializes the active run into a replacement generation before spawn', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-rebuild-'));
    roots.push(root);
    const host = new PiManagedExtensionHost();
    hosts.push(host);
    const first = await host.registerWorker({
      conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    await host.bindRun('conversation-a', 1, 'run-a');
    const replacement = await host.registerWorker({
      conversationId: 'conversation-a', generation: 2, projectId: 'project-a', extensionsDir: root,
    });
    const context = JSON.parse(await readFile(
      replacement.env.MAKELORE_PI_CONTEXT_FILE as string,
      'utf8',
    )) as Record<string, unknown>;
    expect(context).toEqual({
      conversationId: 'conversation-a', workerGeneration: 2, role: 'parent', runId: 'run-a',
    });
    await first.dispose();
    const response = await post(replacement, {
      action: 'lease.acquire', conversationId: 'conversation-a', workerGeneration: 2,
      runId: 'run-a', resourceId: 'replacement-tool',
    });
    expect(response.status).toBe(200);
  });

  it('streams stable subagent details and rejects recursive child dispatch', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-subagent-'));
    roots.push(root);
    const scheduler = new PiSubagentScheduler({
      processBudget: new PiProcessBudget(8),
      openChild: async (input) => ({
        id: input.taskId,
        async run() { return { summary: `done ${input.agentId}` }; },
        async stop() {},
      }),
    });
    const host = new PiManagedExtensionHost();
    const tracked: Array<{ kind: string; id: string }> = [];
    let untracked = 0;
    host.configureSubagents({
      scheduler,
      trackGenerationResource: (input) => {
        tracked.push({ kind: input.kind, id: input.id });
        return () => { untracked += 1; };
      },
    });
    hosts.push(host);
    const parent = await host.registerWorker({
      conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    const child = await host.registerWorker({
      conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
      role: 'child', runId: 'run-a',
    });
    await host.bindRun('conversation-a', 1, 'run-a');
    const request = {
      action: 'subagent.dispatch', conversationId: 'conversation-a', workerGeneration: 1,
      runId: 'run-a', resourceId: 'subagent-tool',
      request: {
        mode: 'single',
        tasks: [{ agentId: 'agent-a', task: 'Inspect', toolProfile: 'read-only' }],
      },
    };
    const response = await post(parent, request);
    expect(response.status).toBe(200);
    const lines = (await response.text()).trim().split('\n').map((line) => JSON.parse(line));
    expect(lines.at(-1)).toMatchObject({
      done: true,
      details: {
        schema: 'subagent.v1', mode: 'single',
        tasks: [{ agentId: 'agent-a', status: 'complete', summary: 'done agent-a' }],
      },
    });
    expect(tracked).toEqual([{ kind: 'child', id: 'subagent-tool' }]);
    expect(untracked).toBe(1);
    expect((await post(child, request)).status).toBe(403);
    await scheduler.close();
  });

  it('propagates generation cancellation without an orphan child or leaked budget', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-child-cancel-'));
    roots.push(root);
    const processBudget = new PiProcessBudget(8);
    let stopped = 0;
    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() { stopped += 1; },
      }),
    });
    let cancelGeneration: (() => void) | undefined;
    let untracked = 0;
    const host = new PiManagedExtensionHost();
    host.configureSubagents({
      scheduler,
      trackGenerationResource: (input) => {
        cancelGeneration = input.cancel;
        return () => { untracked += 1; };
      },
    });
    hosts.push(host);
    const parent = await host.registerWorker({
      conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    await host.bindRun('conversation-a', 1, 'run-a');
    const flight = post(parent, {
      action: 'subagent.dispatch', conversationId: 'conversation-a', workerGeneration: 1,
      runId: 'run-a', resourceId: 'subagent-tool',
      request: {
        mode: 'single',
        tasks: [{ agentId: 'agent-a', task: 'Wait', toolProfile: 'read-only' }],
      },
    });
    await expect.poll(() => processBudget.activeCount).toBe(1);
    cancelGeneration?.();
    const response = await flight;
    const lines = (await response.text()).trim().split('\n').map((line) => JSON.parse(line));
    expect(lines.at(-1)).toMatchObject({
      done: true,
      details: { tasks: [{ status: 'aborted', errorCode: 'SUBAGENT_ABORTED' }] },
    });
    expect(stopped).toBe(1);
    expect(processBudget.activeCount).toBe(0);
    expect(untracked).toBe(1);
    await scheduler.close();
  });

  it('validates worker identity and enforces project-scoped leases over loopback HTTP', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-'));
    roots.push(root);
    const host = new PiManagedExtensionHost();
    hosts.push(host);
    const first = await host.registerWorker({
      conversationId: 'conversation-a1', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    const second = await host.registerWorker({
      conversationId: 'conversation-a2', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    const other = await host.registerWorker({
      conversationId: 'conversation-b1', generation: 1, projectId: 'project-b', extensionsDir: root,
    });
    await Promise.all([
      host.bindRun('conversation-a1', 1, 'run-a1'),
      host.bindRun('conversation-a2', 1, 'run-a2'),
      host.bindRun('conversation-b1', 1, 'run-b1'),
    ]);
    const identity = (conversationId: string, runId: string, resourceId: string) => ({
      action: 'lease.acquire', conversationId, workerGeneration: 1, runId, resourceId,
    });

    const firstResponse = await post(first, identity('conversation-a1', 'run-a1', 'tool-a1'));
    expect(firstResponse.status).toBe(200);
    const firstLease = await firstResponse.json() as { leaseId: string };
    let sameProjectSettled = false;
    const sameProjectFlight = post(second, identity('conversation-a2', 'run-a2', 'tool-a2'))
      .then((response) => {
        sameProjectSettled = true;
        return response;
      });
    const otherResponse = await post(other, identity('conversation-b1', 'run-b1', 'tool-b1'));
    expect(otherResponse.status).toBe(200);
    await Promise.resolve();
    expect(sameProjectSettled).toBe(false);

    const releaseResponse = await post(first, {
      action: 'lease.release',
      conversationId: 'conversation-a1',
      workerGeneration: 1,
      runId: 'run-a1',
      resourceId: 'tool-a1',
      leaseId: firstLease.leaseId,
    });
    expect(releaseResponse.status).toBe(200);
    expect((await sameProjectFlight).status).toBe(200);

    const forged = await post(second, identity('conversation-a2', 'old-run', 'forged'));
    expect(forged.status).toBe(409);
    await first.dispose();
    const staleToken = await post(first, identity('conversation-a1', 'run-a1', 'stale'));
    expect(staleToken.status).toBe(401);
    const currentWorker = await post(other, {
      action: 'lease.release',
      conversationId: 'conversation-b1', workerGeneration: 1, runId: 'run-b1',
      resourceId: 'tool-b1',
      leaseId: (await otherResponse.clone().json() as { leaseId: string }).leaseId,
    });
    expect(currentWorker.status).toBe(200);
  });

  it('releases a crashed worker write lease so only its same-project waiter advances', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-crash-lease-'));
    roots.push(root);
    const leases = new PiProjectWriteLeaseCoordinator();
    const host = new PiManagedExtensionHost(leases);
    hosts.push(host);
    const crashed = await host.registerWorker({
      conversationId: 'conversation-crashed', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    const waiter = await host.registerWorker({
      conversationId: 'conversation-waiter', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    const other = await host.registerWorker({
      conversationId: 'conversation-other', generation: 1, projectId: 'project-b', extensionsDir: root,
    });
    await Promise.all([
      host.bindRun('conversation-crashed', 1, 'run-crashed'),
      host.bindRun('conversation-waiter', 1, 'run-waiter'),
      host.bindRun('conversation-other', 1, 'run-other'),
    ]);
    expect((await post(crashed, {
      action: 'lease.acquire', conversationId: 'conversation-crashed', workerGeneration: 1,
      runId: 'run-crashed', resourceId: 'crashed-write',
    })).status).toBe(200);
    const waiting = post(waiter, {
      action: 'lease.acquire', conversationId: 'conversation-waiter', workerGeneration: 1,
      runId: 'run-waiter', resourceId: 'waiting-write',
    });
    expect((await post(other, {
      action: 'lease.acquire', conversationId: 'conversation-other', workerGeneration: 1,
      runId: 'run-other', resourceId: 'other-write',
    })).status).toBe(200);
    await expect.poll(() => leases.waitingCount('project-a')).toBe(1);

    await crashed.dispose();

    await expect(waiting).resolves.toMatchObject({ status: 200 });
    expect(leases.waitingCount()).toBe(0);
    expect(leases.activeCount).toBe(2);
    await Promise.all([waiter.dispose(), other.dispose()]);
    expect(leases.activeCount).toBe(0);
  });

  it('joins a coding child to the same project write lease as its parent', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-child-lease-'));
    roots.push(root);
    const leases = new PiProjectWriteLeaseCoordinator();
    const host = new PiManagedExtensionHost(leases);
    hosts.push(host);
    const parent = await host.registerWorker({
      conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    const child = await host.registerWorker({
      conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
      role: 'child', runId: 'run-a',
    });
    await host.bindRun('conversation-a', 1, 'run-a');
    const identity = { conversationId: 'conversation-a', workerGeneration: 1, runId: 'run-a' };
    const held = await post(parent, {
      ...identity, action: 'lease.acquire', resourceId: 'parent-write',
    });
    const parentLease = await held.json() as { leaseId: string };
    let childSettled = false;
    const waiting = post(child, {
      ...identity, action: 'lease.acquire', resourceId: 'child-write',
    }).then((response) => {
      childSettled = true;
      return response;
    });
    await expect.poll(() => leases.waitingCount('project-a')).toBe(1);
    expect(childSettled).toBe(false);
    expect((await post(parent, {
      ...identity, action: 'lease.release', resourceId: 'parent-write', leaseId: parentLease.leaseId,
    })).status).toBe(200);
    expect((await waiting).status).toBe(200);
  });
});