#!/usr/bin/env node
/**
 * OPERATING POSTURE A/B harness.
 *
 * For each (variant × model × scenario × trial), sends a one-shot LLM
 * call to the Loxia backend with the `memory` and `taskmanager` native
 * tools wired up. The backend rewrites tool_calls into a `**Calling X**`
 * text block which we parse out to detect intent.
 *
 * Outputs:
 *   - JSON to scripts/op-posture-bench/results-<ts>.json
 *   - Human-readable summary table to stdout
 *
 * Usage:
 *   LOXIA_API_KEY=... node scripts/op-posture-bench/bench.mjs
 *   node scripts/op-posture-bench/bench.mjs --trials 5 --models gpt-4.1-nano,Kimi-K2.6
 *   node scripts/op-posture-bench/bench.mjs --variants A,B  --scenarios S1,S2,S4
 */

import { mkdir, writeFile } from 'node:fs/promises';
import { readFileSync } from 'node:fs';
import { homedir } from 'node:os';
import { join, dirname, resolve } from 'node:path';
import { fileURLToPath } from 'node:url';
import { performance } from 'node:perf_hooks';

import { VARIANTS } from './variants.mjs';
import { SCENARIOS } from './scenarios.mjs';

const __dirname = dirname(fileURLToPath(import.meta.url));
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', 'Kimi-K2.6'];
const DEFAULT_TRIALS = 3;

// ─── argv ──────────────────────────────────────────────────────────

function parseArgs(argv) {
  const args = {
    models: DEFAULT_MODELS,
    trials: DEFAULT_TRIALS,
    variants: Object.keys(VARIANTS),
    scenarios: SCENARIOS.map(s => s.id),
    out: null,
    concurrency: 4,
  };
  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;
      case '--concurrency': args.concurrency = parseInt(v, 10); break;
      case '--help': case '-h':
        printHelp(); process.exit(0);
      default:
        console.error(`unknown flag: ${f}`); process.exit(1);
    }
  }
  return args;
}

function printHelp() {
  console.log(`OPERATING POSTURE A/B bench
  --models <list>     default: ${DEFAULT_MODELS.join(',')}
  --variants <list>   default: A,B,C
  --scenarios <list>  default: S1-preference,...
  --trials <n>        default: ${DEFAULT_TRIALS}
  --concurrency <n>   default: 4
  --out <path>        default: scripts/op-posture-bench/results-<ts>.json
`);
}

// ─── auth ──────────────────────────────────────────────────────────

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; }
}

// ─── tools wired up for the model ──────────────────────────────────

const TOOLS = [
  {
    type: 'function',
    function: {
      name: 'memory',
      description: 'Persistent memory store. Use to list/read/add/append memories that survive across sessions.',
      parameters: {
        type: 'object',
        properties: {
          action: { type: 'string', enum: ['list', 'read', 'add', 'append'] },
          id:     { type: 'string', description: 'Required for read/append.' },
          title:  { type: 'string', description: 'Required for add. Prefix with "plan/" for working plans that auto-inject into your system prompt.' },
          content:{ type: 'string', description: 'Required for add/append.' },
        },
        required: ['action'],
      },
    },
  },
  {
    type: 'function',
    function: {
      name: 'taskmanager',
      description: 'Step-by-step task list for the current task.',
      parameters: {
        type: 'object',
        properties: {
          action: { type: 'string', enum: ['list', 'create', 'update', 'sync', 'complete'] },
          title:  { type: 'string' },
          tasks:  { type: 'array', items: { type: 'object' } },
        },
        required: ['action'],
      },
    },
  },
];

const BASE_SYSTEM_PROMPT = [
  'You are an autonomous agent helping a developer. You have access to a small set of tools.',
  'Be brief. Take tool actions when appropriate. When you call a tool, the result will come back in a follow-up turn.',
].join(' ');

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

async function callLLM({ apiKey, model, variantBody, userMessage }) {
  const systemPrompt = `${BASE_SYSTEM_PROMPT}\n\n${variantBody}`;
  const requestId = `op-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
  const payload = {
    conversationId: requestId,
    message: userMessage,
    messages: [{ role: 'user', content: userMessage }],
    model,
    requestId,
    options: { max_tokens: 600, temperature: 0.3, tools: TOOLS },
    platformProvided: true,
    systemPrompt,
    apiKey,
  };
  const t0 = performance.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 = performance.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 || {},
    model: data.modelUsed || data.model || model,
    latencyMs,
  };
}

// ─── tool-call extraction ──────────────────────────────────────────

/**
 * Backend rewrites tool_calls as fenced JSON inside the content:
 *   **Calling memory**
 *   ```json
 *   { "toolId": "memory", "parameters": { "action": "add", ... } }
 *   ```
 * Returns an array of { toolId, parameters }. Robust to multiple calls
 * in one response, leading whitespace, and bare ```...``` fences.
 */
function extractToolCalls(content) {
  if (typeof content !== 'string') return [];
  const calls = [];
  const re = /\*\*Calling\s+([\w-]+)\*\*\s*[\r\n]+```(?:json)?\s*([\s\S]*?)\s*```/gi;
  let m;
  while ((m = re.exec(content)) !== null) {
    const toolName = m[1];
    try {
      const parsed = JSON.parse(m[2]);
      calls.push({
        toolId: parsed.toolId || toolName,
        parameters: parsed.parameters || parsed,
      });
    } catch {
      // Sometimes models emit a bare params object without the wrapper.
      calls.push({ toolId: toolName, parameters: { _unparsed: m[2].slice(0, 200) } });
    }
  }
  return calls;
}

// ─── scoring ───────────────────────────────────────────────────────

function score(call, expect) {
  // call is the FIRST tool call (or null). expect = { action, titlePrefix? }
  if (expect.action === 'none') {
    return { hit: !call || (call.toolId !== 'memory'), reason: call ? `unexpected ${call.toolId}.${call.parameters?.action}` : 'no call' };
  }
  if (!call) return { hit: false, reason: 'no tool call' };
  if (call.toolId !== 'memory') return { hit: false, reason: `wrong tool: ${call.toolId}` };
  const action = call.parameters?.action;

  // For memory.list expectations we accept either list or read (both
  // satisfy the "look something up" intent).
  if (expect.action === 'memory.list') {
    const ok = action === 'list' || action === 'read';
    return { hit: ok, reason: ok ? `memory.${action}` : `memory.${action} (wanted list/read)` };
  }
  const wanted = expect.action.split('.')[1]; // 'memory.add' -> 'add'
  if (action !== wanted) return { hit: false, reason: `memory.${action} (wanted memory.${wanted})` };

  // titlePrefix check (only for add)
  if (expect.titlePrefix) {
    const title = call.parameters?.title || '';
    if (!title.startsWith(expect.titlePrefix)) {
      return { hit: true, partial: true, reason: `memory.add, title="${title.slice(0, 30)}" (no "${expect.titlePrefix}" prefix)` };
    }
  }
  return { hit: true, reason: `memory.${action}` };
}

// ─── concurrency-limited runner ────────────────────────────────────

async function runAll({ apiKey, models, variants, scenarios, trials, concurrency }) {
  const jobs = [];
  for (const model of models) {
    for (const variantId of variants) {
      const variant = VARIANTS[variantId];
      if (!variant) { console.warn(`unknown variant ${variantId}`); continue; }
      for (const scId of scenarios) {
        const sc = SCENARIOS.find(s => s.id === scId);
        if (!sc) { console.warn(`unknown scenario ${scId}`); continue; }
        for (let t = 0; t < trials; t++) {
          jobs.push({ model, variantId, variant, scenario: sc, trial: t });
        }
      }
    }
  }
  console.error(`running ${jobs.length} calls (concurrency=${concurrency})…`);

  const results = [];
  let done = 0;
  let inflight = 0;
  let idx = 0;
  return new Promise((resolveAll, rejectAll) => {
    const launch = () => {
      while (inflight < concurrency && idx < jobs.length) {
        const job = jobs[idx++];
        inflight++;
        (async () => {
          try {
            const r = await callLLM({
              apiKey,
              model: job.model,
              variantBody: job.variant.body,
              userMessage: job.scenario.userMessage,
            });
            const calls = extractToolCalls(r.content || '');
            const firstCall = calls[0] || null;
            const sc = score(firstCall, job.scenario.expect);
            results.push({
              model: job.model,
              variant: job.variantId,
              scenario: job.scenario.id,
              trial: job.trial,
              hit: sc.hit,
              partial: sc.partial || false,
              reason: sc.reason,
              firstCall,
              numCalls: calls.length,
              finishReason: r.finishReason,
              latencyMs: Math.round(r.latencyMs),
              tokens: r.usage?.total_tokens || (r.usage?.input_tokens || 0) + (r.usage?.output_tokens || 0),
              error: r.error || null,
              contentPreview: (r.content || '').slice(0, 200),
            });
          } catch (e) {
            results.push({
              model: job.model, variant: job.variantId, scenario: job.scenario.id, trial: job.trial,
              hit: false, reason: `EXCEPTION: ${e.message}`, error: e.message,
            });
          } finally {
            inflight--; done++;
            if (done % 5 === 0 || done === jobs.length) {
              process.stderr.write(`\r${done}/${jobs.length}`);
            }
            if (done === jobs.length) {
              process.stderr.write('\n');
              resolveAll(results);
            } else {
              launch();
            }
          }
        })().catch(rejectAll);
      }
    };
    launch();
  });
}

// ─── aggregation + reporting ───────────────────────────────────────

function aggregate(results) {
  const byCell = new Map(); // key: model|variant -> per-scenario stats
  for (const r of results) {
    const k = `${r.model}|${r.variant}`;
    if (!byCell.has(k)) byCell.set(k, { model: r.model, variant: r.variant, total: 0, hits: 0, partials: 0, perScenario: {} });
    const cell = byCell.get(k);
    cell.total++;
    if (r.hit) cell.hits++;
    if (r.partial) cell.partials++;
    if (!cell.perScenario[r.scenario]) cell.perScenario[r.scenario] = { total: 0, hits: 0 };
    cell.perScenario[r.scenario].total++;
    if (r.hit) cell.perScenario[r.scenario].hits++;
  }
  return [...byCell.values()];
}

function printSummary(agg, scenarioIds) {
  console.log('\n=== OPERATING POSTURE A/B ===\n');

  // overall table
  const cellsByModel = new Map();
  for (const c of agg) {
    if (!cellsByModel.has(c.model)) cellsByModel.set(c.model, {});
    cellsByModel.get(c.model)[c.variant] = c;
  }
  console.log('Model                 | A=status-quo | B=stripped | C=minimal');
  console.log('----------------------+--------------+------------+----------');
  for (const [model, vs] of cellsByModel) {
    const row = ['A', 'B', 'C'].map(v => {
      const c = vs[v];
      return c ? `${c.hits}/${c.total} (${Math.round(100 * c.hits / c.total)}%)`.padEnd(12) : 'n/a         ';
    });
    console.log(`${model.padEnd(21)} | ${row[0]} | ${row[1]} | ${row[2]}`);
  }

  // per-scenario breakdown for each variant
  for (const variantId of ['A', 'B', 'C']) {
    console.log(`\nPer-scenario (variant ${variantId}):`);
    const head = ['Model'.padEnd(21), ...scenarioIds.map(s => s.padEnd(20))].join('| ');
    console.log(head);
    console.log('-'.repeat(head.length));
    for (const [model, vs] of cellsByModel) {
      const c = vs[variantId];
      if (!c) continue;
      const cells = scenarioIds.map(sid => {
        const s = c.perScenario[sid];
        return s ? `${s.hits}/${s.total}`.padEnd(20) : '—'.padEnd(20);
      });
      console.log([model.padEnd(21), ...cells].join('| '));
    }
  }
}

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

(async () => {
  const args = parseArgs(process.argv);
  const apiKey = loadApiKey();
  if (!apiKey) { console.error('No LOXIA_API_KEY. Set env or ~/.loxia/keys.json.'); process.exit(1); }

  const results = await runAll({
    apiKey,
    models: args.models,
    variants: args.variants,
    scenarios: args.scenarios,
    trials: args.trials,
    concurrency: args.concurrency,
  });

  const agg = aggregate(results);
  printSummary(agg, 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: { ...args, models: args.models, variants: args.variants, scenarios: args.scenarios, trials: args.trials },
    aggregate: agg,
    results,
  }, null, 2));
  console.log(`\nfull results → ${outPath}`);
})().catch(err => { console.error(err); process.exit(1); });
