const assert = require('node:assert/strict');
const fs = require('node:fs/promises');
const os = require('node:os');
const path = require('node:path');
const http = require('node:http');
const sharp = require('sharp');
const { Chat, ChatMessage } = require('@lmstudio/sdk');
const { toolsProvider } = require('../dist/toolsProvider');
const { recentCandidates } = require('../dist/attachmentIdentity');
const { preprocess } = require('../dist/promptPreprocessor');

(async () => {
  const root = await fs.mkdtemp(path.join(os.tmpdir(), 'aia-multi-'));
  const red = await sharp({create: {width: 24, height: 24, channels: 3, background: 'red'}}).png().toBuffer();
  const blue = await sharp({create: {width: 24, height: 24, channels: 3, background: 'blue'}}).webp().toBuffer();
  for (const [name, bytes] of [['red.png', red], ['blue.webp', blue], ['red-copy.png', red], ['no-meta.png', blue]]) {
    await fs.writeFile(path.join(root, name), bytes);
    if (name !== 'no-meta.png') await fs.writeFile(path.join(root, name.replace(/\.[^.]+$/, '.json')), JSON.stringify({caption: name, sourcePageUrl: 'https://example.invalid/' + name}));
  }
  const settings = {analysisBackend: 'local', strataModel: 'strata', autoAttach: false, carryAnalyzed: true, maxImages: 2, startImage: 1, pageUrl: ''};
  let modelCalls = 0, localSelections = 0;
  const client = {files: {}, llm: {model: async () => {
    localSelections++;
    return {vision: true, identifier: 'test-vlm', respond: async (messages, options) => {
      modelCalls++;
      assert.equal(messages.length, 1);
      assert.equal(messages[0].images.length, 2);
      assert(messages[0].content.indexOf('Image 1: blue.webp') < messages[0].content.indexOf('Image 2: red.png'));
      assert(messages[0].content.includes('https://example.invalid/blue.webp'));
      assert.equal(options.maxTokens, 2048);
      for (const image of messages[0].images) assert.equal((await sharp(Buffer.from(image.identifier.slice(7), 'base64')).metadata()).format, 'png');
      return {content: 'Both images compared together.'};
    }};
  }}};
  let serverUrl = '';
  const ctl = {client, getWorkingDirectory: () => root,
    getPluginConfig: () => ({get: key => settings[key]}),
    getGlobalPluginConfig: () => ({get: key => ({strataBaseUrl: serverUrl, strataApiKey: 'test-key'})[key]}),
  };
  const ctx = {status() {}, warn() {}, signal: new AbortController().signal, callId: 1};
  const multi = (await toolsProvider(ctl))[3];
  const args = {imageNames: ['blue.webp', 'red.png', 'red-copy.png'], prompt: 'Compare colors and appearance.'};
  const output = await multi.implementation(args, ctx);
  assert.equal(modelCalls, 1);
  assert.equal(output.analysis, 'Both images compared together.');
  assert.equal(output.duplicatesRemoved, 1);
  assert.deepEqual(output.images.map(i => i.imageName), ['blue.webp', 'red.png']);
  assert.deepEqual(output.images[1].requestedNames, ['red.png', 'red-copy.png']);
  assert.equal(output._aiaImageCandidates.length, 2);
  const history = Chat.from({messages: [
    {role: 'assistant', content: [{type: 'toolCallRequest', toolCallRequest: {id: 'multi-1', type: 'function', name: 'Multi_Analyze_Local_Images_Advanced', arguments: args}}]},
    {role: 'tool', content: [{type: 'toolCallResult', toolCallId: 'multi-1', content: JSON.stringify(output)}]},
  ]});
  assert.deepEqual(recentCandidates(history).map(c => c.relativePath), ['blue.webp', 'red.png']);
  const preCtl = {...ctl, abortSignal: ctx.signal, guardAbort() {}, debug() {}, pullHistory: async () => history,
    createStatus: () => ({setState() {}, setText() {}})};
  settings.maxImages = 1;
  const first = await preprocess(preCtl, ChatMessage.from('Continue.'));
  assert.equal(first.getFiles(client).length, 1);
  assert(first.getText().includes('blue.webp'));
  history.append(first);
  const second = await preprocess(preCtl, ChatMessage.from('Continue again.'));
  assert.equal(second.getFiles(client).length, 1);
  assert(second.getText().includes('red.png'));
  history.append(second);
  assert.equal((await preprocess(preCtl, ChatMessage.from('Done.'))).getFiles(client).length, 0);
  for (const [names, expected] of [
    [['red.png', 'red-copy.png'], 'not_enough_distinct_images'],
    [['red.png', 'no-meta.png'], 'image_identity_not_verified'],
    [['red.png', '../outside.png'], 'invalid_image_path'],
    [['red.png', 'absent.png'], 'image_not_found'],
  ]) assert.equal((await multi.implementation({...args, imageNames: names}, ctx)).error, expected);
  assert.throws(() => multi.implementation({...args, imageNames: ['red.png']}, ctx));
  assert.throws(() => multi.implementation({...args, imageNames: Array(5).fill('red.png')}, ctx));
  for (let i = 0; i < 3; i++) {
    await fs.writeFile(path.join(root, `large-${i}.png`), Buffer.concat([red, Buffer.alloc(6 * 1024 * 1024, i)]));
    await fs.writeFile(path.join(root, `large-${i}.json`), '{}');
  }
  assert.equal((await multi.implementation({...args, imageNames: ['large-0.png', 'large-1.png', 'large-2.png']}, ctx)).error, 'image_size_limit');
  assert.equal(modelCalls, 1, 'Invalid batches must not run partial analysis.');
  const aborted = new AbortController(); aborted.abort();
  await assert.rejects(() => multi.implementation(args, {...ctx, signal: aborted.signal}), /abort/i);
  let apiCalls = 0;
  let respond;
  const pending = new Promise(resolve => {respond = resolve;});
  const web = http.createServer(async (req, res) => {
    let raw = ''; for await (const chunk of req) raw += chunk;
    apiCalls++;
    const body = JSON.parse(raw);
    assert.equal(req.headers.authorization, 'Bearer test-key');
    assert.equal(body.messages.length, 1);
    assert.equal(body.messages[0].content.length, 3);
    assert.equal(body.messages[0].content[1].image_url.url, `data:image/webp;base64,${blue.toString('base64')}`);
    assert.equal(body.messages[0].content[2].image_url.url, `data:image/png;base64,${red.toString('base64')}`);
    assert.equal(body.stream, false); assert.equal(body.tools, undefined);
    respond(() => {res.setHeader('Content-Type', 'application/json'); res.end(JSON.stringify({choices: [{message: {content: 'Joint Strata comparison.'}}]}));});
  });
  await new Promise(resolve => web.listen(0, '127.0.0.1', resolve));
  try {
    serverUrl = `http://127.0.0.1:${web.address().port}/v1`; settings.analysisBackend = 'strata';
    const before = localSelections;
    let settled = false;
    const request = multi.implementation(args, ctx).then(value => {settled = true; return value;});
    const complete = await pending;
    assert.equal(settled, false, 'Tool must wait for the joint Strata response.');
    complete();
    assert.equal((await request).analysis, 'Joint Strata comparison.');
    assert.equal(apiCalls, 1); assert.equal(localSelections, before);
  } finally {web.closeAllConnections(); await new Promise(resolve => web.close(resolve));}
  console.log('PASS: one joint inference for local and Strata backends, ordered mapping, byte dedup, WebP compatibility, batch handoff across turns, no partial inference on invalid/oversized batches, cancellation and Strata wait semantics.');
  console.log('Local model/controller mocked; real SDK messages, image conversion, files and Strata-format HTTP are exercised.');
})().catch(error => {console.error(error); process.exitCode = 1;});

const assert = require('node:assert/strict');
const fs = require('node:fs/promises');
const os = require('node:os');
const path = require('node:path');
const http = require('node:http');
const sharp = require('sharp');
const { Chat, ChatMessage } = require('@lmstudio/sdk');
const { toolsProvider } = require('../dist/toolsProvider');
const { recentCandidates } = require('../dist/attachmentIdentity');
const { preprocess } = require('../dist/promptPreprocessor');

(async () => {
  const root = await fs.mkdtemp(path.join(os.tmpdir(), 'aia-multi-'));
  const red = await sharp({create: {width: 24, height: 24, channels: 3, background: 'red'}}).png().toBuffer();
  const blue = await sharp({create: {width: 24, height: 24, channels: 3, background: 'blue'}}).webp().toBuffer();
  for (const [name, bytes] of [['red.png', red], ['blue.webp', blue], ['red-copy.png', red], ['no-meta.png', blue]]) {
    await fs.writeFile(path.join(root, name), bytes);
    if (name !== 'no-meta.png') await fs.writeFile(path.join(root, name.replace(/\.[^.]+$/, '.json')), JSON.stringify({caption: name, sourcePageUrl: 'https://example.invalid/' + name}));
  }
  const settings = {analysisBackend: 'local', strataModel: 'strata', autoAttach: false, carryAnalyzed: true, maxImages: 2, startImage: 1, pageUrl: ''};
  let modelCalls = 0, localSelections = 0;
  const client = {files: {}, llm: {model: async () => {
    localSelections++;
    return {vision: true, identifier: 'test-vlm', respond: async (messages, options) => {
      modelCalls++;
      assert.equal(messages.length, 1);
      assert.equal(messages[0].images.length, 2);
      assert(messages[0].content.indexOf('Image 1: blue.webp') < messages[0].content.indexOf('Image 2: red.png'));
      assert(messages[0].content.includes('https://example.invalid/blue.webp'));
      assert.equal(options.maxTokens, 2048);
      for (const image of messages[0].images) assert.equal((await sharp(Buffer.from(image.identifier.slice(7), 'base64')).metadata()).format, 'png');
      return {content: 'Both images compared together.'};
    }};
  }}};
  let serverUrl = '';
  const ctl = {client, getWorkingDirectory: () => root,
    getPluginConfig: () => ({get: key => settings[key]}),
    getGlobalPluginConfig: () => ({get: key => ({strataBaseUrl: serverUrl, strataApiKey: 'test-key'})[key]}),
  };
  const ctx = {status() {}, warn() {}, signal: new AbortController().signal, callId: 1};
  const multi = (await toolsProvider(ctl))[3];
  const args = {imageNames: ['blue.webp', 'red.png', 'red-copy.png'], prompt: 'Compare colors and appearance.'};
  const output = await multi.implementation(args, ctx);
  assert.equal(modelCalls, 1);
  assert.equal(output.analysis, 'Both images compared together.');
  assert.equal(output.duplicatesRemoved, 1);
  assert.deepEqual(output.images.map(i => i.imageName), ['blue.webp', 'red.png']);
  assert.deepEqual(output.images[1].requestedNames, ['red.png', 'red-copy.png']);
  assert.equal(output._aiaImageCandidates.length, 2);
  const history = Chat.from({messages: [
    {role: 'assistant', content: [{type: 'toolCallRequest', toolCallRequest: {id: 'multi-1', type: 'function', name: 'Multi_Analyze_Local_Images_Advanced', arguments: args}}]},
    {role: 'tool', content: [{type: 'toolCallResult', toolCallId: 'multi-1', content: JSON.stringify(output)}]},
  ]});
  assert.deepEqual(recentCandidates(history).map(c => c.relativePath), ['blue.webp', 'red.png']);
  const preCtl = {...ctl, abortSignal: ctx.signal, guardAbort() {}, debug() {}, pullHistory: async () => history,
    createStatus: () => ({setState() {}, setText() {}})};
  settings.maxImages = 1;
  const first = await preprocess(preCtl, ChatMessage.from('Continue.'));
  assert.equal(first.getFiles(client).length, 1);
  assert(first.getText().includes('blue.webp'));
  history.append(first);
  const second = await preprocess(preCtl, ChatMessage.from('Continue again.'));
  assert.equal(second.getFiles(client).length, 1);
  assert(second.getText().includes('red.png'));
  history.append(second);
  assert.equal((await preprocess(preCtl, ChatMessage.from('Done.'))).getFiles(client).length, 0);
  for (const [names, expected] of [
    [['red.png', 'red-copy.png'], 'not_enough_distinct_images'],
    [['red.png', 'no-meta.png'], 'image_identity_not_verified'],
    [['red.png', '../outside.png'], 'invalid_image_path'],
    [['red.png', 'absent.png'], 'image_not_found'],
  ]) assert.equal((await multi.implementation({...args, imageNames: names}, ctx)).error, expected);
  assert.throws(() => multi.implementation({...args, imageNames: ['red.png']}, ctx));
  assert.throws(() => multi.implementation({...args, imageNames: Array(5).fill('red.png')}, ctx));
  for (let i = 0; i < 3; i++) {
    await fs.writeFile(path.join(root, `large-${i}.png`), Buffer.concat([red, Buffer.alloc(6 * 1024 * 1024, i)]));
    await fs.writeFile(path.join(root, `large-${i}.json`), '{}');
  }
  assert.equal((await multi.implementation({...args, imageNames: ['large-0.png', 'large-1.png', 'large-2.png']}, ctx)).error, 'image_size_limit');
  assert.equal(modelCalls, 1, 'Invalid batches must not run partial analysis.');
  const aborted = new AbortController(); aborted.abort();
  await assert.rejects(() => multi.implementation(args, {...ctx, signal: aborted.signal}), /abort/i);
  let apiCalls = 0;
  let respond;
  const pending = new Promise(resolve => {respond = resolve;});
  const web = http.createServer(async (req, res) => {
    let raw = ''; for await (const chunk of req) raw += chunk;
    apiCalls++;
    const body = JSON.parse(raw);
    assert.equal(req.headers.authorization, 'Bearer test-key');
    assert.equal(body.messages.length, 1);
    assert.equal(body.messages[0].content.length, 3);
    assert.equal(body.messages[0].content[1].image_url.url, `data:image/webp;base64,${blue.toString('base64')}`);
    assert.equal(body.messages[0].content[2].image_url.url, `data:image/png;base64,${red.toString('base64')}`);
    assert.equal(body.stream, false); assert.equal(body.tools, undefined);
    respond(() => {res.setHeader('Content-Type', 'application/json'); res.end(JSON.stringify({choices: [{message: {content: 'Joint Strata comparison.'}}]}));});
  });
  await new Promise(resolve => web.listen(0, '127.0.0.1', resolve));
  try {
    serverUrl = `http://127.0.0.1:${web.address().port}/v1`; settings.analysisBackend = 'strata';
    const before = localSelections;
    let settled = false;
    const request = multi.implementation(args, ctx).then(value => {settled = true; return value;});
    const complete = await pending;
    assert.equal(settled, false, 'Tool must wait for the joint Strata response.');
    complete();
    assert.equal((await request).analysis, 'Joint Strata comparison.');
    assert.equal(apiCalls, 1); assert.equal(localSelections, before);
  } finally {web.closeAllConnections(); await new Promise(resolve => web.close(resolve));}
  console.log('PASS: one joint inference for local and Strata backends, ordered mapping, byte dedup, WebP compatibility, batch handoff across turns, no partial inference on invalid/oversized batches, cancellation and Strata wait semantics.');
  console.log('Local model/controller mocked; real SDK messages, image conversion, files and Strata-format HTTP are exercised.');
})().catch(error => {console.error(error); process.exitCode = 1;});