import test from 'node:test';
import assert from 'node:assert/strict';
import { Agent, runAgent } from '../runtime/agent.js';
import {
  MiniCPM5ForCausalLM,
  BusyError,
  ContextOverflowError,
  DeviceLostError,
} from '../runtime/minicpm5.js';

const toolCall = (name, args = {}, id = name, errors) => ({
  id,
  type: 'function',
  function: { name, arguments: args },
  ...(errors ? { errors } : {}),
});

function reply(content = '', calls) {
  return {
    finish_reason: calls ? 'tool_calls' : 'stop',
    message: { role: 'assistant', content, ...(calls ? { tool_calls: calls } : {}) },
    usage: { prompt_tokens: 10, cached_tokens: 2, completion_tokens: 4, reasoning_tokens: 1 },
    timings: { ttft_ms: 1, decode_ms: 3, tokens_per_second: 1000 },
  };
}

function fakeSession(steps) {
  const requests = [];
  let disposed = 0;
  return {
    requests,
    get disposed() { return disposed; },
    dispose() { disposed += 1; },
    generate(history, options) {
      requests.push({ history, options });
      const step = steps[requests.length - 1];
      assert.ok(step, `unexpected generation round ${requests.length}`);
      return {
        result: step.error ? Promise.reject(step.error) : Promise.resolve(step.result ?? step),
        async *[Symbol.asyncIterator]() {
          if (step.error) throw step.error;
          for (const event of step.events ?? []) yield event;
        },
      };
    },
  };
}

async function collect(run, onEvent = () => {}) {
  const events = [];
  for await (const event of run) {
    events.push(event);
    await onEvent(event);
  }
  return { events, result: await run.result };
}

const offeredTool = (name, execute) => ({
  name,
  description: `${name} test tool`,
  parameters: { type: 'object', properties: { value: { type: 'number' } } },
  execute,
});

test('runtime imports in Node without requesting a GPU or downloading weights', () => {
  assert.equal(typeof MiniCPM5ForCausalLM.from_pretrained, 'function');
  assert.equal(typeof MiniCPM5ForCausalLM.checkAvailability, 'function');
  assert.equal(MiniCPM5ForCausalLM.DEFAULT_MODEL_ID, 'openbmb/MiniCPM5-2B-GGUF');
  assert.equal(new BusyError('busy').name, 'BusyError');
  const overflow = new ContextOverflowError(9000, 8192);
  assert.equal(overflow.required, 9000);
  assert.equal(overflow.capacity, 8192);
  assert.equal(new DeviceLostError('unknown', 'test').name, 'DeviceLostError');
});

test('successful tools run sequentially and feed their text back before the final answer', async () => {
  const calls = [toolCall('first', { value: 3 }, 'one'), toolCall('second', { value: 4 }, 'two')];
  const finalReply = reply('The answer is 7.');
  const session = fakeSession([reply('', calls), finalReply]);
  const order = [];
  let active = false;
  const tools = ['first', 'second'].map(name => offeredTool(name, async (args, { signal, call }) => {
    assert.equal(active, false, 'tools must not overlap');
    assert.equal(signal.aborted, false);
    assert.equal(call.function.name, name);
    active = true;
    order.push(`start:${name}`);
    await Promise.resolve();
    order.push(`end:${name}`);
    active = false;
    return { content: String(args.value), details: { source: name } };
  }));
  const agent = new Agent({ model: session, instructions: 'Use the tools.', tools });
  const { events, result } = await collect(agent.run('Add the two values.'));

  assert.deepEqual(order, ['start:first', 'end:first', 'start:second', 'end:second']);
  assert.equal(result.finish_reason, 'stop');
  assert.equal(result.content, 'The answer is 7.');
  assert.equal(result.rounds, 2);
  assert.equal(result.usage.prompt_tokens, 20);
  assert.equal(result.usage.completion_tokens, 8);
  assert.deepEqual(session.requests[1].history.slice(-2), [
    { role: 'tool', tool_call_id: 'one', name: 'first', content: '3' },
    { role: 'tool', tool_call_id: 'two', name: 'second', content: '4' },
  ]);
  assert.deepEqual(session.requests[0].options.tools, session.requests[1].options.tools);
  assert.deepEqual(events.filter(event => event.type === 'tool_end').map(event => event.details), [
    { source: 'first' }, { source: 'second' },
  ]);
  assert.ok(events.every(event => Number.isInteger(event.round)));
  assert.strictEqual(agent.messages.at(-1), finalReply.message);
});

test('malformed and unknown calls become tool errors without executing a tool', async () => {
  let executions = 0;
  const calls = [
    toolCall('known', { value: 'wrong' }, 'invalid', ['value must be a number']),
    toolCall('unavailable', {}, 'unknown'),
  ];
  const agent = new Agent({
    model: fakeSession([reply('', calls), reply('Please correct the request.')]),
    tools: [offeredTool('known', () => { executions += 1; return 'unexpected'; })],
  });
  const { events } = await collect(agent.run('Try those calls.'));
  const ended = events.filter(event => event.type === 'tool_end');

  assert.equal(executions, 0);
  assert.equal(ended.length, 2);
  assert.ok(ended.every(event => event.is_error));
  assert.equal(ended[0].message.content, 'value must be a number');
  assert.match(ended[1].message.content, /No tool named "unavailable"/);
  assert.deepEqual(ended.map(event => event.message.tool_call_id), ['invalid', 'unknown']);
});

test('identical failed calls are refused until another tool succeeds', async () => {
  let failures = 0;
  const fail = () => reply('', [toolCall('fails', { value: 1 })]);
  const session = fakeSession([
    fail(), fail(), reply('', [toolCall('succeeds')]), fail(), reply('Unable to finish.'),
  ]);
  const agent = new Agent({ model: session, tools: [
    offeredTool('fails', () => { failures += 1; throw new Error('Operation failed'); }),
    offeredTool('succeeds', () => 'Success'),
  ] });
  const { events, result } = await collect(agent.run('Attempt the operation.'));
  const ended = events.filter(event => event.type === 'tool_end');

  assert.equal(failures, 2, 'the duplicate must be skipped, then allowed again after success');
  assert.equal(ended[0].message.content, 'Operation failed');
  assert.match(ended[1].message.content, /already failed/);
  assert.equal(ended[2].is_error, false);
  assert.equal(ended[3].message.content, 'Operation failed');
  assert.equal(result.rounds, 5);
});

test('cancelling a running tool answers that call and every remaining call', async () => {
  let firstStarted = false;
  let secondStarted = false;
  const calls = [toolCall('wait', {}, 'running'), toolCall('later', {}, 'pending')];
  const agent = new Agent({ model: fakeSession([reply('', calls)]), tools: [
    offeredTool('wait', async (_args, { signal }) => {
      firstStarted = true;
      await new Promise((_resolve, reject) => {
        signal.addEventListener('abort', () => reject(signal.reason), { once: true });
      });
      return 'Should not complete';
    }),
    offeredTool('later', () => { secondStarted = true; return 'Should not run'; }),
  ] });
  const run = agent.run('Run both tools.');
  const { events, result } = await collect(run, event => {
    if (event.type === 'tool_start') run.cancel();
  });
  const ended = events.filter(event => event.type === 'tool_end');

  assert.equal(firstStarted, true);
  assert.equal(secondStarted, false);
  assert.equal(result.finish_reason, 'cancelled');
  assert.equal(agent.running, false);
  assert.deepEqual(ended.map(event => event.message.tool_call_id), ['running', 'pending']);
  assert.ok(ended.every(event => event.is_error));
  assert.match(ended[0].message.content, /may have partly completed/);
  assert.match(ended[1].message.content, /before this call ran/);
  assert.deepEqual(agent.messages.map(message => message.role), ['user', 'assistant', 'tool', 'tool']);
  assert.doesNotThrow(() => agent.reset(), 'cancellation must settle before resetting');
  assert.deepEqual(agent.messages, []);
});

test('Agent preserves canonical assistant objects and the full history across runs', async () => {
  const first = reply('First answer.');
  first.message.reasoning_content = 'Private reasoning text.';
  Object.freeze(first.message);
  const second = reply('Second answer.');
  const session = fakeSession([first, second]);
  const agent = new Agent({ model: session, instructions: 'Help.' });
  await collect(agent.run('First question.'));
  await collect(agent.run('Second question.'));

  assert.strictEqual(session.requests[1].history[2], first.message);
  assert.deepEqual(session.requests[1].history.map(message => message.role), [
    'system', 'user', 'assistant', 'user',
  ]);
  assert.strictEqual(agent.messages[2], first.message);
  assert.strictEqual(agent.messages[4], second.message);
  agent.reset();
  assert.deepEqual(agent.messages, [{ role: 'system', content: 'Help.' }]);
});

test('Agent reuses and releases its owned session, while a supplied session remains open', async () => {
  const owned = fakeSession([reply('One'), reply('Two')]);
  let created = 0;
  let modelsDisposed = 0;
  const model = {
    create_session() { created += 1; return owned; },
    dispose() { modelsDisposed += 1; },
  };
  const agent = new Agent({ model });
  assert.equal(created, 0, 'the session is allocated lazily');
  await collect(agent.run('One'));
  await collect(agent.run('Two'));
  assert.equal(created, 1);
  await agent.dispose();
  await agent.dispose();
  assert.equal(owned.disposed, 1);
  assert.equal(modelsDisposed, 0, 'the caller retains ownership of the loaded model');

  const supplied = fakeSession([reply('Answer')]);
  const externalAgent = new Agent({ model: supplied });
  await collect(externalAgent.run('Question'));
  await externalAgent.dispose();
  assert.equal(supplied.disposed, 0);
});

test('a cancelled partial reply is emitted but is omitted from canonical Agent history', async () => {
  const partial = { ...reply('Partial answer'), finish_reason: 'cancelled' };
  const agent = new Agent({ model: fakeSession([{
    result: partial,
    events: [{ type: 'text', text: 'Partial answer' }],
  }]), instructions: 'Help.' });
  const { events, result } = await collect(agent.run('Question'));

  assert.equal(events.find(event => event.type === 'text').text, 'Partial answer');
  assert.strictEqual(events.find(event => event.type === 'turn_end').result.message, partial.message);
  assert.equal(result.finish_reason, 'cancelled');
  assert.equal(result.content, '');
  assert.deepEqual(result.messages, []);
  assert.deepEqual(agent.messages, [{ role: 'system', content: 'Help.' }]);
});

test('max_rounds stops an ongoing tool loop after answering the final round of calls', async () => {
  let executions = 0;
  const session = fakeSession([
    reply('', [toolCall('repeat', { value: 1 }, 'first')]),
    reply('', [toolCall('repeat', { value: 2 }, 'second')]),
  ]);
  const agent = new Agent({ model: session, max_rounds: 2, tools: [
    offeredTool('repeat', () => { executions += 1; return 'Continue'; }),
  ] });
  const { result } = await collect(agent.run('Keep going.'));

  assert.equal(result.finish_reason, 'max_rounds');
  assert.equal(result.rounds, 2);
  assert.equal(executions, 2);
  assert.deepEqual(result.messages.map(message => message.role), ['assistant', 'tool', 'assistant', 'tool']);
  assert.equal(agent.messages.at(-1).tool_call_id, 'second');
});

test('runAgent leaves caller history unchanged and returns canonical additions', async () => {
  const history = Object.freeze([{ role: 'user', content: 'Hello' }]);
  const finalReply = reply('Hi');
  const { result } = await collect(runAgent(fakeSession([finalReply]), history));

  assert.equal(history.length, 1);
  assert.equal(result.messages.length, 1);
  assert.strictEqual(result.messages[0], finalReply.message);
});