#!/usr/bin/env node
/**
 * Multi-turn memory-feedback bench.
 *
 * Tests the four memory mechanisms shipped in 5bca65d:
 *   - catalog injection (`## AGENT MEMORY` in system prompt)
 *   - task-transition auto-load (memory attached → inline content)
 *   - conditional nudge (no memory attached → catalog pointer)
 *   - post-compaction nudge (one-shot after compaction)
 *   - append vs add discipline
 *
 * Each scenario is a multi-turn session against the Loxia backend.
 * Pre-populates a per-session in-memory memory + task store, builds
 * the system prompt the way agentScheduler would (catalog block,
 * post-compaction nudge if flag set), runs the agent loop with
 * native tool schemas, scores against an expected behaviour.
 *
 * Variants:
 *   baseline  — features as currently shipped
 *   <variants added when optimizing the weakest mechanism>
 *
 * Usage:
 *   node scripts/memory-bench/bench.mjs --variant baseline
 *   node scripts/memory-bench/bench.mjs --variant baseline --models gpt-4.1-nano
 */
import { readFileSync } from 'node:fs';
import { mkdir, writeFile } from 'node:fs/promises';
import { homedir } from 'node:os';
import { dirname, join, resolve } from 'node:path';
import { fileURLToPath } from 'node:url';

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 = 2;
const MAX_TURNS = 5;

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

function parseArgs(argv) {
  const args = {
    models: DEFAULT_MODELS, trials: DEFAULT_TRIALS,
    scenarios: null, variant: 'baseline', 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 '--scenarios': args.scenarios = v.split(',').map(s => s.trim()); break;
      case '--variant':   args.variant = v; 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; }
}

// ─── Variants — same prompt skeleton, swappable nudge / catalog text ──

/**
 * Build a catalog block from a memory store. Matches the format in
 * agentScheduler.js (descriptive, not prescriptive, minimal).
 */
function buildCatalogBlock(memories, variant) {
  if (!memories || memories.length === 0) return '';
  const lines = ['', '## AGENT MEMORY (catalog — ids and titles only; use `memory read` for content)', ''];
  for (const m of memories) {
    const desc = m.description ? ` "${m.description}"` : '';
    lines.push(`- ${m.id}  ${m.title}${desc}`);
  }
  // VARIANT EXPERIMENT POINTS — left as hooks for optimization sweeps.
  return lines.join('\n');
}

function buildPostCompactionNudge(variant) {
  // The variant string can override the wording. Baseline matches what
  // ships in agentScheduler.js today.
  if (variant === 'compaction-stripped') return '';
  if (variant === 'compaction-tight') {
    return '\n\n## POST-COMPACTION NUDGE (one-time)\nSome older context was summarized. Check `## AGENT MEMORY` for relevant entries before answering.';
  }
  return '\n\n## POST-COMPACTION NUDGE (one-time)\n\nCompaction just summarized older messages. If your current work continues from before the compaction boundary, scan the `## AGENT MEMORY` catalog above — older context that was summarized may still be available as a memory you can `memory read` by id.';
}

const BASE_SYSTEM_PROMPT = [
  'You are an autonomous agent helping a developer. You have access to a small set of tools.',
  'Use them when appropriate. Be concise.',
  '',
  '## OPERATING POSTURE',
  '',
  '- New topic? Scan `## AGENT MEMORY` (below) for relevant entries.',
  '- User states a durable preference, rule, or multi-part ask? `memory add` (title `plan/...` for plans).',
  '- Augmenting an existing memory entry? `memory append` (do NOT create a duplicate).',
  '- Starting a task that references `[memory:mem-...]`? The runtime auto-loads it for you.',
].join('\n');

/**
 * Build the full system prompt the way agentScheduler would for a turn.
 */
function buildSystemPrompt({ memories, tasks, postCompactionFlag, variant }) {
  let p = BASE_SYSTEM_PROMPT;
  if (Array.isArray(tasks) && tasks.length > 0) {
    p += '\n\n## CURRENT TASK LIST\n';
    for (const t of tasks) {
      p += `- ${t.id} [${t.status}] ${t.title}: ${t.description}\n`;
    }
  }
  // Catalog injection — only if memory entries exist.
  p += buildCatalogBlock(memories, variant);
  if (postCompactionFlag) p += buildPostCompactionNudge(variant);
  return p;
}

// ─── Tool schemas given to the model ───────────────────────────────

const TOOLS = [{
  type: 'function',
  function: {
    name: 'memory',
    description: 'Persistent memory store. list (titles only), read (full content), add (new entry — title prefix `plan/...` for plans), append (extend existing).',
    parameters: {
      type: 'object',
      properties: {
        action:  { type: 'string', enum: ['list', 'read', 'add', 'append'] },
        id:      { type: 'string' },
        title:   { type: 'string' },
        content: { type: 'string' },
      },
      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', 'complete'] },
        taskId: { type: 'string' },
        status: { type: 'string', enum: ['pending', 'in_progress', 'completed', 'blocked'] },
        title:  { type: 'string' },
      },
      required: ['action'],
    },
  },
}, {
  type: 'function',
  function: {
    name: 'job-done',
    description: 'Signal that you have finished answering the user.',
    parameters: {
      type: 'object',
      properties: { summary: { type: 'string' } },
      required: ['summary'],
    },
  },
}];

// ─── tool-call extraction (same as op-posture-bench) ─────────────

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,...}]}`.
      // Expand into one call per action so the scoring rubric sees each.
      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;
}

// ─── per-session memory store ─────────────────────────────────────

function makeStore(seedMemories, seedTasks) {
  return {
    memories: seedMemories.map(m => ({ ...m })),
    tasks: seedTasks.map(t => ({ ...t })),
    // Track every tool call the agent made so we can score.
    callLog: [],
  };
}

function executeMemoryCall(store, params) {
  const action = params?.action;
  store.callLog.push({ tool: 'memory', action, params });
  switch (action) {
    case 'list':
      return store.memories.map(m => ({ id: m.id, title: m.title, description: m.description }));
    case 'read': {
      const m = store.memories.find(x => x.id === params.id);
      return m ? { id: m.id, title: m.title, content: m.content } : { error: 'not found' };
    }
    case 'add': {
      const id = `mem-${Date.now()}-${Math.random().toString(36).slice(2, 6)}`;
      const m = { id, title: params.title || '(untitled)', description: (params.content || '').slice(0, 100), content: params.content || '' };
      store.memories.push(m);
      return { id, title: m.title };
    }
    case 'append': {
      const m = store.memories.find(x => x.id === params.id);
      if (!m) return { error: 'not found' };
      m.content = m.content + '\n\n## (appended)\n' + (params.content || '');
      return { id: m.id, title: m.title, contentLength: m.content.length };
    }
    default:
      return { error: `unknown action ${action}` };
  }
}

/**
 * Build the conditional-nudge text for a task transitioning to in_progress
 * without an attached memory marker. Variants swap the wording.
 *
 * Baseline matches what ships in taskManagerTool.js today.
 * V2 (directive): explicit prescribed next action.
 * V3 (echo-catalog): inline list of available memory entries so the
 *   agent doesn't have to look up at the system prompt.
 */
function buildConditionalNudge(memories, variant) {
  if (variant === 'nudge-directive') {
    return 'No memory is attached to this task. Your VERY NEXT tool call must be `memory list` to check for applicable rules, preferences, or prior findings. Only THEN proceed with the work.';
  }
  if (variant === 'nudge-echo-catalog') {
    const lines = ['No memory is attached to this task. Available memory entries that may apply:'];
    for (const m of memories) {
      const desc = m.description ? ` "${m.description}"` : '';
      lines.push(`  • ${m.id}  ${m.title}${desc}`);
    }
    lines.push('Call `memory read <id>` on any that look relevant before doing the work.');
    return lines.join('\n');
  }
  if (variant === 'nudge-combined') {
    const lines = [
      'No memory is attached to this task. Available memory entries:',
    ];
    for (const m of memories) {
      const desc = m.description ? ` — ${m.description}` : '';
      lines.push(`  • ${m.id}  ${m.title}${desc}`);
    }
    lines.push('Your VERY NEXT tool call should be `memory read <id>` on any relevant entry, or proceed if none apply.');
    return lines.join('\n');
  }
  // baseline
  return 'No memory is attached to this task. Scan the `## AGENT MEMORY` catalog in your system prompt for adjacent context that may inform this work.';
}

function executeTaskCall(store, params, variant) {
  const action = params?.action;
  store.callLog.push({ tool: 'taskmanager', action, params });
  if (action === 'update' && params.taskId) {
    const t = store.tasks.find(x => x.id === params.taskId);
    if (!t) return { error: 'not found' };
    const prevStatus = t.status;
    if (params.status) t.status = params.status;
    const result = { id: t.id, title: t.title, status: t.status };
    // Transition handler: auto-load memory if marker present and just-now in_progress.
    if (prevStatus !== 'in_progress' && t.status === 'in_progress') {
      const markerStart = (t.description || '').indexOf('[memory:');
      if (markerStart !== -1) {
        const end = (t.description || '').indexOf(']', markerStart);
        if (end !== -1) {
          const memId = t.description.slice(markerStart + 8, end);
          const mem = store.memories.find(x => x.id === memId);
          if (mem) {
            result.attachedMemory = { id: mem.id, title: mem.title, content: mem.content,
              _hint: 'Auto-loaded because the task references this memory. Use `memory append` if your work yields findings worth recording.' };
          } else {
            result.attachedMemory = { id: memId, error: 'memory referenced in task description was not found' };
          }
        }
      } else {
        result.memoryNudge = buildConditionalNudge(store.memories, variant);
      }
    }
    return result;
  }
  if (action === 'list') return store.tasks;
  return { error: `unsupported action ${action} (this is a bench stub)` };
}

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

async function callLLM({ apiKey, model, messages, systemPrompt }) {
  const requestId = `mem-${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 ──────────────────────────────────────────────────

async function runSession({ apiKey, model, scenario, variant }) {
  const store = makeStore(scenario.seedMemories, scenario.seedTasks);
  let messages = [];
  let postCompactionFlag = scenario.postCompactionFlag;
  const trace = { turns: [], error: null,
    totalPromptTokens: 0, totalOutputTokens: 0, totalLatencyMs: 0 };

  let userTurnIdx = 0;
  if (userTurnIdx < scenario.userTurns.length) {
    messages.push({ role: 'user', content: scenario.userTurns[userTurnIdx++] });
  }

  for (let turn = 0; turn < MAX_TURNS; turn++) {
    const systemPrompt = buildSystemPrompt({
      memories: store.memories, tasks: store.tasks,
      postCompactionFlag, variant,
    });
    // After first turn, the post-compaction flag is consumed.
    postCompactionFlag = false;

    const r = await callLLM({ apiKey, model, messages, 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, 200) });
    messages.push({ role: 'assistant', content: r.content });

    if (calls.length === 0) {
      // Plain text answer. End session.
      break;
    }

    let jobDone = false;
    const toolResultParts = [];
    for (const call of calls) {
      if (call.toolId === 'job-done') { jobDone = true; break; }
      if (call.toolId === 'memory') {
        const out = executeMemoryCall(store, call.parameters);
        toolResultParts.push(`memory ${call.parameters?.action} → ${JSON.stringify(out).slice(0, 600)}`);
      } else if (call.toolId === 'taskmanager') {
        const out = executeTaskCall(store, call.parameters, variant);
        toolResultParts.push(`taskmanager ${call.parameters?.action} → ${JSON.stringify(out).slice(0, 800)}`);
      } else {
        toolResultParts.push(`unknown tool: ${call.toolId}`);
      }
    }
    if (jobDone) break;
    if (toolResultParts.length > 0) {
      messages.push({
        role: 'user', type: 'tool-result',
        content: '[Tool Results]\n' + toolResultParts.join('\n---\n'),
        responseTurnId: `turn-${turn}`,
      });
    }
    // Hand the next user turn forward, if any.
    if (userTurnIdx < scenario.userTurns.length) {
      messages.push({ role: 'user', content: scenario.userTurns[userTurnIdx++] });
    }
  }

  return { trace, store };
}

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

function scoreSession(scenario, store) {
  const calls = store.callLog;
  const exp = scenario.expect;
  if (exp.action === 'none') {
    // Negative control — any memory.add or memory.append is a fail.
    const wrote = calls.some(c => c.tool === 'memory' && (c.action === 'add' || c.action === 'append'));
    return { hit: !wrote, reason: wrote ? `unexpected memory write` : 'no write — correct' };
  }
  const [tool, action] = exp.action.split('.');
  // Find ANY call that matches tool+action AND any id constraint —
  // an agent that hit the right memory eventually counts as a hit,
  // even if it stuttered on its first attempt (this happens with
  // smaller models that get the id format wrong on the first try).
  let match = calls.find(c => {
    if (c.tool !== tool || c.action !== action) return false;
    if (exp.memoryId && (action === 'read' || action === 'append')) {
      return c.params?.id === exp.memoryId;
    }
    return true;
  });
  if (!match && exp.orAction) {
    const [t2, a2] = exp.orAction.split('.');
    const alt = calls.find(c => {
      if (c.tool !== t2 || c.action !== a2) return false;
      // If memoryId is required, the orAction must satisfy it.
      if (exp.memoryId && (a2 === 'read' || a2 === 'append')) {
        return c.params?.id === exp.memoryId;
      }
      return true;
    });
    if (alt) return { hit: true, reason: `matched orAction ${exp.orAction}${exp.memoryId ? ' on ' + exp.memoryId : ''}` };
  }
  if (!match && Array.isArray(exp.orActions)) {
    for (const alt of exp.orActions) {
      const [ta, aa] = alt.split('.');
      const found = calls.find(c => {
        if (c.tool !== ta || c.action !== aa) return false;
        if (exp.memoryId && (aa === 'read' || aa === 'append')) {
          return c.params?.id === exp.memoryId;
        }
        return true;
      });
      if (found) return { hit: true, reason: `matched orAction ${alt}` };
    }
  }
  if (!match) {
    // Diagnose: was the call attempted with a wrong id?
    if (exp.memoryId) {
      const wrong = calls.find(c => c.tool === tool && c.action === action);
      if (wrong) return { hit: false, reason: `${exp.action} on wrong id: ${wrong.params?.id} (wanted ${exp.memoryId})` };
    }
    return { hit: false, reason: `expected ${exp.action}, not called` };
  }
  return { hit: true, reason: `matched ${exp.action}` };
}

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

async function runAll({ apiKey, models, trials, scenarios, variant }) {
  const scList = scenarios
    ? SCENARIOS.filter(s => scenarios.includes(s.id))
    : SCENARIOS;
  const results = [];
  const total = models.length * scList.length * trials;
  let done = 0;
  for (const model of models) {
    for (const scenario of scList) {
      for (let t = 0; t < trials; t++) {
        const { trace, store } = await runSession({ apiKey, model, scenario, variant });
        const score = scoreSession(scenario, store);
        results.push({
          model, scenario: scenario.id, mechanism: scenario.mechanism,
          variant, trial: t, hit: score.hit, reason: score.reason,
          callLog: store.callLog.slice(0, 10),
          totalPromptTokens: trace.totalPromptTokens,
          totalOutputTokens: trace.totalOutputTokens,
          turns: trace.turns.length,
          error: trace.error,
        });
        done++;
        process.stderr.write(`\r${done}/${total} (model=${model} sc=${scenario.id} trial=${t})  `);
      }
    }
  }
  process.stderr.write('\n');
  return results;
}

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

function printReport(results, variant) {
  // Per (model, scenario) hit rate.
  const byCell = new Map();
  for (const r of results) {
    const k = `${r.model}|${r.scenario}`;
    if (!byCell.has(k)) byCell.set(k, { hits: 0, total: 0, mechanism: r.mechanism });
    const c = byCell.get(k);
    c.total++; if (r.hit) c.hits++;
  }
  console.log(`\n=== MEMORY-BENCH (variant=${variant}) ===\n`);
  console.log('Model        | Scenario              | Mechanism                   | Hits  ');
  console.log('-------------+-----------------------+-----------------------------+-------');
  for (const [k, v] of byCell) {
    const [model, scenario] = k.split('|');
    console.log(`${model.padEnd(12)} | ${scenario.padEnd(21)} | ${v.mechanism.padEnd(27)} | ${v.hits}/${v.total}  (${Math.round(100 * v.hits / v.total)}%)`);
  }
  // Per mechanism overall.
  console.log('\nBy mechanism (across all models):');
  const byMech = new Map();
  for (const r of results) {
    if (!byMech.has(r.mechanism)) byMech.set(r.mechanism, { hits: 0, total: 0 });
    const m = byMech.get(r.mechanism);
    m.total++; if (r.hit) m.hits++;
  }
  for (const [mech, m] of byMech) {
    console.log(`  ${mech.padEnd(32)} ${m.hits}/${m.total}  (${Math.round(100 * m.hits / m.total)}%)`);
  }
}

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

(async () => {
  const args = parseArgs(process.argv);
  const apiKey = loadApiKey();
  if (!apiKey) { console.error('No LOXIA_API_KEY.'); process.exit(1); }
  console.error(`memory-bench: variant=${args.variant} models=${args.models.join(',')} trials=${args.trials}`);
  const results = await runAll({
    apiKey, models: args.models, trials: args.trials,
    scenarios: args.scenarios, variant: args.variant,
  });
  printReport(results, args.variant);

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