import type {
  AIProvider,
  AIProviderRegistry,
  ChatMessage,
  ChatRequestOptions,
} from '@genesislcap/foundation-ai';
import { assert, createLogicSuite } from '@genesislcap/foundation-testing';
import type { AgentConfig } from '../../config/config';
import { sumUsage } from '../../utils/sum-usage';
import { usageRows } from '../../utils/usage-rows';
// Side-effect import — MUST come before `./chat-driver` so the driver subclasses
// jsdom's EventTarget rather than Node's native one. Mirrors chat-driver.test.ts.
import './align-event-globals';
import { ChatDriver } from './chat-driver';

// ---------------------------------------------------------------------------
// Sub-agent trace capture when ONE tool call invokes SEVERAL sub-agents.
//
// The capture slot used to be a single `{ trace?: ChatMessage[] }` per tool call,
// assigned by each `requestSubAgent`. That covers N tool calls × 1 sub-agent — the
// case its docblock named — but not 1 tool call × N sub-agents, which is what a
// code-driven scheduler or a retry produces. All but the last trace was dropped, so
// those children ran, were billed, and then had no record in history.
//
// The failure was invisible to every existing check: `sumUsage` and `usageRows` both
// recurse into `subAgentTrace`, so they summed a truncated input and still reconciled
// with each other perfectly. Measured at 1 of 7 traces kept on a seven-way fan-out,
// ~2.3x under the true cost.
//
// These tests assert the traces SURVIVE, which is upstream of any cost assertion.
// ---------------------------------------------------------------------------

/** Answers by rule: concurrent children under one turn make a FIFO queue nondeterministic. */
const ruleProvider = (): AIProvider => {
  let parentTurns = 0;
  return {
    chat: async (
      _history: ChatMessage[],
      _userMessage: string,
      options?: ChatRequestOptions,
    ): Promise<ChatMessage> => {
      const tools = (options?.tools ?? []).map((t) => t.name);
      // Child turn: finishes via its completion tool, and reports usage so the
      // reconciliation assertions have something to add up.
      if (tools.includes('work')) {
        return {
          role: 'assistant',
          content: '',
          model: 'claude-sonnet-5',
          cost: 0.01,
          inputTokens: 100,
          outputTokens: 20,
          toolCalls: [{ id: 'w1', name: 'work', args: {} }],
        };
      }
      if (tools.includes('delegate') && parentTurns === 0) {
        parentTurns += 1;
        return {
          role: 'assistant',
          content: '',
          model: 'claude-sonnet-5',
          cost: 0.02,
          inputTokens: 200,
          outputTokens: 30,
          toolCalls: [{ id: 'd0', name: 'delegate', args: {} }],
        };
      }
      return { role: 'assistant', content: 'done', model: 'claude-sonnet-5' };
    },
  };
};

const makeRegistry = (provider: AIProvider): AIProviderRegistry => ({
  get: () => provider,
  default: () => provider,
  defaultName: () => 'test',
  names: () => ['test'],
  getStatus: async () => null,
  listStatuses: async () => [],
});

const agent = (overrides: Partial<AgentConfig> & { name: string }): AgentConfig =>
  ({ description: 'test agent', ...overrides }) as AgentConfig;

/** A child that reports one unit of usage and completes. */
const worker = (name: string): AgentConfig =>
  agent({
    name,
    toolDefinitions: [
      { name: 'work', description: 'work', parameters: { type: 'object', properties: {} } },
    ],
    toolHandlers: {
      work: async (_args, ctx) => {
        ctx.completeSubAgent?.({ ok: true });
        return 'worked';
      },
    },
  });

/**
 * A parent whose single `delegate` tool call invokes `names` in turn — the shape a
 * code-driven scheduler produces, and the one the old single slot truncated.
 */
const boss = (children: AgentConfig[], invoke: string[]): AgentConfig =>
  agent({
    name: 'boss',
    subAgents: children,
    toolDefinitions: [
      { name: 'delegate', description: 'delegate', parameters: { type: 'object', properties: {} } },
    ],
    toolHandlers: {
      delegate: async (_args, ctx) => {
        await Promise.all(invoke.map((n) => ctx.requestSubAgent!(n, { task: 'go' })));
        return 'delegated';
      },
    },
  });

const run = async (config: AgentConfig): Promise<readonly ChatMessage[]> => {
  const driver = new ChatDriver(makeRegistry(ruleProvider()), {
    maxToolIterations: 20,
    maxFoldOperations: 5,
    sessionKey: '',
  });
  driver.applyAgent(config);
  await driver.sendMessage('go');
  return driver.getHistory();
};

const tracesIn = (history: readonly ChatMessage[]): readonly ChatMessage[][] =>
  history.flatMap((m) =>
    (m.toolCalls ?? []).flatMap((tc) => (tc.subAgentTrace ? [tc.subAgentTrace] : [])),
  );

const suite = createLogicSuite('ChatDriver sub-agent trace capture');

// ---------------------------------------------------------------------------

suite('keeps every trace when one tool call invokes three sub-agents', async () => {
  const names = ['gen_a', 'gen_b', 'gen_c'];
  const history = await run(boss(names.map(worker), names));

  const traces = tracesIn(history);
  assert.is(traces.length, 1, 'one tool call, so one concatenated trace');

  // Each child contributes at least its own assistant turn. Before the fix this was
  // one child's worth regardless of how many ran.
  const seen = new Set(traces[0].map((m) => m.agentName).filter(Boolean));
  assert.equal(
    [...seen].sort(),
    names,
    `every invoked child must appear — got ${JSON.stringify([...seen])}`,
  );
});

suite('prices every child — usageRows emits a row per child and reconciles', async () => {
  // The reconciliation test extended to a fan-out. The single-child version passed
  // before the fix AND after it, which is exactly why it did not catch this.
  const names = ['gen_a', 'gen_b', 'gen_c'];
  const history = await run(boss(names.map(worker), names));

  const rows = usageRows(history);
  const childRows = rows.filter((r) => r.subAgentDepth === 1);
  assert.is(childRows.length, names.length, 'one row per child that ran');

  const rowTotal = rows.reduce((n, r) => n + (r.costUsd ?? 0) + (r.externalCostUsd ?? 0), 0);
  assert.ok(
    Math.abs(rowTotal - sumUsage(history).costUsd) < 1e-12,
    `rows ${rowTotal} vs sumUsage ${sumUsage(history).costUsd}`,
  );
  // Parent turn ($0.02) + three children ($0.01 each). Asserted as a number so a
  // regression that silently drops a child fails here rather than only in the
  // reconciliation above, which would still agree with a truncated input.
  assert.ok(Math.abs(rowTotal - 0.05) < 1e-12, `expected 0.05, got ${rowTotal}`);
});

suite('keeps both traces when the same sub-agent is invoked twice (a retry)', async () => {
  // The retry shape: one handler, one child name, two invocations. Under the old slot
  // the first attempt's trace was overwritten by the second.
  const history = await run(boss([worker('gen_a')], ['gen_a', 'gen_a']));

  const traces = tracesIn(history);
  assert.is(traces.length, 1);
  const childTurns = traces[0].filter((m) => m.agentName === 'gen_a' && m.cost != null);
  assert.is(childTurns.length, 2, 'both attempts survive, not just the last');

  const rows = usageRows(history).filter((r) => r.subAgentDepth === 1);
  assert.is(rows.length, 2, 'and both are priced');
});

suite('keeps child traces when the parent handler throws after they ran', async () => {
  // Children that completed were billed by the provider. If the error path drops
  // their traces, the run under-reports exactly as the single-slot bug did — and
  // just as silently, since `sumUsage` and `usageRows` would still agree with each
  // other over the truncated input. PR review.
  const names = ['gen_a', 'gen_b'];
  const parent = agent({
    name: 'boss',
    subAgents: names.map(worker),
    toolDefinitions: [
      { name: 'delegate', description: 'delegate', parameters: { type: 'object', properties: {} } },
    ],
    toolHandlers: {
      delegate: async (_args, ctx) => {
        await Promise.all(names.map((n) => ctx.requestSubAgent!(n, { task: 'go' })));
        throw new Error('post-processing the children failed');
      },
    },
  });

  const history = await run(parent);
  const traces = tracesIn(history);
  assert.is(traces.length, 1, 'the failed tool call still carries its trace');
  const seen = [...new Set(traces[0].map((m) => m.agentName).filter(Boolean))].sort();
  assert.equal(seen, names, 'both children survive the throw');

  const childRows = usageRows(history).filter((r) => r.subAgentDepth === 1);
  assert.is(childRows.length, 2, 'and both are still priced');
});

suite('leaves subAgentTrace undefined when no sub-agent ran', async () => {
  // Presence is the signal readers key off (`usage-rows.ts`, the UI's `when(...)`),
  // so an empty array would be a different and wrong claim.
  const history = await run(
    agent({
      name: 'boss',
      subAgents: [worker('gen_a')],
      toolDefinitions: [
        {
          name: 'delegate',
          description: 'delegate',
          parameters: { type: 'object', properties: {} },
        },
      ],
      toolHandlers: { delegate: async () => 'did it myself' },
    }),
  );

  for (const m of history) {
    for (const tc of m.toolCalls ?? []) {
      assert.is(tc.subAgentTrace, undefined, 'not an empty array');
    }
  }
  assert.equal(
    usageRows(history).filter((r) => r.subAgentDepth === 1),
    [],
    'no child rows',
  );
});

suite.run();
