#!/usr/bin/env node
/**
 * Grounding-model benchmark harness.
 *
 * Runs every adapter × every reasoning-effort × every case (N trials)
 * and aggregates accuracy, latency, token usage, and cost. Outputs:
 *
 *   - human-readable table on stdout
 *   - machine-readable JSON file at --out (default: bench-results-<ts>.json)
 *
 * Usage:
 *   FOUNDRY_ENDPOINT=https://...services.ai.azure.com/ \
 *   FOUNDRY_KEY=<key> \
 *     node scripts/grounding-bench/bench.mjs \
 *       --model kimi-k2.6 \
 *       --effort none,low,medium,high \
 *       --trials 3 \
 *       --case all \
 *       --out bench-$(date +%Y%m%d-%H%M).json
 *
 *   # Dry run against the mock (no spend, validates wiring):
 *     node scripts/grounding-bench/bench.mjs --model mock --effort low --trials 2
 *
 * The bench script is the ONLY consumer of the grounding subsystem
 * for Phase 0. Same imports the future DesktopTool will use — proves
 * the surface from day 1.
 */

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

import {
  createGroundingModel,
  isRegistered,
  REASONING_EFFORT_VALUES,
  ModelId,
  ErrorCode,
  parseCoords,
  distance,
  pointInRect,
  summarise,
} from '../../src/services/grounding/index.js';
import { buildScreenshot, buttonCenter } from './imageBuilder.js';
import { CASES } from './cases.js';

const __dirname = dirname(fileURLToPath(import.meta.url));
const REPO_ROOT  = resolve(__dirname, '..', '..');

const DEFAULT_TRIALS = 3;
const DEFAULT_MODEL_ID = ModelId.KIMI_K2_6;

// ─── argv parsing (no `commander` dep — keep this script self-contained) ────

function parseArgs(argv) {
  const args = {
    model:   DEFAULT_MODEL_ID,
    effort:  ['low'],
    trials:  DEFAULT_TRIALS,
    cases:   'all',
    out:     null,
  };
  for (let i = 2; i < argv.length; i += 2) {
    const flag = argv[i];
    const val  = argv[i + 1];
    switch (flag) {
      case '--model':  args.model = val; break;
      case '--effort': args.effort = val.split(',').map(s => s.trim()).filter(Boolean); break;
      case '--trials': args.trials = parseInt(val, 10); break;
      case '--case':
      case '--cases':  args.cases = val; break;
      case '--out':    args.out = val; break;
      case '--help':
      case '-h':       printHelpAndExit(0);
      default:
        console.error(`unknown flag: ${flag}`);
        printHelpAndExit(1);
    }
  }
  return args;
}

function printHelpAndExit(code) {
  console.log(`grounding bench — usage:
  node scripts/grounding-bench/bench.mjs [--model id] [--effort csv] [--trials N] [--case id|all] [--out path]

  --model    grounding model id (default: ${DEFAULT_MODEL_ID}; also: ${ModelId.MOCK})
  --effort   comma-separated reasoning_effort values to sweep
             (any of: ${REASONING_EFFORT_VALUES.join(', ')})
  --trials   trials per (model, effort, case) cell — default ${DEFAULT_TRIALS}
  --case     case id from cases.js, or "all" — default all
  --out      results JSON path — default ./bench-results-<ts>.json

  ENV (required when --model is not 'mock'):
    FOUNDRY_ENDPOINT  https://<account>.services.ai.azure.com/
    FOUNDRY_KEY       Foundry account key
`);
  process.exit(code);
}

// ─── adapter construction ────────────────────────────────────────────

function makeAdapter(modelId) {
  if (!isRegistered(modelId)) {
    console.error(`unknown model id: ${modelId}`);
    process.exit(2);
  }
  if (modelId === ModelId.MOCK) {
    // Mock returns a stable wrong-but-near coord — useful for verifying
    // the runner machinery without API spend.
    return createGroundingModel(modelId, {
      responses: {},
      latencyMs: 50,
      tokens: 100,
    });
  }
  const endpoint = process.env.FOUNDRY_ENDPOINT;
  const apiKey   = process.env.FOUNDRY_KEY;
  if (!endpoint || !apiKey) {
    console.error('FOUNDRY_ENDPOINT and FOUNDRY_KEY env vars are required for live model runs');
    process.exit(3);
  }
  return createGroundingModel(modelId, { endpoint, apiKey });
}

// ─── trial execution ─────────────────────────────────────────────────

async function runOneTrial({ adapter, kase, effort }) {
  const target = kase.buttons.find(b => b.id === kase.target_button_id);
  if (!target) throw new Error(`case ${kase.id}: no target_button_id ${kase.target_button_id}`);

  const screenshot = await buildScreenshot({
    width: kase.width, height: kase.height, buttons: kase.buttons,
  });

  const intent = kase.intent_template(target);
  const truthCenter = buttonCenter(target);

  const t0 = performance.now();
  let result;
  let errorCode = null;
  try {
    result = await adapter.ground({
      screenshot,
      intent,
      imageSize: { width: kase.width, height: kase.height },
      knobs: { reasoning_effort: effort },
    });
  } catch (err) {
    errorCode = err.code || ErrorCode.PROVIDER_ERROR;
    const elapsedMs = performance.now() - t0;
    return { hit: false, err: null, elapsedMs, tokens: 0, errorCode, pred: null, truthCenter };
  }
  const elapsedMs = performance.now() - t0;

  // Prefer coords parsed by the adapter; fall back to a raw-answer reparse
  // in case the adapter implementation changes in the future.
  const pred = result.coords || parseCoords(result.answer);
  const hit  = pred ? pointInRect(pred, target.rect) : false;
  const err  = pred ? distance(pred, truthCenter) : null;
  const tokens = (result.usage?.input_tokens || 0) + (result.usage?.output_tokens || 0);

  return { hit, err, elapsedMs, tokens, errorCode: null, pred, truthCenter };
}

// ─── orchestration ───────────────────────────────────────────────────

async function main() {
  const args = parseArgs(process.argv);
  const adapter = makeAdapter(args.model);

  const selectedCases = args.cases === 'all'
    ? CASES.slice()
    : CASES.filter(c => c.id === args.cases);
  if (selectedCases.length === 0) {
    console.error(`no cases matched "${args.cases}". Available: ${CASES.map(c => c.id).join(', ')}`);
    process.exit(4);
  }

  // Validate efforts up-front against the adapter so failures surface early.
  validateEfforts(args.effort);

  printHeader(args, selectedCases);

  /** @type {Array<{model: string, effort: string, caseId: string, trials: any[]}>} */
  const cells = [];

  for (const effort of args.effort) {
    for (const kase of selectedCases) {
      const trials = [];
      for (let t = 0; t < args.trials; t++) {
        const trial = await runOneTrial({ adapter, kase, effort });
        trials.push(trial);
        process.stdout.write(formatTrialLine({ effort, kase, t, trial }));
      }
      cells.push({ model: args.model, effort, caseId: kase.id, trials });
      printCellSummary({ effort, kase, trials });
    }
  }

  await writeReport(cells, args);
}

// ─── presentation helpers ────────────────────────────────────────────

function validateEfforts(efforts) {
  for (const e of efforts) {
    if (!REASONING_EFFORT_VALUES.includes(e)) {
      console.error(`unknown effort "${e}". Valid: ${REASONING_EFFORT_VALUES.join(', ')}`);
      process.exit(5);
    }
  }
}

function printHeader(args, cases) {
  console.log('───────────────────────────────────────────────────────────────────────');
  console.log(`Grounding bench`);
  console.log(`  model:   ${args.model}`);
  console.log(`  efforts: ${args.effort.join(', ')}`);
  console.log(`  trials:  ${args.trials} per cell`);
  console.log(`  cases:   ${cases.map(c => c.id).join(', ')}`);
  console.log('───────────────────────────────────────────────────────────────────────');
}

function formatTrialLine({ effort, kase, t, trial }) {
  const status = trial.errorCode
    ? `ERR(${trial.errorCode})`
    : trial.hit ? 'HIT' : 'MISS';
  const errStr = trial.err == null ? '  --  ' : `${trial.err.toFixed(0).padStart(4)}px`;
  return `  [${effort.padEnd(6)}|${kase.id.padEnd(28)}|#${t}] ` +
         `${status.padEnd(12)} err=${errStr} ${trial.elapsedMs.toFixed(0).padStart(5)}ms tok=${String(trial.tokens).padStart(5)}\n`;
}

function printCellSummary({ effort, kase, trials }) {
  const s = summarise(trials);
  const hitsPct = (s.hitRate * 100).toFixed(0).padStart(3);
  console.log(`  ──> [${effort.padEnd(6)}|${kase.id.padEnd(28)}]  ${hitsPct}% hit  ` +
              `mean err ${s.meanErrPx.toFixed(0).padStart(4)}px  ` +
              `p50 ${(s.p50LatencyMs / 1000).toFixed(1)}s  ` +
              `p95 ${(s.p95LatencyMs / 1000).toFixed(1)}s  ` +
              `mean tok ${s.meanTokens.toFixed(0)}\n`);
}

async function writeReport(cells, args) {
  const ts = new Date().toISOString().replace(/[:.]/g, '-');
  const outPath = args.out
    ? resolve(args.out)
    : join(REPO_ROOT, `bench-results-${ts}.json`);

  // Aggregate at multiple grouping levels for easy downstream analysis.
  const report = {
    runAt: new Date().toISOString(),
    args,
    cells: cells.map(c => ({ ...c, summary: summarise(c.trials) })),
    byEffort: groupAndSummarise(cells, c => c.effort),
    byCase:   groupAndSummarise(cells, c => c.caseId),
  };

  await mkdir(dirname(outPath), { recursive: true });
  await writeFile(outPath, JSON.stringify(report, null, 2), 'utf8');

  console.log('───────────────────────────────────────────────────────────────────────');
  console.log('Aggregate summaries\n');
  console.log('By effort:');
  for (const e of report.byEffort) printAggLine(e);
  console.log('\nBy case:');
  for (const c of report.byCase) printAggLine(c);
  console.log(`\nFull report: ${outPath}`);
}

function groupAndSummarise(cells, keyFn) {
  const groups = new Map();
  for (const c of cells) {
    const key = keyFn(c);
    if (!groups.has(key)) groups.set(key, []);
    groups.get(key).push(...c.trials);
  }
  return [...groups.entries()].map(([key, trials]) => ({
    key, ...summarise(trials),
  }));
}

function printAggLine({ key, n, hitRate, meanErrPx, p50LatencyMs, p95LatencyMs, meanTokens }) {
  console.log(
    `  ${key.padEnd(30)} n=${String(n).padStart(3)}  ` +
    `hit=${(hitRate * 100).toFixed(0).padStart(3)}%  ` +
    `err=${meanErrPx.toFixed(0).padStart(4)}px  ` +
    `p50=${(p50LatencyMs / 1000).toFixed(1)}s  p95=${(p95LatencyMs / 1000).toFixed(1)}s  ` +
    `tok=${meanTokens.toFixed(0)}`
  );
}

// ─── go ──────────────────────────────────────────────────────────────

main().catch(err => {
  console.error('\nBench failed:', err.code || '', err.message);
  if (process.env.DEBUG) console.error(err.stack);
  process.exit(1);
});
