#!/usr/bin/env node
// Post-hoc session telemetry for /new and /prd (opt-in via `-stats`).
//
// Reads the live Claude Code transcripts — the orchestrator's main transcript
// plus every subagent transcript — and aggregates REAL token usage and wall-clock
// time by AGENT ROLE (coder, code-reviewer, qa-sentinel, doc-reviewer, …). The
// data already exists in the transcripts (`message.usage` + `timestamp`), so this
// costs ZERO model tokens: it runs in Node, after the run, outside the model.
//
// Transcript layout (verified):
//   ~/.claude/projects/<slug>/<session_id>.jsonl                  ← orchestrator (main)
//   ~/.claude/projects/<slug>/<session_id>/subagents/agent-*.jsonl ← one file per subagent
//   where slug = cwd with every "/" replaced by "-".
//
// Role mapping: each `Agent` tool_use in the main transcript carries
// `input.subagent_type` (= role) and `input.prompt`. Every subagent file's first
// user message IS that prompt, so we link file→role by prompt-prefix match
// (verified 20/20 on a real /new run). Unlinkable files fall back to role
// "unknown" with a slug-derived label — their token/time data stays accurate.
//
// Cache accounting: the four usage counters are kept SEPARATE on purpose.
// `cache_read` is billed ~10% of fresh input; conflating them would overstate
// cost and mislead "what consumed the most". A single headline `cost_units`
// figure weights them: output + fresh_input + cache_creation×1.25 + cache_read×0.1.
//
// Usage:
//   node analyze-session-tokens.js --skill new --run-id batch-FEAT-0001
//   node analyze-session-tokens.js --skill prd --run-id prd-checkout --out-dir docs/metrics
//   # test/override (bypass env + cwd derivation):
//   node analyze-session-tokens.js --skill new --run-id X --session <id> --project-dir <path>
//
// Output:
//   - <out-dir>/sessions/<run_id>.md  — human-readable per-role breakdown + totals
//   - stdout: same summary, PLUS a final line `TELEMETRY_JSON=<json>` carrying the
//     totals object for the skill to embed in docs/metrics/skill-runs.jsonl.
//
// Fail-safe: ANY error prints `stats: SKIPPED (<reason>)` and exits 0. This script
// must NEVER abort a /new or /prd run.

'use strict';

const fs = require('fs');
const os = require('os');
const path = require('path');

// ---- arg parsing -----------------------------------------------------------

function getArg(name, fallback) {
  const i = process.argv.indexOf(`--${name}`);
  return i !== -1 && process.argv[i + 1] ? process.argv[i + 1] : fallback;
}

// Bail out softly — telemetry is never allowed to break a run.
function skip(reason) {
  console.log(`stats: SKIPPED (${reason})`);
  process.exit(0);
}

const skill = getArg('skill', 'new');
const runId = getArg('run-id', null);
const outDir = getArg('out-dir', path.join('docs', 'metrics'));
const sessionId = getArg('session', process.env.CLAUDE_CODE_SESSION_ID || null);
const projectDir = getArg('project-dir', null);

if (!runId) skip('no --run-id provided');
if (!sessionId) skip('no session id ($CLAUDE_CODE_SESSION_ID unset, no --session)');

// ---- locate transcripts ----------------------------------------------------

function projectSlug(cwd) {
  // Claude Code stores transcripts under a slug = cwd with "/" → "-".
  return cwd.replace(/\//g, '-');
}

let baseDir;
if (projectDir) {
  baseDir = projectDir;
} else {
  baseDir = path.join(os.homedir(), '.claude', 'projects', projectSlug(process.cwd()));
}

const mainTranscript = path.join(baseDir, `${sessionId}.jsonl`);
const subDir = path.join(baseDir, sessionId, 'subagents');

if (!fs.existsSync(mainTranscript)) skip(`main transcript not found: ${mainTranscript}`);

// ---- jsonl helpers ---------------------------------------------------------

function* readLines(file) {
  let raw;
  try {
    raw = fs.readFileSync(file, 'utf8');
  } catch (err) {
    return;
  }
  for (const line of raw.split(/\r?\n/)) {
    if (!line.trim()) continue;
    let obj;
    try {
      obj = JSON.parse(line);
    } catch (err) {
      continue; // tolerate partial / malformed lines
    }
    yield obj;
  }
}

function contentBlocks(rec) {
  const c = rec && rec.message && rec.message.content;
  if (Array.isArray(c)) return c;
  if (typeof c === 'string') return [{ type: 'text', text: c }];
  return [];
}

function blocksText(blocks) {
  return blocks
    .map((b) => (typeof b === 'string' ? b : b && typeof b.text === 'string' ? b.text : ''))
    .join('')
    .trim();
}

// Aggregate the four usage counters + wall-clock span of one transcript file.
function aggregateFile(file) {
  const acc = {
    fresh_input: 0,
    cache_read: 0,
    cache_creation: 0,
    output: 0,
    turns: 0,
    firstTs: null,
    lastTs: null,
  };
  for (const rec of readLines(file)) {
    const ts = rec.timestamp ? Date.parse(rec.timestamp) : NaN;
    if (!Number.isNaN(ts)) {
      if (acc.firstTs === null || ts < acc.firstTs) acc.firstTs = ts;
      if (acc.lastTs === null || ts > acc.lastTs) acc.lastTs = ts;
    }
    if (rec.type !== 'assistant') continue;
    const u = (rec.message && rec.message.usage) || null;
    if (!u) continue;
    acc.fresh_input += u.input_tokens || 0;
    acc.cache_read += u.cache_read_input_tokens || 0;
    acc.cache_creation += u.cache_creation_input_tokens || 0;
    acc.output += u.output_tokens || 0;
    acc.turns += 1;
  }
  return acc;
}

function costUnits(a) {
  return Math.round(a.output + a.fresh_input + a.cache_creation * 1.25 + a.cache_read * 0.1);
}

function wallMs(a) {
  return a.firstTs !== null && a.lastTs !== null ? a.lastTs - a.firstTs : 0;
}

// ---- role map from main transcript -----------------------------------------

// Collect every Agent spawn: { role, promptPrefix } in spawn order.
const agentSpawns = [];
for (const rec of readLines(mainTranscript)) {
  for (const blk of contentBlocks(rec)) {
    if (blk && blk.type === 'tool_use' && blk.name === 'Agent' && blk.input) {
      const role = blk.input.subagent_type || 'unknown';
      const prompt = (blk.input.prompt || '').trim();
      agentSpawns.push({ role, prompt });
    }
  }
}

function matchRole(firstPrompt) {
  if (!firstPrompt) return null;
  for (const s of agentSpawns) {
    if (!s.prompt) continue;
    // bidirectional prefix match tolerates truncation on either side
    if (firstPrompt.startsWith(s.prompt.slice(0, 50)) || s.prompt.startsWith(firstPrompt.slice(0, 50))) {
      return s.role;
    }
  }
  return null;
}

// ---- aggregate subagents + main --------------------------------------------

const byRole = new Map(); // role → { agg, agents }
function addToRole(role, agg) {
  if (!byRole.has(role)) {
    byRole.set(role, {
      fresh_input: 0,
      cache_read: 0,
      cache_creation: 0,
      output: 0,
      turns: 0,
      agents: 0,
      wallSumMs: 0,
    });
  }
  const r = byRole.get(role);
  r.fresh_input += agg.fresh_input;
  r.cache_read += agg.cache_read;
  r.cache_creation += agg.cache_creation;
  r.output += agg.output;
  r.turns += agg.turns;
  r.agents += 1;
  r.wallSumMs += wallMs(agg);
}

let subFiles = [];
if (fs.existsSync(subDir)) {
  try {
    subFiles = fs
      .readdirSync(subDir)
      .filter((f) => f.endsWith('.jsonl'))
      .map((f) => path.join(subDir, f));
  } catch (err) {
    subFiles = [];
  }
}

let unmatchedFiles = 0;
for (const f of subFiles) {
  const agg = aggregateFile(f);
  if (agg.turns === 0 && wallMs(agg) === 0) continue; // empty/aborted subagent
  // role: first user prompt of the subagent file
  let firstPrompt = '';
  for (const rec of readLines(f)) {
    if (rec.type === 'user') {
      firstPrompt = blocksText(contentBlocks(rec));
      break;
    }
  }
  const role = matchRole(firstPrompt);
  if (!role) unmatchedFiles += 1;
  addToRole(role || 'unknown', agg);
}

// orchestrator (main transcript) as its own row
const mainAgg = aggregateFile(mainTranscript);
addToRole('orchestrator', mainAgg);

// team-mode nesting caveat: Agent calls in main with no local subagent file
// (background / nested orchestrators write to a different session dir).
const localSubCount = subFiles.length;
const spawnedCount = agentSpawns.length;
const missingNested = Math.max(0, spawnedCount - localSubCount);

// ---- totals ----------------------------------------------------------------

const totals = {
  fresh_input: 0,
  cache_read: 0,
  cache_creation: 0,
  output: 0,
  turns: 0,
  agents: 0,
};
for (const r of byRole.values()) {
  totals.fresh_input += r.fresh_input;
  totals.cache_read += r.cache_read;
  totals.cache_creation += r.cache_creation;
  totals.output += r.output;
  totals.turns += r.turns;
  totals.agents += r.agents;
}
totals.cost_units = costUnits(totals);

// run wall-clock = span of the main transcript (first→last timestamp)
const runWallMs = wallMs(mainAgg);

// ---- formatting helpers -----------------------------------------------------

function fmtTokens(n) {
  if (n >= 1e6) return `${(n / 1e6).toFixed(1)}M`;
  if (n >= 1e3) return `${(n / 1e3).toFixed(1)}k`;
  return String(n);
}
function fmtDuration(ms) {
  if (!ms || ms < 0) return '-';
  const s = Math.round(ms / 1000);
  const m = Math.floor(s / 60);
  const rs = s % 60;
  return m > 0 ? `${m}m${rs.toString().padStart(2, '0')}s` : `${rs}s`;
}
function pad(s, w) {
  s = String(s);
  return s.length >= w ? s : s + ' '.repeat(w - s.length);
}
function padL(s, w) {
  s = String(s);
  return s.length >= w ? s : ' '.repeat(w - s.length) + s;
}

// rows sorted by cost_units desc
const roleRows = [...byRole.entries()]
  .map(([role, r]) => ({ role, ...r, cost_units: costUnits(r) }))
  .sort((a, b) => b.cost_units - a.cost_units);

// ---- write per-run markdown summary ----------------------------------------

const lines = [];
lines.push(`# Session telemetry — ${runId}`);
lines.push('');
lines.push(`- **skill**: ${skill}`);
lines.push(`- **session**: ${sessionId}`);
lines.push(`- **run wall-clock** (orchestrator span, includes idle): ${fmtDuration(runWallMs)}`);
lines.push(`- **agents**: ${totals.agents - 1} subagent(s) + 1 orchestrator`);
lines.push(`- **total cost units**: ${fmtTokens(totals.cost_units)}  _(output + fresh_in + cache_creation×1.25 + cache_read×0.1)_`);
lines.push('');
lines.push('## By role');
lines.push('');
const H = ['role', 'agents', 'wall(sum)', 'output', 'fresh_in', 'cache_cr', 'cache_rd', 'cost_units', '% cost'];
const W = [20, 7, 10, 9, 9, 9, 9, 11, 7];
lines.push(H.map((h, i) => pad(h, W[i])).join(' | '));
lines.push(W.map((w) => '-'.repeat(w)).join('-|-'));
for (const r of roleRows) {
  const pct = totals.cost_units ? `${((r.cost_units / totals.cost_units) * 100).toFixed(0)}%` : '-';
  lines.push(
    [
      pad(r.role, W[0]),
      padL(r.agents, W[1]),
      padL(fmtDuration(r.wallSumMs), W[2]),
      padL(fmtTokens(r.output), W[3]),
      padL(fmtTokens(r.fresh_input), W[4]),
      padL(fmtTokens(r.cache_creation), W[5]),
      padL(fmtTokens(r.cache_read), W[6]),
      padL(fmtTokens(r.cost_units), W[7]),
      padL(pct, W[8]),
    ].join(' | ')
  );
}
lines.push(W.map((w) => '-'.repeat(w)).join('-|-'));
lines.push(
  [
    pad('TOTAL', W[0]),
    padL(totals.agents, W[1]),
    padL(fmtDuration(runWallMs), W[2]),
    padL(fmtTokens(totals.output), W[3]),
    padL(fmtTokens(totals.fresh_input), W[4]),
    padL(fmtTokens(totals.cache_creation), W[5]),
    padL(fmtTokens(totals.cache_read), W[6]),
    padL(fmtTokens(totals.cost_units), W[7]),
    padL('100%', W[8]),
  ].join(' | ')
);
lines.push('');
lines.push('> wall(sum) is the sum of per-agent spans and reflects ACTIVE compute time; the');
lines.push('> run wall-clock is the orchestrator first→last span and INCLUDES human idle gaps.');
lines.push('> Subagents may run in parallel, so wall(sum) can exceed the run wall-clock.');
lines.push('> cache_read tokens are ~10% the cost of fresh input — kept separate on purpose.');
if (unmatchedFiles > 0) {
  lines.push('>');
  lines.push(`> ⚠ ${unmatchedFiles} subagent file(s) could not be mapped to a role (bucketed as "unknown").`);
}
if (missingNested > 0) {
  lines.push('>');
  lines.push(
    `> ⚠ team-mode caveat: ${missingNested} Agent spawn(s) in the main transcript have no local`
  );
  lines.push(
    `> subagent file (likely background / nested orchestrators in separate sessions) — their`
  );
  lines.push(`> tokens are NOT included in these totals. Sequential-mode runs are fully accurate.`);
}
lines.push('');

// ---- emit ------------------------------------------------------------------

const sessionsDir = path.join(outDir, 'sessions');
try {
  fs.mkdirSync(sessionsDir, { recursive: true });
  fs.writeFileSync(path.join(sessionsDir, `${runId}.md`), lines.join('\n'));
} catch (err) {
  skip(`could not write summary: ${err.message}`);
}

// human summary to stdout
console.log(lines.join('\n'));

// machine-readable totals for the skill to embed in skill-runs.jsonl
const jsonOut = {
  run_wall_ms: runWallMs,
  by_role: Object.fromEntries(
    roleRows.map((r) => [
      r.role,
      {
        agents: r.agents,
        output: r.output,
        fresh_input: r.fresh_input,
        cache_creation: r.cache_creation,
        cache_read: r.cache_read,
        cost_units: r.cost_units,
        wall_sum_ms: r.wallSumMs,
      },
    ])
  ),
  totals: {
    agents: totals.agents,
    output: totals.output,
    fresh_input: totals.fresh_input,
    cache_creation: totals.cache_creation,
    cache_read: totals.cache_read,
    cost_units: totals.cost_units,
    turns: totals.turns,
  },
  unmatched_subagents: unmatchedFiles,
  nested_agents_excluded: missingNested,
};
console.log(`TELEMETRY_JSON=${JSON.stringify(jsonOut)}`);
process.exit(0);
