import test from 'node:test';
import assert from 'node:assert/strict';
import { PassThrough, Writable } from 'node:stream';
import fs from 'node:fs';
import os from 'node:os';
import path from 'node:path';
import { Bridge } from '../bridge.js';

const start = { prompt: 'dummy', system_instruction: 'trusted', remaining_ms: 1000, model: { id: 'dummy', context_window_tokens: 10000, max_output_tokens: 1000 } };
function harness(t: any, script: (f: any, send: (type: string, payload: any, reply?: string) => void) => void, remaining = 1000) {
  const work = fs.mkdtempSync(path.join(os.tmpdir(), 'pi-bridge-')); t.after(() => fs.rmSync(work, { recursive: true, force: true }));
  const input = new PassThrough(), output = new PassThrough();
  const bridge = new Bridge({ stdin: input, stdout: output, workDir: work, onChildReap: async () => {} });
  const frames: any[] = []; let bytes = '', seq = 0;
  const send = (type: string, payload: any, reply_to?: string) => input.write(JSON.stringify({ v: 1, type, id: `b_${++seq}`, ...(reply_to ? { reply_to } : {}), payload }) + '\n');
  output.on('data', part => {
    bytes += part.toString(); let end;
    while ((end = bytes.indexOf('\n')) >= 0) { const f = JSON.parse(bytes.slice(0, end)); bytes = bytes.slice(end + 1); frames.push(f); script(f, send); }
  });
  const done = bridge.start(); send('start', { ...start, remaining_ms: remaining });
  return { input, bridge, frames, done, work };
}
const final = { content: [{ type: 'text', text: 'answer' }], stop_reason: 'stop', usage: { input_tokens: 2, output_tokens: 2 } };
test('zero artifacts completes only after collection_ready and stays alive until EOF', async t => {
  let complete!: () => void; const observed = new Promise<void>(r => complete = r);
  const h = harness(t, (f, send) => {
    if (f.type === 'model_request') send('model_response', { result: final, error: null }, f.id);
    if (f.type === 'collection_start') send('collection_ready', {}, f.id);
    if (f.type === 'complete') complete();
  });
  await observed; await new Promise<void>(resolve => setImmediate(resolve)); assert.equal(h.bridge.getState(), 'COMPLETE');
  let settled = false; h.done.then(() => settled = true);
  await new Promise(r => setTimeout(r, 10)); assert(!settled);
  h.input.end(); await h.done;
  assert.equal(h.frames.at(-1).payload.accepted_transfer_count, 0);
});
for (const mode of ['wrong-reply', 'wrong-type', 'duplicate', 'extra-field', 'eof', 'upstream', 'length', 'empty-answer']) {
  test(`terminal model/protocol failure: ${mode}`, async t => {
    const h = harness(t, (f, send) => {
      if (f.type !== 'model_request') return;
      if (mode === 'eof') { h.input.end(); return; }
      const result = mode === 'length' ? { ...final, stop_reason: 'length' } : mode === 'empty-answer' ? { ...final, content: [] } : final;
      const payload = mode === 'upstream' ? { result: null, error: { code: 'upstream_failed', message: 'SECRET_PAYLOAD' } } : { result, error: null };
      send(mode === 'wrong-type' ? 'tool_response' : 'model_response', mode === 'extra-field' ? { ...payload, bad: 1 } : payload, mode === 'wrong-reply' ? 'a_unknown' : f.id);
      if (mode === 'duplicate') send('model_response', payload, f.id);
    });
    await assert.rejects(h.done); assert.equal(h.bridge.getState(), 'FAILED');
    await new Promise(r => setTimeout(r, 5));
    assert(!h.frames.some(f => f.type === 'complete'));
    assert(!JSON.stringify(h.frames.filter(f => f.type === 'error')).includes('SECRET'));
  });
}
test('collection correlation and stalled collection retain deadline', async t => {
  for (const wrong of [true, false]) {
    const h = harness(t, (f, send) => {
      if (f.type === 'model_request') send('model_response', { result: final, error: null }, f.id);
      if (wrong && f.type === 'collection_start') send('collection_ready', {}, 'a_unknown');
    }, 200);
    await assert.rejects(h.done);
    assert.equal(h.bridge.getState(), 'FAILED');
  }
});
test('artifact begin/end acknowledgements must match request, transfer and phase', async t => {
  for (const phase of ['begin', 'end']) {
    const h = harness(t, (f, send) => {
      if (f.type === 'model_request') send('model_response', { result: final, error: null }, f.id);
      if (f.type === 'collection_start') send('collection_ready', {}, f.id);
      if (f.type === 'artifact_begin') send('artifact_ack', { transfer_id: f.payload.transfer_id, decision: phase === 'begin' ? 'stored' : 'accept', warning: null }, f.id);
      if (f.type === 'artifact_end') send('artifact_ack', { transfer_id: 'a_wrong', decision: 'stored', warning: null }, f.id);
    });
    fs.mkdirSync(path.join(h.work, 'artifacts'), { recursive: true }); fs.writeFileSync(path.join(h.work, 'artifacts', 'empty'), '');
    await assert.rejects(h.done); assert(!h.frames.some(f => f.type === 'complete'));
  }
});

test('concurrent remote calls correlate independently and stay bounded', async t => {
  let modelFrame: any, sendReply!: (type: string, payload: any, reply?: string) => void;
  let modelReady!: () => void; const ready = new Promise<void>(resolve => modelReady = resolve);
  const h = harness(t, (f, send) => {
    sendReply = send;
    if (f.type === 'model_request') { modelFrame = f; modelReady(); }
    if (f.type === 'tool_request') send('tool_response', { result: { page_id: f.payload.parameters.page_id, title: '', space: '', url: '', markdown: '', truncated: false }, error: null }, f.id);
    if (f.type === 'collection_start') send('collection_ready', {}, f.id);
    if (f.type === 'complete') h.input.end();
  });
  await ready;
  const calls = ['1', '2', '3'].map(page_id => h.bridge.sendToolRequest('confluence_view', { page_id }));
  await assert.rejects(h.bridge.sendToolRequest('confluence_view', { page_id: '4' }), /limit/);
  assert.deepEqual((await Promise.all(calls)).map(x => x.page_id), ['1', '2', '3']);
  sendReply('model_response', { result: final, error: null }, modelFrame.id);
  await h.done;
  assert.equal(h.frames.filter(f => f.type === 'tool_request').length, 3);
});
test('model call total cannot exceed fifty turns', async t => {
  let calls = 0;
  const h = harness(t, (f, send) => {
    if (f.type !== 'model_request') return;
    calls++;
    send('model_response', { result: { content: [{ type: 'tool_call', id: `sdk_${calls}`, name: 'read', arguments: { path: '/nonexistent-pi-fixture' } }], stop_reason: 'tool_calls', usage: { input_tokens: 1, output_tokens: 1 } }, error: null }, f.id);
  }, 2000);
  await assert.rejects(h.done); assert.equal(calls, 50); assert(!h.frames.some(f => f.type === 'complete'));
});

for (const delayed of [false, true]) {
  test(`accepted complete stays terminal beyond remaining_ms, delayed model=${delayed}`, async t => {
    let observed!: () => void; const complete = new Promise<void>(resolve => observed = resolve);
    const h = harness(t, (f, send) => {
      if (f.type === 'model_request') {
        const respond = () => send('model_response', { result: final, error: null }, f.id);
        if (delayed) setTimeout(respond, 200); else respond();
      }
      if (f.type === 'collection_start') send('collection_ready', {}, f.id);
      if (f.type === 'complete') observed();
    }, 400);
    let settled = false; let failure: unknown;
    h.done.then(() => settled = true, error => failure = error);
    await complete;
    await new Promise(resolve => setTimeout(resolve, 500));
    h.bridge.failRun('query_timeout');
    assert.equal(h.bridge.getState(), 'COMPLETE'); assert.equal(failure, undefined); assert(!settled);
    assert.deepEqual(h.frames.map(f => f.type), ['model_request', 'collection_start', 'complete']);
    h.input.end(); await h.done;
  });
}
test('stalled model still expires before completion', async t => {
  const h = harness(t, () => {}, 80);
  await assert.rejects(h.done, /deadline expired/);
  assert.equal(h.bridge.getState(), 'FAILED'); assert(await h.bridge.flushOutput());
  assert.equal(h.frames.at(-1).type, 'error');
});
test('synchronous EOF cannot turn a failed complete write into success', async t => {
  const work = fs.mkdtempSync(path.join(os.tmpdir(), 'pi-complete-write-'));
  t.after(() => fs.rmSync(work, { recursive: true, force: true }));
  const input = new PassThrough(); let seq = 0; const frames: any[] = [];
  const send = (type: string, payload: any, reply_to: string) => input.write(JSON.stringify({ v: 1, type, id: `b_${++seq}`, reply_to, payload }) + '\n');
  const output = new Writable({ write(chunk, _encoding, callback) {
    const f = JSON.parse(chunk.toString()); frames.push(f);
    if (f.type === 'model_request') send('model_response', { result: final, error: null }, f.id);
    if (f.type === 'collection_start') send('collection_ready', {}, f.id);
    if (f.type === 'complete') { input.end(); setTimeout(() => callback(new Error('DUMMY_PRIVATE_STREAM_ERROR')), 20); }
    else callback();
  } });
  const bridge = new Bridge({ stdin: input, stdout: output, workDir: work, onChildReap: async () => {} });
  const done = bridge.start(); input.write(JSON.stringify({ v: 1, type: 'start', id: `b_${++seq}`, payload: start }) + '\n');
  await assert.rejects(done); assert.equal(bridge.getState(), 'FAILED');
  assert.equal(await bridge.flushOutput(), false);
  assert.deepEqual(frames.map(f => f.type), ['model_request', 'collection_start', 'complete']);
});
test('failure flushing waits for slow callbacks and stays bounded if output stalls', async () => {
  for (const stalled of [false, true]) {
    let release!: () => void; const chunks: string[] = [];
    const output = new Writable({ write(chunk, _encoding, callback) {
      chunks.push(chunk.toString()); release = () => callback();
      if (!stalled) setTimeout(release, 60);
    } });
    const input = new PassThrough(); const bridge = new Bridge({ stdin: input, stdout: output });
    const done = bridge.start(); bridge.failRun('invalid_input'); await assert.rejects(done);
    const started = Date.now(); const flushed = await bridge.flushOutput(stalled ? 40 : 500);
    assert.equal(flushed, !stalled); assert(Date.now() - started < 600);
    if (!stalled) assert(Date.now() - started >= 40);
    assert.equal(JSON.parse(chunks.join('')).type, 'error');
    if (stalled) release(); input.end();
  }
});
test('collection warnings retain repeated and overflow counts on the wire', async t => {
  let captured: any[] = [];
  const h = harness(t, (f, send) => {
    if (f.type === 'model_request') send('model_response', { result: final, error: null }, f.id);
    if (f.type === 'collection_start') { captured = f.payload.warnings; send('collection_ready', {}, f.id); }
    if (f.type === 'complete') h.input.end();
  });
  for (let i = 0; i < 105; i++) h.bridge.addWarning({ code: `warning_${i}`, message: 'Bounded warning.' });
  h.bridge.addWarning({ code: 'warning_0', message: 'Bounded warning.' });
  await h.done; assert.equal(captured.length, 100);
  assert.match(captured[0].message, /Repeated 2 times/);
  assert.equal(captured.at(-1).code, 'warnings_aggregated'); assert.match(captured.at(-1).message, /^6 additional/);
});

test('deadline remains active until the complete write succeeds', async t => {
  const work = fs.mkdtempSync(path.join(os.tmpdir(), 'pi-complete-stall-'));
  t.after(() => fs.rmSync(work, { recursive: true, force: true }));
  const input = new PassThrough(); let seq = 0, release!: () => void; const frames: any[] = [];
  const send = (type: string, payload: any, reply_to: string) => input.write(JSON.stringify({ v: 1, type, id: `b_${++seq}`, reply_to, payload }) + '\n');
  const output = new Writable({ write(chunk, _encoding, callback) {
    const f = JSON.parse(chunk.toString()); frames.push(f);
    if (f.type === 'model_request') send('model_response', { result: final, error: null }, f.id);
    if (f.type === 'collection_start') send('collection_ready', {}, f.id);
    if (f.type === 'complete') release = () => callback(); else callback();
  } });
  const bridge = new Bridge({ stdin: input, stdout: output, workDir: work, onChildReap: async () => {} });
  const done = bridge.start(); input.write(JSON.stringify({ v: 1, type: 'start', id: `b_${++seq}`, payload: { ...start, remaining_ms: 150 } }) + '\n');
  await assert.rejects(done, /deadline expired/); assert.equal(bridge.getState(), 'FAILED');
  release(); assert(await bridge.flushOutput()); input.end();
  assert.deepEqual(frames.map(f => f.type), ['model_request', 'collection_start', 'complete']);
});

test('EOF while complete is still queued fails before starting the terminal frame', async t => {
  const work = fs.mkdtempSync(path.join(os.tmpdir(), 'pi-queued-complete-'));
  t.after(() => fs.rmSync(work, { recursive: true, force: true }));
  fs.mkdirSync(path.join(work, 'artifacts')); fs.writeFileSync(path.join(work, 'artifacts', 'empty'), '');
  const input = new PassThrough(); let seq = 0, release!: () => void; const frames: any[] = [];
  let stored!: () => void; const acknowledged = new Promise<void>(resolve => stored = resolve);
  const send = (type: string, payload: any, reply_to: string) => input.write(JSON.stringify({ v: 1, type, id: `b_${++seq}`, reply_to, payload }) + '\n');
  const output = new Writable({ write(chunk, _encoding, callback) {
    const f = JSON.parse(chunk.toString()); frames.push(f);
    if (f.type === 'model_request') send('model_response', { result: final, error: null }, f.id);
    if (f.type === 'collection_start') send('collection_ready', {}, f.id);
    if (f.type === 'artifact_begin') send('artifact_ack', { transfer_id: f.payload.transfer_id, decision: 'accept', warning: null }, f.id);
    if (f.type === 'artifact_end') {
      release = () => callback();
      send('artifact_ack', { transfer_id: f.payload.transfer_id, decision: 'stored', warning: null }, f.id); stored();
    } else callback();
  } });
  const bridge = new Bridge({ stdin: input, stdout: output, workDir: work, onChildReap: async () => {} });
  const done = bridge.start(); const failed = assert.rejects(done, /connection was lost/);
  input.write(JSON.stringify({ v: 1, type: 'start', id: `b_${++seq}`, payload: start }) + '\n');
  await acknowledged; await new Promise<void>(resolve => setImmediate(resolve));
  input.end(); await failed; release(); assert(await bridge.flushOutput());
  assert(!frames.some(f => f.type === 'complete')); assert.equal(frames.at(-1).type, 'error');
});