import { describe, it, beforeEach, afterEach, mock } from 'node:test';
import assert from 'node:assert/strict';
import { mkdtemp, writeFile, chmod, rm, readFile } from 'fs/promises';
import { tmpdir } from 'os';
import { join } from 'path';
import { extractBalancedJSON, extractJSONText, extractMeaningfulError, mapWithConcurrency } from '../llm.js';

const ENV_KEYS = [
  'AI_ASSISTANT_CLI', 'MODEL', 'OPENCODE_MODEL', 'PATH', 'AI_ASSISTANT_TIMEOUT_MS', 'AI_ASSISTANT_MAX_BUFFER',
  'FAKE_AI_STDOUT', 'FAKE_AI_STDERR', 'FAKE_AI_EXIT', 'FAKE_AI_STDIN_PATH', 'FAKE_AI_ARGS_PATH',
];

let saved = {};
let tempDir;
beforeEach(() => {
  saved = {};
  for (const k of ENV_KEYS) { saved[k] = process.env[k]; delete process.env[k]; }
  tempDir = null;
});
afterEach(async () => {
  for (const k of ENV_KEYS) {
    if (saved[k] === undefined) delete process.env[k];
    else process.env[k] = saved[k];
  }
  if (tempDir) await rm(tempDir, { recursive: true, force: true });
  mock.restoreAll();
});

async function installFakeCLI(command = 'codex') {
  tempDir = await mkdtemp(join(tmpdir(), 'ai-cli-test-'));
  const script = join(tempDir, command);
  await writeFile(script, `#!/bin/sh
if [ -n "$FAKE_AI_ARGS_PATH" ]; then printf '%s\\n' "$*" > "$FAKE_AI_ARGS_PATH"; fi
if [ -n "$FAKE_AI_STDIN_PATH" ]; then /bin/cat > "$FAKE_AI_STDIN_PATH"; else /bin/cat >/dev/null; fi
if [ -n "$FAKE_AI_STDERR" ]; then printf '%s' "$FAKE_AI_STDERR" >&2; fi
if [ -n "$FAKE_AI_STDOUT" ]; then printf '%s' "$FAKE_AI_STDOUT"; fi
exit "\${FAKE_AI_EXIT:-0}"
`);
  await chmod(script, 0o755);
  process.env.PATH = tempDir;
  return { stdinPath: join(tempDir, 'stdin.txt'), argsPath: join(tempDir, 'args.txt') };
}

describe('chat - assistant CLI', async () => {
  const { chat } = await import('../llm.js');

  it('runs the detected CLI with MODEL and sends the prompts through stdin', async () => {
    const { stdinPath, argsPath } = await installFakeCLI('codex');
    process.env.MODEL = 'gpt-5-mini';
    process.env.FAKE_AI_STDOUT = 'cli response';
    process.env.FAKE_AI_STDIN_PATH = stdinPath;
    process.env.FAKE_AI_ARGS_PATH = argsPath;

    const result = await chat('sys', 'user');

    assert.equal(result, 'cli response');
    assert.match(await readFile(argsPath, 'utf8'), /exec --model gpt-5-mini/);
    const prompt = await readFile(stdinPath, 'utf8');
    assert.match(prompt, /<system>\nsys\n<\/system>/);
    assert.match(prompt, /<user>\nuser\n<\/user>/);
  });

  it('can force opencode with AI_ASSISTANT_CLI', async () => {
    const { argsPath } = await installFakeCLI('opencode');
    process.env.AI_ASSISTANT_CLI = 'opencode';
    process.env.MODEL = 'google/gemini-2.5-pro';
    process.env.FAKE_AI_STDOUT = 'ok';
    process.env.FAKE_AI_ARGS_PATH = argsPath;

    const result = await chat('sys', 'user');

    assert.equal(result, 'ok');
    assert.match(await readFile(argsPath, 'utf8'), /run --model google\/gemini-2.5-pro/);
  });

  it('runs Antigravity through agy with MODEL and prompt argument', async () => {
    const { stdinPath, argsPath } = await installFakeCLI('agy');
    process.env.AI_ASSISTANT_CLI = 'agy';
    process.env.MODEL = 'gemini-2.5-pro';
    process.env.FAKE_AI_STDOUT = 'antigravity response';
    process.env.FAKE_AI_STDIN_PATH = stdinPath;
    process.env.FAKE_AI_ARGS_PATH = argsPath;

    const result = await chat('sys', 'user');

    assert.equal(result, 'antigravity response');
    const args = await readFile(argsPath, 'utf8');
    assert.match(args, /-p .*--model gemini-2.5-pro/s);
    assert.match(args, /<system>\nsys\n<\/system>/);
    assert.equal(await readFile(stdinPath, 'utf8'), '');
  });

  it('throws an error when the CLI fails instead of exiting the process', async () => {
    await installFakeCLI('codex');
    process.env.FAKE_AI_EXIT = '2';
    process.env.FAKE_AI_STDERR = 'provider overloaded';
    const exitMock = mock.method(process, 'exit', () => { throw new Error('exit should not be called'); });

    await assert.rejects(() => chat('sys', 'user'), /provider overloaded/);

    assert.equal(exitMock.mock.calls.length, 0);
  });
});

describe('chatJSON', async () => {
  const { chatJSON } = await import('../llm.js');

  it('parses plain JSON response', async () => {
    await installFakeCLI('codex');
    process.env.FAKE_AI_STDOUT = '[{"level":"critical"}]';

    const result = await chatJSON('sys', 'user');

    assert.deepEqual(result, [{ level: 'critical' }]);
  });

  it('strips markdown code block before parsing', async () => {
    await installFakeCLI('codex');
    process.env.FAKE_AI_STDOUT = '```json\n[{"level":"info"}]\n```';

    const result = await chatJSON('sys', 'user');

    assert.deepEqual(result, [{ level: 'info' }]);
  });

  it('extracts JSON array from surrounding prose', async () => {
    await installFakeCLI('codex');
    process.env.FAKE_AI_STDOUT = '**Reviewing findings**\n\n[{"level":"warning","suggestion":"x"}]\n\nDone.';

    const result = await chatJSON('sys', 'user');

    assert.deepEqual(result, [{ level: 'warning', suggestion: 'x' }]);
  });

  it('extracts JSON object from surrounding prose', async () => {
    await installFakeCLI('codex');
    process.env.FAKE_AI_STDOUT = '**Begin Combine**\n{"merged_text":"repo block\\n\\nsource block"}';

    const result = await chatJSON('sys', 'user');

    assert.deepEqual(result, { merged_text: 'repo block\n\nsource block' });
  });

  it('returns [] when JSON is invalid', async () => {
    await installFakeCLI('codex');
    process.env.FAKE_AI_STDOUT = 'not json';

    const result = await chatJSON('sys', 'user');

    assert.deepEqual(result, []);
  });
});

describe('extractBalancedJSON', () => {
  it('returns the whole object for a simple object from index 0', () => {
    const text = '{"a":1}';

    assert.equal(extractBalancedJSON(text, 0), '{"a":1}');
  });

  it('returns the full balanced segment for deeply nested object/array', () => {
    const text = '{"a":[1,{"b":[2,{"c":3}]}],"d":4}';

    assert.equal(extractBalancedJSON(text, 0), '{"a":[1,{"b":[2,{"c":3}]}],"d":4}');
  });

  it('does not let braces inside a string value break balancing', () => {
    const text = '{"a":"}{"}';

    assert.equal(extractBalancedJSON(text, 0), '{"a":"}{"}');
  });

  it('handles an escaped quote inside a string value', () => {
    const text = '{"a":"\\""}';

    assert.equal(extractBalancedJSON(text, 0), '{"a":"\\""}');
  });

  it('returns null for truncated/incomplete JSON', () => {
    const text = '{"a":1';

    assert.equal(extractBalancedJSON(text, 0), null);
  });

  it('extracts a balanced array when starting at a "["', () => {
    const text = '[1,[2,3],{"a":4}]';

    assert.equal(extractBalancedJSON(text, 0), '[1,[2,3],{"a":4}]');
  });

  it('excludes trailing content after the balanced segment', () => {
    const text = '{"a":1} trailing text {"b":2}';

    assert.equal(extractBalancedJSON(text, 0), '{"a":1}');
  });
});

describe('extractJSONText', () => {
  it('strips a fenced ```json block', () => {
    const text = '```json\n{"a":1}\n```';

    const result = extractJSONText(text);

    assert.deepEqual(JSON.parse(result), { a: 1 });
  });

  it('extracts a JSON object after leading prose', () => {
    const text = 'Here are the findings:\n{"level":"critical"}';

    const result = extractJSONText(text);

    assert.deepEqual(JSON.parse(result), { level: 'critical' });
  });

  it('extracts an array embedded in surrounding text', () => {
    const text = 'prefix [1,2,3] suffix';

    const result = extractJSONText(text);

    assert.deepEqual(JSON.parse(result), [1, 2, 3]);
  });

  it('returns an already-pure JSON string as-is', () => {
    const text = '{"a":1,"b":[2,3]}';

    const result = extractJSONText(text);

    assert.equal(result, '{"a":1,"b":[2,3]}');
    assert.deepEqual(JSON.parse(result), { a: 1, b: [2, 3] });
  });

  it('returns the de-fenced original text when no valid JSON is found', () => {
    const text = '```\nnot json at all\n```';

    const result = extractJSONText(text);

    assert.equal(result, 'not json at all');
  });
});

describe('extractMeaningfulError', () => {
  it('抽出尾端真正的錯誤，而非開頭的 codex banner／回顯 prompt', () => {
    const raw = [
      'OpenAI Codex v0.142.5',
      '--------',
      'workdir: /workspace/actions/ai-code-review',
      'model: gpt-5.4-mini',
      'reasoning effort: none',
      '--------',
      'user',
      '請依照以下系統指示處理使用者內容，並只輸出要求的最終結果。',
      'ERROR codex_api::endpoint::responses_websocket: failed to connect to websocket: HTTP error: 401 Unauthorized',
      'ERROR: Your access token could not be refreshed because your refresh token was revoked. Please log out and sign in again.',
    ].join('\n');

    const result = extractMeaningfulError(raw);

    assert.match(result, /401 Unauthorized/);
    assert.match(result, /refresh token was revoked/);
    assert.doesNotMatch(result, /workdir:/);
    assert.doesNotMatch(result, /請依照以下系統指示/);
  });

  it('抽不到錯誤行時退取尾段（不取開頭）', () => {
    const raw = 'A'.repeat(1200) + '\nTAIL-CONTENT';
    const result = extractMeaningfulError(raw, 100);
    assert.ok(result.length <= 100);
    assert.match(result, /TAIL-CONTENT$/);
  });

  it('容錯處理空輸入', () => {
    assert.equal(extractMeaningfulError(''), '');
    assert.equal(extractMeaningfulError(null), '');
  });
});

describe('mapWithConcurrency', () => {
  it('回傳與輸入同索引對應的結果（保序）', async () => {
    const out = await mapWithConcurrency([1, 2, 3, 4], 2, async (n) => n * 10);
    assert.deepEqual(out, [10, 20, 30, 40]);
  });

  it('遵守併發上限（同時執行數不超過 limit）', async () => {
    let active = 0, peak = 0;
    const wait = () => new Promise(r => setTimeout(r, 5));
    await mapWithConcurrency([1, 2, 3, 4, 5, 6], 2, async () => {
      active += 1; peak = Math.max(peak, active);
      await wait();
      active -= 1;
    });
    assert.ok(peak <= 2, `peak=${peak} 應 <= 2`);
  });

  it('limit 大於項目數時仍全部執行', async () => {
    const out = await mapWithConcurrency(['a', 'b'], 10, async (s) => s.toUpperCase());
    assert.deepEqual(out, ['A', 'B']);
  });

  it('limit<=0 表示不限制（全部同時並行）', async () => {
    let active = 0, peak = 0;
    const wait = () => new Promise(r => setTimeout(r, 5));
    await mapWithConcurrency([1, 2, 3, 4, 5], 0, async () => {
      active += 1; peak = Math.max(peak, active);
      await wait();
      active -= 1;
    });
    assert.equal(peak, 5, `peak=${peak} 應等於項目數（不限制）`);
  });

  it('空輸入回傳空陣列', async () => {
    assert.deepEqual(await mapWithConcurrency([], 3, async () => 1), []);
    assert.deepEqual(await mapWithConcurrency(null, 3, async () => 1), []);
  });
});
