import type { ChatMessage } from '@genesislcap/foundation-ai';
import { assert, createLogicSuite } from '@genesislcap/foundation-testing';
import { flattenSubAgentMessages } from './flatten-sub-agent-messages';

const msg = (overrides: Partial<ChatMessage> = {}): ChatMessage => ({
  role: 'assistant',
  content: '',
  ...overrides,
});

/** A parent assistant message whose single tool call delegated to a sub-agent. */
const delegating = (id: string, agentName: string, trace: ChatMessage[]): ChatMessage =>
  msg({
    agentName,
    toolCalls: [{ id, name: 'requestSubAgent', args: {}, subAgentTrace: trace }],
  });

const suite = createLogicSuite('flattenSubAgentMessages');

suite('returns an empty array for an empty list', () => {
  assert.equal(flattenSubAgentMessages([]), []);
});

suite('emits top-level messages unchanged and tags none as sub-agent', () => {
  const out = flattenSubAgentMessages([
    msg({ role: 'user', content: 'hi', agentName: 'UI Builder' }),
    msg({ content: 'hello', agentName: 'UI Builder' }),
  ]);
  assert.is(out.length, 2);
  assert.is(out[0].kind, 'message');
  assert.is(out[0].agentName, 'UI Builder');
  // No depth/marker fields on top-level entries.
  assert.is(out[0].subAgentDepth, undefined);
  assert.is(out[0].subAgentName, undefined);
  assert.is(out[1].subAgentOf, undefined);
});

suite('hoists a sub-agent trace to top-level entries after the spawning message', () => {
  const trace = [
    msg({
      agentName: 'UI Architecture Planner',
      content: '',
      toolCalls: [{ id: 'g', name: 'grep_source', args: {} }],
    }),
    msg({
      role: 'tool',
      agentName: 'UI Architecture Planner',
      toolResult: { toolCallId: 'g', content: 'result' },
    }),
  ];
  const out = flattenSubAgentMessages([
    delegating('tc1', 'UI Builder', trace),
    msg({
      role: 'tool',
      agentName: 'UI Builder',
      toolResult: { toolCallId: 'tc1', content: 'plan done' },
    }),
  ]);

  // parent tool-call msg, then 2 hoisted children, then the parent tool-result msg.
  assert.is(out.length, 4);
  assert.is(out[0].agentName, 'UI Builder'); // parent unchanged
  assert.is(out[1].agentName, 'UI Builder › UI Architecture Planner');
  assert.is(out[2].agentName, 'UI Builder › UI Architecture Planner');
  assert.is(out[3].agentName, 'UI Builder'); // parent tool-result, back at top level
});

suite('marks hoisted entries with depth, raw name, and the spawning tool-call id', () => {
  const out = flattenSubAgentMessages([
    delegating('tc1', 'UI Builder', [msg({ agentName: 'UI Architecture Planner', content: 'x' })]),
  ]);
  const child = out[1];
  assert.is(child.subAgentDepth, 1);
  assert.is(child.subAgentName, 'UI Architecture Planner');
  assert.is(child.subAgentOf, 'tc1');
});

suite(
  'moves the trace out — the emitted parent tool call no longer carries it (no duplication)',
  () => {
    const trace = [msg({ agentName: 'Planner', content: 'only once' })];
    const out = flattenSubAgentMessages([delegating('tc1', 'UI Builder', trace)]);
    // The trace is stripped from the emitted tool call ...
    assert.is(out[0].toolCalls![0].subAgentTrace, undefined);
    // ... and surfaces exactly once, as a hoisted top-level entry.
    const hoisted = out.filter((e) => e.subAgentDepth === 1);
    assert.is(hoisted.length, 1);
    assert.is(hoisted[0].content, 'only once');
  },
);

suite('preserves per-message usage fields on hoisted entries', () => {
  const out = flattenSubAgentMessages([
    delegating('tc1', 'UI Builder', [
      msg({
        agentName: 'Planner',
        model: 'claude-sonnet-4-6',
        inputTokens: 100,
        outputTokens: 20,
        cost: 0.5,
      }),
    ]),
  ]);
  const child = out[1];
  assert.is(child.model, 'claude-sonnet-4-6');
  assert.is(child.inputTokens, 100);
  assert.is(child.outputTokens, 20);
  assert.is(child.cost, 0.5);
});

suite('composes the breadcrumb for nested sub-agents', () => {
  const grandchild = [msg({ agentName: 'Grounding', content: 'deep' })];
  const childTrace = [delegating('tc2', 'Planner', grandchild)];
  const out = flattenSubAgentMessages([delegating('tc1', 'UI Builder', childTrace)]);

  // top: UI Builder | depth1: Planner | depth2: Grounding
  const depth1 = out.find((e) => e.subAgentDepth === 1)!;
  const depth2 = out.find((e) => e.subAgentDepth === 2)!;
  assert.is(depth1.agentName, 'UI Builder › Planner');
  assert.is(depth2.agentName, 'UI Builder › Planner › Grounding');
  assert.is(depth2.subAgentName, 'Grounding');
  assert.is(depth2.subAgentOf, 'tc2');
});

suite('correlates two sub-agents spawned in the same parent turn', () => {
  const parent = msg({
    agentName: 'UI Builder',
    toolCalls: [
      {
        id: 'a',
        name: 'requestSubAgent',
        args: {},
        subAgentTrace: [msg({ agentName: 'Planner', content: 'p' })],
      },
      {
        id: 'b',
        name: 'requestSubAgent',
        args: {},
        subAgentTrace: [msg({ agentName: 'Reviewer', content: 'r' })],
      },
    ],
  });
  const out = flattenSubAgentMessages([parent]);
  const planner = out.find((e) => e.subAgentName === 'Planner')!;
  const reviewer = out.find((e) => e.subAgentName === 'Reviewer')!;
  assert.is(planner.subAgentOf, 'a');
  assert.is(reviewer.subAgentOf, 'b');
  assert.is(planner.agentName, 'UI Builder › Planner');
  assert.is(reviewer.agentName, 'UI Builder › Reviewer');
});

suite('falls back gracefully when the spawning message has no agentName', () => {
  const out = flattenSubAgentMessages([
    delegating('tc1', undefined as unknown as string, [
      msg({ agentName: 'Planner', content: 'x' }),
    ]),
  ]);
  // No parent name to prefix → show the sub-agent's own name unbreadcrumbed.
  assert.is(out[1].agentName, 'Planner');
  assert.is(out[1].subAgentName, 'Planner');
});

suite.run();
