import assert from 'node:assert/strict';
import test from 'node:test';
import type { AgentTask } from '@pig/db';
import { z } from 'zod';
import { defineTool, PrimeOpenAIProvider } from '../src/provider';

const task = {
  id: '10000000-0000-4000-8000-000000000001',
  kind: 'enrich_account',
  subject: '20000000-0000-4000-8000-000000000002',
  reason: 'Extract the cited description.',
  payload: { sourceUrl: 'https://example.com/source' },
  priority: 0,
  budget: 2,
  attempts: 1,
  maxAttempts: 3,
  dueAt: new Date(),
  leasedUntil: new Date(),
  leasedBy: 'test',
  startedAt: new Date(),
  finishedAt: null,
  outcome: null,
  error: null,
  requestedByUserId: null,
  createdAt: new Date(),
} satisfies AgentTask;

test('Prime requests disable Nemotron reasoning and expose only supplied PIG tools', async () => {
  const bodies: Record<string, unknown>[] = [];
  let calls = 0;
  const fetchImpl: typeof fetch = async (_input, init) => {
    bodies.push(JSON.parse(String(init?.body)) as Record<string, unknown>);
    calls += 1;
    return Response.json(
      calls === 1
        ? {
            choices: [
              {
                message: {
                  content: null,
                  tool_calls: [
                    {
                      id: 'call_1',
                      type: 'function',
                      function: { name: 'pig_read', arguments: '{}' },
                    },
                  ],
                },
              },
            ],
            usage: { prompt_tokens: 10, completion_tokens: 5 },
          }
        : {
            choices: [{ message: { content: 'The cited record was inspected.' } }],
            usage: { prompt_tokens: 15, completion_tokens: 6 },
          },
    );
  };
  const provider = new PrimeOpenAIProvider({ apiKey: 'test', fetchImpl });

  const result = await provider.run({
    task,
    tools: [
      defineTool({
        name: 'pig_read',
        description: 'Read application data.',
        inputSchema: z.object({}).strict(),
        execute: async () => ({ name: 'Example' }),
      }),
    ],
  });

  assert.equal(result.summary, 'The cited record was inspected.');
  assert.equal(result.inputTokens, 25);
  assert.equal(result.outputTokens, 11);
  assert.equal(bodies.length, 2);
  for (const body of bodies) {
    assert.equal(body.model, 'nvidia/nemotron-3-nano-30b-a3b');
    assert.equal(body.reasoning_effort, 'none');
    assert.equal(body.parallel_tool_calls, false);
    const tools = body.tools as { function: { name: string } }[];
    assert.deepEqual(tools.map((tool) => tool.function.name), ['pig_read']);
    assert.ok(!JSON.stringify(tools).includes('bash'));
    assert.ok(!JSON.stringify(tools).includes('filesystem'));
  }
});
