#!/usr/bin/env node
/**
 * v2 live-probe — production-shape multi-turn validation.
 *
 * Talks to running Loxia at :8080. Per session:
 *   1. POST create_agent (full production tool surface)
 *   2. Switch to AGENT mode
 *   3. Seed memories via REST (server-assigns ids; we track the map)
 *   4. Send userTurns[0]; wait idle
 *   5. Send userTurns[1]; wait idle ... etc.
 *   6. Read full conversation from disk
 *   7. Score per-phase per the scenario's rubric
 *   8. Unload agent (preserves conversation file for diagnosis)
 *
 * The CURRENT nudge variant comes from whatever Loxia was launched
 * with (LOXIA_MEMORY_NUDGE_VARIANT env var). The probe records the
 * `--label` arg so you can name the output file accordingly.
 *
 * Usage:
 *   # Variant A run (after Loxia restart with LOXIA_MEMORY_NUDGE_VARIANT=descriptive)
 *   node scripts/memory-bench/live-probe-v2.mjs --label baseline
 *
 *   # Variant B run (after Loxia restart with the env var unset, default=directive)
 *   node scripts/memory-bench/live-probe-v2.mjs --label directive
 *
 *   # Scope down to one scenario / model:
 *   node scripts/memory-bench/live-probe-v2.mjs --scenarios V2-S1-midtask-rule --models Kimi-K2.6 --trials 3
 */

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_V2 } from './scenarios-v2.mjs';

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

const DEFAULT_MODELS = ['Kimi-K2.6', 'gpt-5.4', 'gpt-4.1-mini'];
const DEFAULT_TRIALS = 5;

function parseArgs(argv) {
  const args = {
    models: DEFAULT_MODELS, trials: DEFAULT_TRIALS,
    scenarios: null, label: 'live-v2', 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 = 180_000) {
  const deadline = Date.now() + timeoutMs;
  let consecutiveIdle = 0;
  while (Date.now() < deadline) {
    try {
      const a = await api(`/api/agents/${agentId}/activity`);
      const idle = a.isActive === false && /no-pending-work|no.messages|complete|finished|done/i.test(a.reason || '');
      if (idle) { consecutiveIdle++; 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);
      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);
      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;
}

// ─── per-phase scoring ────────────────────────────────────────────

function scorePhase(phase, callLogPhase, agentTextPhase) {
  if (!phase) return { hit: true, reason: 'no expectation' };
  const reasons = [];
  let hit = false;

  // Action match
  if (Array.isArray(phase.anyOfActions)) {
    for (const a of phase.anyOfActions) {
      const [tool, action] = a.split('.');
      if (callLogPhase.find(c => c.tool === tool && c.action === action)) {
        hit = true;
        reasons.push(`called ${a}`);
        break;
      }
    }
  }
  // Content-fingerprint match (any of)
  if (!hit && Array.isArray(phase.contentFingerprint) && agentTextPhase) {
    const lower = agentTextPhase.toLowerCase();
    const fp = phase.contentFingerprint.find(f => lower.includes(f.toLowerCase()));
    if (fp) {
      hit = true;
      reasons.push(`fingerprint "${fp}"`);
    }
  }
  if (!hit) reasons.push('no action/fingerprint match');
  return { hit, reason: reasons.join(' + ') };
}

function scoreNegative(neg, callLog) {
  if (!neg) return { hit: true, reason: 'no neg expectation' };
  const reasons = [];
  if (Array.isArray(neg.forbiddenActions)) {
    for (const a of neg.forbiddenActions) {
      const [tool, action] = a.split('.');
      const found = callLog.find(c => c.tool === tool && c.action === action);
      if (found) return { hit: false, reason: `forbidden ${a} was called` };
    }
  }
  if (neg.maxCallsPerAction) {
    for (const [actionStr, max] of Object.entries(neg.maxCallsPerAction)) {
      const [tool, action] = actionStr.split('.');
      const count = callLog.filter(c => c.tool === tool && c.action === action).length;
      if (count > max) return { hit: false, reason: `${actionStr} called ${count} times (>${max})` };
    }
  }
  reasons.push('clean');
  return { hit: true, reason: reasons.join(', ') };
}

// ─── one session: multi-turn driver ───────────────────────────────

async function runOneSession({ scenario, model, apiKey }) {
  const systemPrompt =
    'You are an autonomous coding agent helping a developer. Use your tools when relevant. Be concise. When done with the user\'s ask, call `jobdone`.';
  const created = await orchestrator('create_agent', {
    name: `v2-${scenario.id}-${Date.now()}`,
    systemPrompt, model,
    capabilities: ['memory', 'taskmanager', 'jobdone'],
    dynamicModelRouting: false,
    platformProvided: true,
  });
  const agentId = created.id || created.agentId;

  const seedIdMap = new Map();
  // Per-phase boundary tracking: turnBoundary[i] = message index where userTurns[i] sits
  const turnBoundaries = [];
  let totalCalls = [];
  let perPhaseCalls = [];
  let perPhaseText = [];
  let error = null;

  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 },
      });
      const assignedId = resp.memory?.id || resp.id || resp.data?.id || resp.data?.memory?.id;
      if (mem.id && assignedId) seedIdMap.set(mem.id, assignedId);
    }

    // Send each scripted user turn in sequence, waiting for idle between.
    for (let i = 0; i < scenario.userTurns.length; i++) {
      await orchestrator('send_message', {
        agentId, message: scenario.userTurns[i], mode: 'agent', apiKey,
      });
      await waitForIdle(agentId, 180_000);
    }

    // Read full conversation, split into phases by user-turn boundaries.
    const messages = await readConversation(agentId);
    if (!messages) throw new Error('no conversation file');

    // Find indexes of each scripted user turn.
    const userTurnIdxs = [];
    let cursor = 0;
    for (const ut of scenario.userTurns) {
      const idx = messages.findIndex((m, k) =>
        k >= cursor && m.role === 'user' && (m.content || '').includes(ut.slice(0, 60))
      );
      if (idx < 0) throw new Error(`user turn not found in conversation: ${ut.slice(0, 30)}…`);
      userTurnIdxs.push(idx);
      cursor = idx + 1;
    }
    // Phase i = messages between userTurnIdxs[i] (exclusive) and
    // userTurnIdxs[i+1] (exclusive) OR end of conversation.
    for (let i = 0; i < userTurnIdxs.length; i++) {
      const startIdx = userTurnIdxs[i] + 1;
      const endIdx = i + 1 < userTurnIdxs.length ? userTurnIdxs[i + 1] : messages.length;
      const phase = messages.slice(startIdx, endIdx);
      const phaseCalls = [];
      const phaseTextParts = [];
      for (const m of phase) {
        if (m.role === 'assistant') {
          phaseTextParts.push(m.content || '');
          for (const c of extractToolCalls(m.content || '')) {
            phaseCalls.push({ tool: c.toolId, action: c.parameters?.action, params: c.parameters });
          }
        } else if (m.role === 'user' && m.type === 'tool-result') {
          // Capture jobdone summaries which often carry the agent's answer.
          if ((m.content || '').includes('jobdone')) phaseTextParts.push(m.content || '');
        }
      }
      perPhaseCalls.push(phaseCalls);
      perPhaseText.push(phaseTextParts.join('\n'));
      totalCalls.push(...phaseCalls);
    }
  } catch (e) {
    error = e.message;
  } finally {
    try { await orchestrator('unload_agent', { agentId }); } catch { /* best effort */ }
  }

  return { agentId, error, perPhaseCalls, perPhaseText, totalCalls };
}

// ─── scenario-level scoring ───────────────────────────────────────

function scoreSession(scenario, session) {
  if (session.error) return { hit: false, reason: `error: ${session.error}`, sub: {} };
  const exp = scenario.expect;
  const sub = {};

  // S1 shape: { recall, save }
  if (exp.recall || exp.save) {
    sub.recall = scorePhase(exp.recall, session.perPhaseCalls[0] || [], session.perPhaseText[0] || '');
    sub.save = scorePhase(exp.save, session.perPhaseCalls[1] || [], session.perPhaseText[1] || '');
    // Hit if BOTH phases hit.
    const hit = sub.recall.hit && sub.save.hit;
    return { hit, reason: `recall=${sub.recall.hit ? 'Y' : 'N'} save=${sub.save.hit ? 'Y' : 'N'}`, sub };
  }
  // S2 shape: { perTask: [...] }
  if (Array.isArray(exp.perTask)) {
    let hitsPerTask = 0;
    for (let i = 0; i < exp.perTask.length; i++) {
      const r = scorePhase(exp.perTask[i], session.perPhaseCalls[i] || [], session.perPhaseText[i] || '');
      sub[`task${i + 1}`] = r;
      if (r.hit) hitsPerTask++;
    }
    // Partial credit reported; hit only if ALL phases hit.
    return {
      hit: hitsPerTask === exp.perTask.length,
      partialHits: hitsPerTask, totalPhases: exp.perTask.length,
      reason: `${hitsPerTask}/${exp.perTask.length} tasks-with-memory`,
      sub,
    };
  }
  // S3 shape: { negative }
  if (exp.negative) {
    const r = scoreNegative(exp.negative, session.totalCalls);
    return { hit: r.hit, reason: r.reason, sub: { negative: r } };
  }
  return { hit: false, reason: 'unknown expectation shape', sub };
}

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

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

  const scList = args.scenarios
    ? SCENARIOS_V2.filter(s => args.scenarios.includes(s.id))
    : SCENARIOS_V2;

  console.error(`live-probe-v2: label=${args.label} models=${args.models.join(',')} scenarios=${scList.map(s => s.id).join(',')} trials=${args.trials}`);

  const results = [];
  const total = scList.length * args.models.length * args.trials;
  let done = 0;

  for (const scenario of scList) {
    for (const model of args.models) {
      for (let t = 0; t < args.trials; t++) {
        const session = await runOneSession({ scenario, model, apiKey });
        const score = scoreSession(scenario, session);
        done++;
        process.stderr.write(`\r${done}/${total} (sc=${scenario.id} model=${model} trial=${t}) ${score.hit ? 'HIT' : 'MISS'} ${(score.reason||'').slice(0,60)}                       \n`);
        results.push({
          label: args.label, scenario: scenario.id, mechanism: scenario.mechanism,
          model, trial: t, hit: score.hit, reason: score.reason,
          sub: score.sub,
          partialHits: score.partialHits, totalPhases: score.totalPhases,
          totalCallCount: session.totalCalls?.length || 0,
          error: session.error, agentId: session.agentId,
        });
      }
    }
  }
  process.stderr.write('\n');

  // Aggregate per (label, scenario, model).
  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, partial: 0, totalPhases: 0 });
    const c = byCell.get(k);
    c.total++; if (r.hit) c.hits++;
    if (r.partialHits) { c.partial += r.partialHits; c.totalPhases += r.totalPhases; }
  }
  console.log(`\n=== LIVE PROBE V2 (label=${args.label}) ===\n`);
  console.log('Scenario                | Model        | Hits           | Phase partial');
  console.log('------------------------+--------------+----------------+--------------');
  for (const [k, v] of byCell) {
    const [sc, model] = k.split('|');
    const hitPct = Math.round(100 * v.hits / v.total);
    const phasePct = v.totalPhases ? Math.round(100 * v.partial / v.totalPhases) : null;
    const phaseCol = phasePct == null ? 'n/a' : `${v.partial}/${v.totalPhases} (${phasePct}%)`;
    const hitCol = `${v.hits}/${v.total} (${hitPct}%)`.padEnd(14);
    console.log(`${sc.padEnd(23)} | ${model.padEnd(12)} | ${hitCol} | ${phaseCol}`);
  }

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