import { afterEach, describe, expect, it } from 'vitest';
import Fastify, { type FastifyInstance } from 'fastify';
import fs from 'node:fs';
import os from 'node:os';
import path from 'node:path';
import { MindDB } from '@waggle/core';
import type { AgentLoopConfig, AgentResponse, ToolDefinition } from '@waggle/agent';
import { AgentRunRegistry } from '../../src/local/agent-run-registry.js';
import { bindChatCollaborationTools } from '../../src/local/chat-collaboration.js';
import { SignalBus } from '../../src/local/signal-bus.js';

const resources: Array<{
  dir: string;
  personal: MindDB;
  workspace: MindDB;
  registry: AgentRunRegistry;
  server: FastifyInstance;
}> = [];

const collaborationNames = [
  'spawn_agent', 'list_agents', 'get_agent_result',
  'compose_workflow', 'orchestrate_workflow', 'list_harnesses', 'run_harness',
];

function tool(name: string): ToolDefinition {
  return { name, description: name, parameters: { type: 'object', properties: {} }, execute: async () => name };
}

function deferred<T>() {
  let resolve!: (value: T) => void;
  let reject!: (error: Error) => void;
  const promise = new Promise<T>((done, fail) => { resolve = done; reject = fail; });
  return { promise, resolve, reject };
}

async function waitFor(predicate: () => boolean, message: string): Promise<void> {
  for (let attempt = 0; attempt < 100; attempt++) {
    if (predicate()) return;
    await new Promise((resolve) => setTimeout(resolve, 5));
  }
  throw new Error(message);
}

function setup() {
  const dir = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-chat-collaboration-'));
  const personal = new MindDB(path.join(dir, 'personal.mind'));
  const workspace = new MindDB(path.join(dir, 'workspace.mind'));
  const registry = new AgentRunRegistry(path.join(dir, 'agent-runs.json'));
  const signalBus = new SignalBus();
  const server = Fastify({ logger: false });
  server.decorate('localConfig', {
    dataDir: dir, port: 0, host: '127.0.0.1', litellmUrl: 'http://llm.test',
  });
  server.decorate('agentRunRegistry', registry);
  server.decorate('signalBus', signalBus);
  server.decorate('multiMind', { personal } as never);
  server.decorate('mindCache', {
    acquire: (workspaceId: string) => {
      if (workspaceId !== 'workspace-a') throw new Error(`wrong workspace: ${workspaceId}`);
      return workspace;
    },
    release: () => undefined,
  } as never);
  server.decorate('agentState', {
    litellmApiKey: 'test-key',
    skills: [],
    // Deliberately wrong. The request-bound bridge must never consult it.
    activeWorkspaceId: 'workspace-wrong',
  } as never);
  resources.push({ dir, personal, workspace, registry, server });
  return { dir, personal, workspace, registry, signalBus, server };
}

function bind(
  server: FastifyInstance,
  runLoop: (config: AgentLoopConfig) => Promise<AgentResponse>,
  sessionId = 'chat-session-a',
) {
  const visibleTools = [...collaborationNames.map(tool), tool('read_file'), tool('bash')];
  return bindChatCollaborationTools({
    server,
    visibleTools,
    workerTools: visibleTools,
    workspaceId: 'workspace-a',
    parentSessionId: sessionId,
    parentTask: 'Coordinate specialists on the release',
    model: 'model-default',
    runLoop,
    securityContext: {
      blockedTools: ['bash'],
      allowedToolNames: new Set(visibleTools.map((item) => item.name)),
    },
  });
}

afterEach(async () => {
  for (const resource of resources.splice(0)) {
    try { await resource.server.close(); } catch { /* not listening */ }
    resource.registry.close();
    resource.workspace.close();
    resource.personal.close();
    fs.rmSync(resource.dir, { recursive: true, force: true });
  }
});

describe('request-bound chat collaboration', () => {
  it('creates a durable scoped Room, Dance chain, dual memory, and registry-backed list/get', async () => {
    const { dir, registry, signalBus, server } = setup();
    const runnerCalls: AgentLoopConfig[] = [];
    const tools = bind(server, async (config) => {
      runnerCalls.push(config);
      config.onToolUse?.('read_file', { path: 'README.md' });
      return {
        content: 'Scoped specialist result',
        toolsUsed: ['read_file'],
        usage: { inputTokens: 11, outputTokens: 7 },
      };
    });

    const spawn = tools.find((item) => item.name === 'spawn_agent')!;
    const output = await spawn.execute({
      name: 'Release researcher', role: 'researcher', task: 'Inspect the release evidence',
    });
    const worker = registry.list({ source: 'chat_subagent', workspaceId: 'workspace-a' })
      .find((run) => run.kind === 'worker')!;
    const room = registry.get(worker.roomId)!;

    expect(output).toContain(`**Run ID:** ${worker.id}`);
    expect(worker).toMatchObject({
      kind: 'worker', workspaceId: 'workspace-a', status: 'completed',
      result: { summary: 'Scoped specialist result', sessionId: 'chat-session-a' },
      metrics: { toolsUsed: ['read_file'], inputTokens: 11, outputTokens: 7 },
      memoryRefs: { status: 'complete' },
    });
    expect(room).toMatchObject({ kind: 'room', workspaceIds: ['workspace-a'], status: 'completed' });
    expect(runnerCalls).toHaveLength(1);
    expect(runnerCalls[0].tools.map((item) => item.name)).toContain('read_file');
    expect(runnerCalls[0].tools.map((item) => item.name)).not.toContain('spawn_agent');
    expect(runnerCalls[0].tools.map((item) => item.name)).not.toContain('bash');

    const dance = signalBus.query({ teamId: `room::${worker.roomId}` }).reverse();
    expect(dance.map((item) => item.subtype)).toEqual([
      'task_delegation', 'task_claim', 'discovery', 'routed_share',
    ]);
    expect(dance[1].referenceId).toBe(dance[0].id);
    expect(dance[3].referenceId).toBe(dance[0].id);
    expect(dance.every((item) => item.content.workspaceId === 'workspace-a')).toBe(true);

    const secondOutput = await spawn.execute({
      name: 'Release writer', role: 'writer', task: 'Draft the release summary',
    });
    const sessionWorkers = registry.list({ source: 'chat_subagent', workspaceId: 'workspace-a' })
      .filter((run) => run.kind === 'worker');
    expect(secondOutput).toContain('Scoped specialist result');
    expect(sessionWorkers).toHaveLength(2);
    expect(new Set(sessionWorkers.map((run) => run.roomId)).size).toBe(2);

    const list = await tools.find((item) => item.name === 'list_agents')!.execute({});
    const get = await tools.find((item) => item.name === 'get_agent_result')!.execute({ agent_id: worker.id });
    expect(list).toContain(worker.id);
    expect(get).toContain('Scoped specialist result');

    const otherSessionTools = bind(server, async () => ({
      content: '', toolsUsed: [], usage: { inputTokens: 0, outputTokens: 0 },
    }), 'chat-session-b');
    expect(await otherSessionTools.find((item) => item.name === 'list_agents')!.execute({}))
      .toContain('No sub-agents spawned yet');

    const restored = new AgentRunRegistry(path.join(dir, 'agent-runs.json'));
    expect(restored.get(worker.id)?.status).toBe('completed');
    restored.close();
  });

  it('cancels one delegate truthfully and leaves a separate request runnable', async () => {
    const { registry, signalBus, server } = setup();
    const first = deferred<AgentResponse>();
    const firstTools = bind(server, (config) => {
      config.signal?.addEventListener('abort', () => first.reject(new Error('aborted')), { once: true });
      return first.promise;
    }, 'cancel-session');
    const running = firstTools.find((item) => item.name === 'spawn_agent')!.execute({
      name: 'Cancelable', role: 'researcher', task: 'Wait for cancellation',
    });
    await waitFor(
      () => registry.list({ source: 'chat_subagent', workspaceId: 'workspace-a' })
        .some((run) => run.kind === 'worker' && run.status === 'running'),
      'delegate did not start',
    );
    const worker = registry.list({ source: 'chat_subagent', workspaceId: 'workspace-a' })
      .find((run) => run.kind === 'worker')!;
    await registry.control(worker.id, 'cancel');
    expect(await running).toContain('Sub-Agent Error');
    expect(registry.get(worker.id)?.status).toBe('cancelled');
    expect(registry.get(worker.roomId)?.status).toBe('cancelled');
    expect(signalBus.query({ teamId: `room::${worker.roomId}` })[0]).toMatchObject({
      subtype: 'routed_share', content: { phase: 'cancelled' },
    });

    const secondTools = bind(server, async () => ({
      content: 'still works', toolsUsed: [], usage: { inputTokens: 1, outputTokens: 1 },
    }), 'independent-session');
    expect(await secondTools.find((item) => item.name === 'spawn_agent')!.execute({
      name: 'Independent', role: 'writer', task: 'Finish independently',
    })).toContain('still works');
  });

  it('exposes workflow cancellation only at the shared Room boundary', async () => {
    const { registry, signalBus, server } = setup();
    const tools = bind(server, (config) => new Promise<AgentResponse>((_resolve, reject) => {
      config.signal?.addEventListener('abort', () => reject(new Error('workflow aborted')), { once: true });
    }), 'workflow-cancel-session');
    const pending = tools.find((item) => item.name === 'orchestrate_workflow')!.execute({
      task: 'Cancelable workflow',
      inline_template: {
        name: 'Cancelable room', description: 'One shared cancellation boundary', aggregation: 'concatenate',
        steps: [{ name: 'Worker', role: 'researcher', task: 'Wait' }],
      },
    });
    await waitFor(
      () => registry.list({ source: 'workflow', workspaceId: 'workspace-a' })
        .some((run) => run.kind === 'worker' && run.status === 'running'),
      'workflow worker did not start',
    );
    const runs = registry.list({ source: 'workflow', workspaceId: 'workspace-a' });
    const room = runs.find((run) => run.kind === 'room')!;
    const worker = runs.find((run) => run.kind === 'worker')!;
    expect(worker.capabilities.cancel).toBe(false);
    expect(room.capabilities.cancel).toBe(true);
    await registry.control(room.id, 'cancel');
    await pending;
    expect(registry.get(room.id)?.status).toBe('cancelled');
    expect(registry.get(worker.id)?.status).toBe('cancelled');
    expect(signalBus.query({ teamId: `room::${room.id}` })[0]).toMatchObject({
      subtype: 'routed_share', content: { phase: 'cancelled' },
    });
  });

  it('tracks a parallel workflow as one Room with model-specific durable workers', async () => {
    const { registry, signalBus, server } = setup();
    const calls: Array<{ config: AgentLoopConfig; finish: ReturnType<typeof deferred<AgentResponse>> }> = [];
    const tools = bind(server, (config) => {
      const finish = deferred<AgentResponse>();
      calls.push({ config, finish });
      return finish.promise;
    }, 'workflow-session');
    const workflow = tools.find((item) => item.name === 'orchestrate_workflow')!;
    const pending = workflow.execute({
      task: 'Research and draft in parallel',
      inline_template: {
        name: 'Parallel pair',
        description: 'Two independent specialists',
        aggregation: 'concatenate',
        steps: [
          { name: 'Research', role: 'researcher', task: 'Research', model: 'model-research' },
          { name: 'Draft', role: 'writer', task: 'Draft', model: 'model-draft' },
        ],
      },
    });
    await waitFor(() => calls.length === 2, 'workflow workers did not start concurrently');
    const activeWorkers = registry.list({ source: 'workflow', workspaceId: 'workspace-a' })
      .filter((run) => run.kind === 'worker');
    expect(activeWorkers).toHaveLength(2);
    expect(activeWorkers.every((run) => run.status === 'running')).toBe(true);
    expect(new Set(calls.map((call) => call.config.model))).toEqual(new Set(['model-research', 'model-draft']));

    calls[0].finish.resolve({
      content: 'Research result', toolsUsed: ['read_file'], usage: { inputTokens: 3, outputTokens: 4 },
    });
    calls[1].finish.resolve({
      content: 'Draft result', toolsUsed: [], usage: { inputTokens: 5, outputTokens: 6 },
    });
    const output = await pending;
    expect(output).toContain('Research result');
    expect(output).toContain('Draft result');

    const runs = registry.list({ source: 'workflow', workspaceId: 'workspace-a' });
    const room = runs.find((run) => run.kind === 'room')!;
    const workers = runs.filter((run) => run.kind === 'worker');
    expect(room).toMatchObject({ status: 'completed', memoryRefs: { status: 'complete' } });
    expect(workers.every((run) => run.status === 'completed')).toBe(true);
    expect(workers.every((run) => run.memoryRefs.status === 'complete')).toBe(true);
    for (const worker of workers) {
      const chain = signalBus.query({ teamId: `room::${worker.roomId}` })
        .filter((item) => item.content.runId === worker.id);
      expect(chain.map((item) => item.subtype)).toEqual(expect.arrayContaining([
        'task_delegation', 'task_claim', 'routed_share',
      ]));
    }
  });
});