#!/usr/bin/env node
/**
 * Stub-on-record A/B harness.
 *
 * Loops a synthetic agent against the Loxia backend over multiple turns,
 * with one of two variants for how PRIOR-turn code-map results are
 * rendered into the LLM call:
 *
 *   A (status-quo):  full content preserved in every prior tool-result.
 *   B (stubbed):     prior code-map results replaced with a stub. Current
 *                    turn's content preserved.
 *
 * Each scenario instructs the agent to re-examine the same fixture file
 * across multiple turns. Variant B should produce smaller prompt-token
 * counts without harming the agent's ability to keep working.
 *
 * Outputs:
 *   - JSON to scripts/stub-bench/results-<ts>.json
 *   - human summary table to stdout
 *
 * Usage:
 *   node scripts/stub-bench/bench.mjs
 *   node scripts/stub-bench/bench.mjs --models gpt-4.1-mini --trials 3
 */
import { readFileSync } from 'node:fs';
import { mkdir, writeFile, readFile } from 'node:fs/promises';
import { homedir } from 'node:os';
import { dirname, resolve, join } from 'node:path';
import { fileURLToPath } from 'node:url';

import { applyCodeMapStubbing } from '../../src/services/codeMapStubRenderer.js';

const __dirname = dirname(fileURLToPath(import.meta.url));
const REPO_ROOT = resolve(__dirname, '..', '..');
const BASE_URL = process.env.LOXIA_BASE_URL ||
  'https://loxia-api-g7hrb8bxdae8a2h7.z02.azurefd.net';

const DEFAULT_MODELS = ['gpt-4.1-nano', 'gpt-4.1-mini'];
const DEFAULT_TRIALS = 2;
const MAX_TURNS = 6;

function parseArgs(argv) {
  const args = {
    models: DEFAULT_MODELS, trials: DEFAULT_TRIALS,
    variants: ['A', 'B'], scenarios: null, out: null,
  };
  for (let i = 2; i < argv.length; i += 2) {
    const f = argv[i], v = argv[i + 1];
    switch (f) {
      case '--models':    args.models = v.split(',').map(s => s.trim()); break;
      case '--trials':    args.trials = parseInt(v, 10); break;
      case '--variants':  args.variants = v.split(',').map(s => s.trim()); break;
      case '--scenarios': args.scenarios = v.split(',').map(s => s.trim()); break;
      case '--out':       args.out = v; break;
      default: console.error(`unknown flag: ${f}`); process.exit(1);
    }
  }
  return args;
}

function loadApiKey() {
  if (process.env.LOXIA_API_KEY) return process.env.LOXIA_API_KEY;
  try {
    const p = join(homedir(), '.loxia', 'keys.json');
    const j = JSON.parse(readFileSync(p, 'utf8'));
    return j?.backend?.key || null;
  } catch { return null; }
}

// ─── Fixture loading ──────────────────────────────────────────────

async function loadFixture() {
  const path = join(__dirname, 'fixture', 'auth.js');
  const content = await readFile(path, 'utf8');
  return { path: 'fixture/auth.js', content };
}

// ─── Tool implementations (just for code-map skeleton/read-range) ──

function fakeSkeleton(content) {
  // Cheap pseudo-skeleton: keep class/function declarations + braces.
  const lines = content.split('\n');
  const sig = [];
  for (let i = 0; i < lines.length; i++) {
    const l = lines[i].trim();
    if (l.startsWith('class ') || l.startsWith('function ') ||
        l.startsWith('async ') || l.startsWith('//') ||
        l.includes('module.exports') ||
        (l.endsWith('{') && !l.startsWith('if ') && !l.startsWith('for '))) {
      sig.push(`${i + 1}: ${lines[i]}`);
    }
  }
  return sig.join('\n');
}

function readRange(content, startLine, endLine) {
  const lines = content.split('\n');
  const a = Math.max(0, startLine - 1);
  const b = Math.min(lines.length, endLine);
  return lines.slice(a, b).map((l, i) => `${a + 1 + i}: ${l}`).join('\n');
}

// ─── LLM call ─────────────────────────────────────────────────────

const TOOLS = [{
  type: 'function',
  function: {
    name: 'code-map',
    description: 'Read code structure. skeleton: top-level signatures. read-range: specific line range.',
    parameters: {
      type: 'object',
      properties: {
        action:    { type: 'string', enum: ['skeleton', 'read-range'] },
        filePath:  { type: 'string' },
        startLine: { type: 'number' },
        endLine:   { type: 'number' },
      },
      required: ['action', 'filePath'],
    },
  },
}, {
  type: 'function',
  function: {
    name: 'job-done',
    description: 'Signal that you have finished the task.',
    parameters: {
      type: 'object',
      properties: { summary: { type: 'string' } },
      required: ['summary'],
    },
  },
}];

const SYSTEM_PROMPT_BASE =
  'You are an agent inspecting a small codebase. Use the `code-map` tool to look at structure ' +
  'and line ranges. When you have answered the user, call `job-done`. Be brief — one tool ' +
  'call per turn.';

function extractToolCalls(content) {
  if (typeof content !== 'string') return [];
  const calls = [];
  let pos = 0;
  while (pos < content.length) {
    const markerIdx = content.indexOf('**Calling ', pos);
    if (markerIdx === -1) break;
    const nameStart = markerIdx + '**Calling '.length;
    const nameEnd = content.indexOf('**', nameStart);
    if (nameEnd === -1) break;
    const toolName = content.slice(nameStart, nameEnd);
    const fenceStart = content.indexOf('```', nameEnd);
    if (fenceStart === -1) break;
    const bodyStart = content.indexOf('\n', fenceStart) + 1;
    const fenceEnd = content.indexOf('```', bodyStart);
    if (fenceEnd === -1) break;
    let body = content.slice(bodyStart, fenceEnd).trim();
    if (body.startsWith('json')) body = body.slice(4).trim();
    try {
      const parsed = JSON.parse(body);
      // Plural-canonical (e7d496c): tools may emit `{actions:[{type,...}]}`
      // or legacy `{action,...}`. Expand actions[] into one call each so
      // downstream scoring sees every operation. The action field is
      // normalized to lowercase `action` (from `type`) on the way out
      // for back-compat with this bench's existing rubric.
      const params = parsed.parameters || parsed;
      const actions = params?.actions;
      if (Array.isArray(actions) && actions.length > 0) {
        for (const a of actions) {
          calls.push({
            toolId: parsed.toolId || toolName,
            parameters: { ...a, action: a.type || a.action },
          });
        }
      } else {
        calls.push({
          toolId: parsed.toolId || toolName,
          parameters: params,
        });
      }
    } catch {
      calls.push({ toolId: toolName, parameters: { _unparsed: body.slice(0, 200) } });
    }
    pos = fenceEnd + 3;
  }
  return calls;
}

async function callLLM({ apiKey, model, messages, systemPrompt }) {
  const requestId = `stb-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
  const last = messages[messages.length - 1];
  const payload = {
    conversationId: requestId,
    message: typeof last?.content === 'string' ? last.content : '',
    messages,
    model, requestId,
    options: { max_tokens: 700, temperature: 0.2, tools: TOOLS },
    platformProvided: true,
    systemPrompt,
    apiKey,
  };
  const t0 = Date.now();
  const res = await fetch(`${BASE_URL}/llm/chat`, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'Authorization': `Bearer ${apiKey}`,
      'X-Request-ID': requestId,
    },
    body: JSON.stringify(payload),
  });
  const latencyMs = Date.now() - t0;
  if (!res.ok) {
    const txt = await res.text();
    return { error: `HTTP ${res.status}: ${txt.slice(0, 200)}`, latencyMs };
  }
  const data = await res.json();
  return {
    content: data.content || '',
    finishReason: data.finishReason || 'stop',
    usage: data.usage || {},
    latencyMs,
  };
}

// ─── one session loop ─────────────────────────────────────────────

async function runSession({ apiKey, model, scenario, fixture, variant }) {
  const systemPrompt = SYSTEM_PROMPT_BASE;
  let messages = [{ role: 'user', content: scenario.userMessage }];
  const trace = {
    turns: [], finished: false, toolCalls: 0, codeMapCalls: 0,
    totalPromptTokens: 0, totalOutputTokens: 0, totalLatencyMs: 0,
    error: null,
  };

  for (let turn = 0; turn < MAX_TURNS; turn++) {
    // Render: apply stubbing for variant B (preserves the last
    // tool-result, stubs all prior code-map ones).
    const renderedMessages = variant === 'B'
      ? applyCodeMapStubbing(messages)
      : messages;

    const r = await callLLM({ apiKey, model, messages: renderedMessages, systemPrompt });
    if (r.error) {
      trace.error = r.error;
      break;
    }
    trace.totalPromptTokens += r.usage.input_tokens || r.usage.prompt_tokens || 0;
    trace.totalOutputTokens += r.usage.output_tokens || r.usage.completion_tokens || 0;
    trace.totalLatencyMs += r.latencyMs;

    const calls = extractToolCalls(r.content);
    trace.turns.push({
      turn, finishReason: r.finishReason, toolCalls: calls.length,
      preview: r.content.slice(0, 160),
    });

    // Append assistant message to history.
    messages.push({ role: 'assistant', content: r.content });

    if (calls.length === 0) {
      // No tool call — agent gave up or answered directly.
      break;
    }
    trace.toolCalls += calls.length;

    // Process each tool call → append tool-result message.
    let jobDone = false;
    const toolResultParts = [];
    for (const call of calls) {
      if (call.toolId === 'job-done') {
        trace.finished = true;
        jobDone = true;
        break;
      }
      if (call.toolId === 'code-map') {
        trace.codeMapCalls++;
        const action = call.parameters?.action;
        let body = '';
        if (action === 'skeleton') {
          body = `code-map skeleton:\n  filePath: ${fixture.path}\n${fakeSkeleton(fixture.content)}`;
        } else if (action === 'read-range') {
          const sl = Number(call.parameters?.startLine || 1);
          const el = Number(call.parameters?.endLine || sl);
          body = `code-map read-range:\n  filePath: ${fixture.path}\n  startLine: ${sl}\n  endLine: ${el}\n${readRange(fixture.content, sl, el)}`;
        } else {
          body = `code-map: unknown action ${action}`;
        }
        toolResultParts.push(body);
      } else {
        toolResultParts.push(`unknown tool: ${call.toolId}`);
      }
    }
    if (jobDone) break;

    if (toolResultParts.length > 0) {
      messages.push({
        role: 'user', type: 'tool-result',
        content: '[Tool Results — 1 batch: code-map]\n' + toolResultParts.join('\n---\n'),
        responseTurnId: `turn-${turn}`,
      });
    }
  }

  return trace;
}

// ─── scenarios ────────────────────────────────────────────────────

const SCENARIOS = [
  {
    id: 'S1-revisit-skeleton',
    description: 'User asks twice about the same file structure — Monopoly re-touch pattern.',
    userMessage:
      'Look at fixture/auth.js and tell me what methods AuthService exposes. ' +
      'After listing them, also tell me which one is responsible for revoking sessions, ' +
      'and which line range that method occupies.',
  },
  {
    id: 'S2-skeleton-then-ranges',
    description: 'User asks for skeleton, then explicit follow-up ranges.',
    userMessage:
      'In fixture/auth.js: (1) list all the methods of AuthService; ' +
      '(2) show me the body of login and verify; ' +
      '(3) tell me whether logout actually deletes from tokenStore.',
  },
  {
    id: 'S3-single-call',
    description: 'Single-call simple task — negative control. Stubbing should not hurt anything.',
    userMessage:
      'Look at fixture/auth.js and tell me the name of the class it exports.',
  },
];

// ─── runner ───────────────────────────────────────────────────────

async function runAll({ apiKey, models, variants, trials, scenarios }) {
  const fixture = await loadFixture();
  const scList = scenarios
    ? SCENARIOS.filter(s => scenarios.includes(s.id))
    : SCENARIOS;
  const results = [];
  const total = models.length * variants.length * scList.length * trials;
  let done = 0;
  for (const model of models) {
    for (const variant of variants) {
      for (const scenario of scList) {
        for (let t = 0; t < trials; t++) {
          const trace = await runSession({ apiKey, model, scenario, fixture, variant });
          done++;
          process.stderr.write(`\r${done}/${total} (model=${model} variant=${variant} sc=${scenario.id} trial=${t})  `);
          results.push({ model, variant, scenario: scenario.id, trial: t, ...trace });
        }
      }
    }
  }
  process.stderr.write('\n');
  return results;
}

// ─── reporting ────────────────────────────────────────────────────

function aggregate(results) {
  // Group by (model, scenario, variant).
  const byCell = new Map();
  for (const r of results) {
    const k = `${r.model}|${r.scenario}|${r.variant}`;
    if (!byCell.has(k)) byCell.set(k, []);
    byCell.get(k).push(r);
  }
  const agg = [];
  for (const [key, list] of byCell) {
    const [model, scenario, variant] = key.split('|');
    const avg = arr => arr.reduce((s, x) => s + x, 0) / arr.length;
    agg.push({
      model, scenario, variant, n: list.length,
      finishedPct:    100 * list.filter(r => r.finished).length / list.length,
      avgPromptTok:   Math.round(avg(list.map(r => r.totalPromptTokens))),
      avgOutputTok:   Math.round(avg(list.map(r => r.totalOutputTokens))),
      avgToolCalls:   +avg(list.map(r => r.toolCalls)).toFixed(2),
      avgCodeMapCalls:+avg(list.map(r => r.codeMapCalls)).toFixed(2),
      errors:         list.filter(r => r.error).length,
    });
  }
  return agg;
}

function printReport(agg, models, scenarios) {
  console.log('\n=== STUB-ON-RECORD A/B ===\n');
  // Pair A/B per (model, scenario) and compute deltas.
  const byPair = new Map();
  for (const a of agg) {
    const k = `${a.model}|${a.scenario}`;
    if (!byPair.has(k)) byPair.set(k, {});
    byPair.get(k)[a.variant] = a;
  }
  console.log('Model        | Scenario             | A=tokens | B=tokens | Δ%   | A=calls | B=calls | A=done | B=done');
  console.log('-------------+----------------------+----------+----------+------+---------+---------+--------+-------');
  for (const [key, pair] of byPair) {
    const [model, scenario] = key.split('|');
    const a = pair.A, b = pair.B;
    if (!a || !b) continue;
    const dPct = a.avgPromptTok === 0 ? 0 : Math.round(100 * (b.avgPromptTok - a.avgPromptTok) / a.avgPromptTok);
    const sign = dPct < 0 ? '' : '+';
    console.log(
      `${model.padEnd(12)} | ${scenario.padEnd(20)} | ` +
      `${String(a.avgPromptTok).padStart(8)} | ${String(b.avgPromptTok).padStart(8)} | ` +
      `${(sign + dPct + '%').padStart(4)} | ` +
      `${String(a.avgToolCalls).padStart(7)} | ${String(b.avgToolCalls).padStart(7)} | ` +
      `${String(Math.round(a.finishedPct) + '%').padStart(6)} | ${String(Math.round(b.finishedPct) + '%').padStart(6)}`
    );
  }
}

// ─── main ─────────────────────────────────────────────────────────

(async () => {
  const args = parseArgs(process.argv);
  const apiKey = loadApiKey();
  if (!apiKey) { console.error('No LOXIA_API_KEY.'); process.exit(1); }

  console.error(`stub-bench: models=${args.models.join(',')} variants=${args.variants.join(',')} trials=${args.trials}`);
  const results = await runAll({
    apiKey, models: args.models, variants: args.variants,
    trials: args.trials, scenarios: args.scenarios,
  });
  const agg = aggregate(results);
  printReport(agg, args.models, args.scenarios);

  const ts = new Date().toISOString().replace(/[:.]/g, '-');
  const outPath = args.out || join(__dirname, `results-${ts}.json`);
  await mkdir(dirname(outPath), { recursive: true });
  await writeFile(outPath, JSON.stringify({
    timestamp: new Date().toISOString(),
    args,
    aggregate: agg,
    results,
  }, null, 2));
  console.log(`\nfull results → ${outPath}`);
})().catch(err => { console.error(err); process.exit(1); });
