#!/usr/bin/env node
/**
 * Live-agent A/B probe — talks to the running Loxia at localhost:8080.
 *
 * Validates the bench's predictions against PRODUCTION code paths
 * (real agentScheduler, memoryService, taskManagerTool, conversation
 * pipeline).
 *
 * For each (scenario × model × trial):
 *   1. POST create_agent (CHAT default)
 *   2. POST /api/agents/:id/mode {mode:'agent'} so the agent processes
 *      autonomously without manual ticks
 *   3. POST /api/agents/:id/memories for each seed memory
 *   4. If scenario seeds tasks, send a setup message and wait for idle
 *   5. Send the trigger message
 *   6. Poll until idle
 *   7. Read the conversation file directly, extract tool calls, score
 *   8. Delete the agent
 *
 * Usage:
 *   node scripts/memory-bench/live-probe.mjs
 *   node scripts/memory-bench/live-probe.mjs --scenarios M3-transition-nudge --trials 3
 *   node scripts/memory-bench/live-probe.mjs --models gpt-4.1-mini --label before
 */

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

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

const __dirname = dirname(fileURLToPath(import.meta.url));
const HOST = process.env.LOXIA_HOST || 'http://localhost:8080';
const SESSION_ID = `bench-${Date.now()}`;

const DEFAULT_MODELS = ['gpt-4.1-nano', 'gpt-4.1-mini'];
const DEFAULT_SCENARIOS = ['M1-catalog-recall', 'M3-transition-nudge'];
const DEFAULT_TRIALS = 3;

function parseArgs(argv) {
  const args = {
    models: DEFAULT_MODELS, scenarios: DEFAULT_SCENARIOS,
    trials: DEFAULT_TRIALS, label: 'live', 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 '--scenarios': args.scenarios = v.split(',').map(s => s.trim()); break;
      case '--trials':    args.trials = parseInt(v, 10); break;
      case '--label':     args.label = v; break;
      case '--out':       args.out = v; break;
      default: console.error(`unknown flag: ${f}`); process.exit(1);
    }
  }
  return args;
}

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

async function orchestrator(action, payload) {
  const res = await fetch(`${HOST}/api/orchestrator`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ sessionId: SESSION_ID, action, payload, projectDir: process.cwd() }),
  });
  const j = await res.json();
  if (!j.success) throw new Error(`orchestrator ${action} failed: ${j.error}`);
  return j.data;
}

async function api(path, opts = {}) {
  const res = await fetch(`${HOST}${path}`, {
    headers: { 'Content-Type': 'application/json', ...(opts.headers || {}) },
    method: opts.method || 'GET',
    body: opts.body ? JSON.stringify(opts.body) : undefined,
  });
  if (!res.ok) {
    const text = await res.text();
    throw new Error(`${opts.method || 'GET'} ${path} → HTTP ${res.status}: ${text.slice(0, 200)}`);
  }
  return res.json();
}

async function setAgentMode(agentId, mode) {
  await api(`/api/agents/${agentId}/mode`, { method: 'POST', body: { mode } });
}

async function waitForIdle(agentId, timeoutMs = 90_000) {
  const deadline = Date.now() + timeoutMs;
  let consecutiveIdle = 0;
  while (Date.now() < deadline) {
    try {
      const a = await api(`/api/agents/${agentId}/activity`);
      // Agent is idle when isActive is false AND reason is one of the
      // terminal-shape reasons (no-pending-work, chat-mode-no-messages).
      const idle = a.isActive === false && /no-pending-work|no.messages|complete|finished/i.test(a.reason || '');
      if (idle) {
        consecutiveIdle++;
        // Two consecutive idle checks → really done (avoids racing
        // between turn boundaries).
        if (consecutiveIdle >= 2) return true;
      } else {
        consecutiveIdle = 0;
      }
    } catch { /* keep polling */ }
    await new Promise(r => setTimeout(r, 1500));
  }
  return false;
}

const USER_DATA_DIR = process.env.LOXIA_USER_DATA ||
  join(homedir(), 'AppData', 'Local', 'loxia-autopilot');

async function readConversation(agentId) {
  const candidates = [
    join(USER_DATA_DIR, 'state', 'agents', `agent-${agentId}-conversations.json`),
    join(USER_DATA_DIR, 'state', 'agents', `${agentId}-conversations.json`),
  ];
  for (const p of candidates) {
    try {
      const s = await readFile(p, 'utf8');
      const j = JSON.parse(s);
      // Convention: { conversations: { full: { messages: [...] }, ... } }
      const full = j.conversations?.full || j.conversations?.[Object.keys(j.conversations || {})[0]];
      return full?.messages || j.messages || [];
    } catch { /* try next */ }
  }
  return null;
}

function extractToolCalls(content) {
  if (typeof content !== 'string') return [];
  const calls = [];
  let pos = 0;
  while (pos < content.length) {
    const m = content.indexOf('**Calling ', pos);
    if (m === -1) break;
    const ns = m + '**Calling '.length;
    const ne = content.indexOf('**', ns);
    if (ne === -1) break;
    const name = content.slice(ns, ne);
    const fs = content.indexOf('```', ne);
    if (fs === -1) break;
    const bs = content.indexOf('\n', fs) + 1;
    const fe = content.indexOf('```', bs);
    if (fe === -1) break;
    let body = content.slice(bs, fe).trim();
    if (body.startsWith('json')) body = body.slice(4).trim();
    try {
      const p = JSON.parse(body);
      // Production wraps multi-call payloads as parameters.actions[].
      // Bench-scoring rubric expects { tool, action } flat — normalize.
      const actions = p.parameters?.actions;
      if (Array.isArray(actions)) {
        for (const a of actions) {
          calls.push({ toolId: p.toolId || name, parameters: { ...a, action: a.type || a.action } });
        }
      } else {
        calls.push({ toolId: p.toolId || name, parameters: p.parameters || p });
      }
    } catch { /* unparsed */ }
    pos = fe + 3;
  }
  return calls;
}

function score(scenario, callLog, seedIdMap, agentText = '') {
  const exp = scenario.expect;
  if (exp.action === 'none') {
    const wrote = callLog.some(c => c.tool === 'memory' && (c.action === 'add' || c.action === 'append'));
    return { hit: !wrote, reason: wrote ? 'unexpected memory write' : 'correct: no memory write' };
  }

  // Translate scenario seed-id → server-assigned id at scoring time so
  // we match against whatever the real memoryService minted.
  const expectedId = exp.memoryId
    ? (seedIdMap?.get(exp.memoryId) || exp.memoryId)
    : null;

  const matches = (c, t, a) => {
    if (c.tool !== t || c.action !== a) return false;
    if (expectedId && (a === 'read' || a === 'append')) {
      return c.params?.id === expectedId;
    }
    return true;
  };

  const [tool, action] = exp.action.split('.');
  let match = callLog.find(c => matches(c, tool, action));
  if (!match && exp.orAction) {
    const [t2, a2] = exp.orAction.split('.');
    match = callLog.find(c => matches(c, t2, a2));
  }
  if (!match && Array.isArray(exp.orActions)) {
    for (const alt of exp.orActions) {
      const [ta, aa] = alt.split('.');
      const found = callLog.find(c => matches(c, ta, aa));
      if (found) { match = found; break; }
    }
  }
  if (!match) {
    // Final fallback: content-fingerprint match in the agent's text.
    // The mechanism is "agent uses the memory's info"; if the answer
    // verbatim contains distinctive content from the seed memory, the
    // catalog injection clearly delivered.
    if (Array.isArray(exp.answerContentAny) && agentText) {
      const lower = agentText.toLowerCase();
      const hit = exp.answerContentAny.find(fp => lower.includes(fp.toLowerCase()));
      if (hit) return { hit: true, reason: `agent answer referenced memory content: "${hit}"` };
    }
    const wrong = callLog.find(c => c.tool === tool && c.action === action);
    if (wrong && expectedId) return { hit: false, reason: `${exp.action} on wrong id: ${wrong.params?.id} (wanted ${expectedId})` };
    return { hit: false, reason: `expected ${exp.action}, not called` };
  }
  return { hit: true, reason: `matched ${match.tool}.${match.action}` };
}

// ─── one trial ────────────────────────────────────────────────────

async function runOneTrial({ scenario, model, apiKey }) {
  const systemPrompt = 'You are an autonomous coding agent helping a developer. Use your tools when relevant. Be brief. When done, call `jobdone`.';
  const created = await orchestrator('create_agent', {
    name: `live-${scenario.id}-${Date.now()}`,
    systemPrompt, model,
    capabilities: ['memory', 'taskmanager', 'jobdone'],
    dynamicModelRouting: false,
    platformProvided: true,
  });
  // The API endpoints expect the full agent ID INCLUDING the `agent-`
  // prefix (e.g. POST /api/agents/agent-foo-123/mode). Don't strip.
  const agentId = created.id || created.agentId;

  let scoreResult = { hit: false, reason: 'aborted' };
  let callLog = [];
  let usage = { promptTokens: 0, outputTokens: 0 };

  // Map { seed-id-in-scenario → server-assigned-id-after-POST }
  // — the memoryService refuses to honor caller-supplied ids and
  // mints its own, so scoring needs the translation.
  const seedIdMap = new Map();

  try {
    await setAgentMode(agentId, 'agent');

    for (const mem of scenario.seedMemories) {
      const resp = await api(`/api/agents/${agentId}/memories`, {
        method: 'POST',
        body: { title: mem.title, description: mem.description, content: mem.content },
      });
      // Common server shapes: { memory: { id } } | { id } | { data: { id } }
      const assignedId = resp.memory?.id || resp.id || resp.data?.id || resp.data?.memory?.id;
      if (mem.id && assignedId) seedIdMap.set(mem.id, assignedId);
    }

    // If we need seeded tasks, prompt the agent to create them first.
    if (scenario.seedTasks.length > 0) {
      const setupMsg = 'Setup phase: create these tasks WITHOUT doing any other work. ' +
        scenario.seedTasks.map(t =>
          `Task title="${t.title}", description="${t.description}", status=pending.`
        ).join(' ') +
        ' After creating them, call jobdone with summary="tasks created".';
      await orchestrator('send_message', {
        agentId, message: setupMsg, mode: 'agent', apiKey,
      });
      await waitForIdle(agentId, 60_000);
    }

    // Trigger.
    await orchestrator('send_message', {
      agentId, message: scenario.userTurns[0], mode: 'agent', apiKey,
    });
    await waitForIdle(agentId, 120_000);

    // Read conversation, extract calls AFTER the trigger only.
    const messages = await readConversation(agentId);
    if (!messages) {
      scoreResult = { hit: false, reason: 'no conversation file' };
    } else {
      // Find the trigger message and only consider assistant turns after it.
      let triggerIdx = -1;
      for (let i = messages.length - 1; i >= 0; i--) {
        const m = messages[i];
        if (m.role === 'user' && (m.content || '').includes(scenario.userTurns[0].slice(0, 60))) {
          triggerIdx = i; break;
        }
      }
      const after = triggerIdx >= 0 ? messages.slice(triggerIdx + 1) : messages;
      // Build a flat text blob of all assistant text after the trigger
      // — including jobdone summaries — for content-fingerprint scoring.
      const agentText = after
        .filter(m => m.role === 'assistant' || (m.role === 'user' && m.type === 'tool-result' && /jobdone/i.test(m.content || '')))
        .map(m => m.content || '')
        .join('\n');
      for (const m of after) {
        if (m.role !== 'assistant') continue;
        for (const c of extractToolCalls(m.content || '')) {
          callLog.push({ tool: c.toolId, action: c.parameters?.action, params: c.parameters });
        }
      }
      scoreResult = score(scenario, callLog, seedIdMap, agentText);
    }
  } catch (e) {
    scoreResult = { hit: false, reason: `error: ${e.message}` };
  } finally {
    // Unload (not delete) so the conversation file survives for
    // diagnosis. The user can manually clean these up after the run.
    try {
      await orchestrator('unload_agent', { agentId });
    } catch (e) { /* best effort */ }
  }

  return { agentId, scoreResult, callLog, usage };
}

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

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

  console.error(`live-probe: label=${args.label} scenarios=${args.scenarios.join(',')} models=${args.models.join(',')} trials=${args.trials}`);

  const scList = SCENARIOS.filter(s => args.scenarios.includes(s.id));
  const results = [];
  let done = 0;
  const total = scList.length * args.models.length * args.trials;

  for (const scenario of scList) {
    for (const model of args.models) {
      for (let t = 0; t < args.trials; t++) {
        const r = await runOneTrial({ scenario, model, apiKey });
        done++;
        process.stderr.write(`\r${done}/${total} (sc=${scenario.id} model=${model} trial=${t}) ${r.scoreResult.hit ? 'HIT ' : 'MISS'} ${r.scoreResult.reason.slice(0,60)}                       \n`);
        results.push({
          label: args.label, scenario: scenario.id, mechanism: scenario.mechanism,
          model, trial: t, hit: r.scoreResult.hit, reason: r.scoreResult.reason,
          callLog: r.callLog.slice(0, 10), agentId: r.agentId,
        });
      }
    }
  }
  process.stderr.write('\n');

  // Aggregate.
  const byCell = new Map();
  for (const r of results) {
    const k = `${r.scenario}|${r.model}`;
    if (!byCell.has(k)) byCell.set(k, { hits: 0, total: 0 });
    const c = byCell.get(k);
    c.total++; if (r.hit) c.hits++;
  }
  console.log(`\n=== LIVE PROBE (label=${args.label}) ===\n`);
  console.log('Scenario              | Model        | Hits   ');
  console.log('----------------------+--------------+--------');
  for (const [k, v] of byCell) {
    const [sc, model] = k.split('|');
    console.log(`${sc.padEnd(21)} | ${model.padEnd(12)} | ${v.hits}/${v.total} (${Math.round(100*v.hits/v.total)}%)`);
  }
  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++;
  }
  console.log('\nBy mechanism:');
  for (const [mech, m] of byMech) {
    console.log(`  ${mech.padEnd(32)} ${m.hits}/${m.total}  (${Math.round(100*m.hits/m.total)}%)`);
  }

  const ts = new Date().toISOString().replace(/[:.]/g, '-');
  const outPath = args.out || join(__dirname, `live-${args.label}-${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); });
