'use strict';

// Per-session token-usage summary.
//
// Produces the numbers shown next to the model-name chip in the top bar:
//   { total, ctx, ctxWindow, ctxPct, exact, totalLabel, breakdown }
//
// Sources, in order of fidelity:
//   - Claude Code: ~/.claude/projects/<hash>/<id>.jsonl — every assistant entry
//     carries message.usage (input/output/cache_read/cache_creation tokens). We
//     sum them for `total` and take the last turn for `ctx`. EXACT.
//   - Codex: ~/.codex/sessions/.../rollout-*.jsonl — periodic `token_count`
//     events whose `total_token_usage` is already cumulative and which also
//     carry `model_context_window`. We read the LAST such event. EXACT.
//   - Everything else (Cursor/Gemini/Aider/OpenCode/Wall-E): estimate by
//     tokenizing the parsed message text. Marked `exact:false` so the UI can
//     render a `~` prefix.
//
// Claude files can be tens of MB and we recompute on every turn-finish, so the
// Claude path scans INCREMENTALLY: JSONL is append-only, so we keep a running
// total + byte offset per file and only read the bytes appended since last time.
// Codex events are cumulative, so we only ever read the file tail.

const fs = require('fs');

const DEFAULT_CTX_WINDOW = 200000;
// Codex: a cumulative token_count event normally sits within the last few KB, so
// read a small tail first and grow only if it isn't found (huge final tool blob).
const CODEX_TAIL_BYTES = 64 * 1024;
const CODEX_TAIL_BYTES_MAX = 4 * 1024 * 1024;

function _contextWindowOverride(value) {
  const n = Math.round(Number(value) || 0);
  return Number.isFinite(n) && n > 0 ? n : 0;
}

// Fallback context windows by model id (Codex carries its own window in-band, so
// this is only consulted for Claude + the estimate tier). Conservative defaults.
// `override` is the preferred path for provider/catalog metadata. The static map
// is only a cold-start safety net for transcripts seen before the registry cache
// has refreshed.
function ctxWindowForModel(modelId, override) {
  const fromRegistry = _contextWindowOverride(override);
  if (fromRegistry) return fromRegistry;
  const id = String(modelId || '').toLowerCase();
  if (!id) return DEFAULT_CTX_WINDOW;
  if (/\[1m\]|-1m\b|context-1m/.test(id)) return 1000000;
  // Cold-start defaults only. The live window is read from the model registry
  // when available, and per-session transcript evidence can still correct it.
  if (/^claude-(opus-4-8|fable-5|mythos-5)\b/.test(id)) return 1000000;
  if (/^claude-/.test(id)) return 200000;
  if (/^gpt-5|^o[34]\b/.test(id)) return 256000;
  if (/^gpt-4/.test(id)) return 128000;
  if (/^gemini/.test(id)) return 1000000;
  if (/^(deepseek|kimi|moonshot|qwen)/.test(id)) return 128000;
  return DEFAULT_CTX_WINDOW;
}

// Claude's usable context window is NOT a fixed property of the model id — Claude Code
// runs a per-session 200K vs 1M ("extended context") mode, so a static model→window map
// is fundamentally incomplete. Read the window from the session's own authoritative
// signals instead (the Claude analog of Codex's in-band model_context_window):
//   • AUTO compact_boundary `preTokens` — Claude Code auto-compacts near the limit, so
//     its recorded pre-compaction token count is its OWN tight measurement of the window.
//     This is authoritative: it can correct the static estimate UP *or* DOWN. (Manual
//     /compact is ignored — the user can compact at any size.)
//   • Largest single-turn prompt (input+cache_read+cache_creation) — a hard lower bound,
//     since the API cannot accept a prompt larger than the window. Only raises, never
//     lowers, the static default (it's a floor, not a measurement of the true window).
// The evidence is classified to the nearest standard window with a tolerance, so an
// auto-compact that overshoots a round 1M (e.g. 1,008,316) reads as the 1M tier instead
// of rounding up to 2M. A fresh session with no evidence falls back to the model map.
const _CTX_WINDOW_TIERS = [200000, 1000000, 2000000, 5000000, 10000000];
const _CTX_WINDOW_TOLERANCE = 1.05;
function _snapToCtxTier(observed) {
  const obs = Math.max(0, Number(observed) || 0);
  for (const tier of _CTX_WINDOW_TIERS) { if (obs <= tier * _CTX_WINDOW_TOLERANCE) return tier; }
  return obs;
}
function _effectiveCtxWindow(baseWindow, evidence) {
  const base = Number(baseWindow) > 0 ? Number(baseWindow) : DEFAULT_CTX_WINDOW;
  const ev = evidence || {};
  const autoAt = Math.max(0, Number(ev.autoCompactAt) || 0);
  const maxPrompt = Math.max(0, Number(ev.maxPrompt) || 0);
  // Authoritative auto-compact measurement wins outright (up or down).
  if (autoAt > 0) return _snapToCtxTier(autoAt);
  // Otherwise a prompt larger than the static estimate proves the estimate is too small.
  if (maxPrompt > base) return Math.max(_snapToCtxTier(maxPrompt), base);
  return base;
}

// Compact human label: 980 -> "980", 1234 -> "1.2k", 1_200_000 -> "1.2M".
function formatTokens(n) {
  const v = Math.max(0, Math.round(Number(n) || 0));
  if (v < 1000) return String(v);
  if (v < 1_000_000) {
    const k = v / 1000;
    return (k < 10 ? k.toFixed(1) : Math.round(k)) + 'k';
  }
  const m = v / 1_000_000;
  return (m < 10 ? m.toFixed(1) : Math.round(m)) + 'M';
}

function _isClaude(provider) { return String(provider || '').toLowerCase() === 'claude'; }
function _isCodex(provider) { return String(provider || '').toLowerCase() === 'codex'; }

// Match a brace-balanced JSON object beginning at the `{` that follows `key` at
// or after `from`. Returns the parsed object or null. Used for Codex events,
// whose token objects are small.
function _matchObjectAfter(text, key, from = 0) {
  const at = text.indexOf(key, from);
  if (at === -1) return null;
  let i = text.indexOf('{', at);
  if (i === -1) return null;
  const start = i;
  let depth = 0;
  let inStr = false;
  let esc = false;
  for (; i < text.length; i++) {
    const ch = text[i];
    if (inStr) {
      if (esc) esc = false;
      else if (ch === '\\') esc = true;
      else if (ch === '"') inStr = false;
      continue;
    }
    if (ch === '"') inStr = true;
    else if (ch === '{') depth++;
    else if (ch === '}') {
      depth--;
      if (depth === 0) {
        try { return JSON.parse(text.slice(start, i + 1)); }
        catch { return null; }
      }
    }
  }
  return null;
}

const _CODEX_TOKEN_COUNT_PAYLOAD_RE = /"payload"\s*:\s*\{\s*"type"\s*:\s*"token_count"/;

function _codexSummaryFromTokenCountLines(text) {
  if (!text) return null;
  let end = text.length;
  while (end > 0) {
    const start = text.lastIndexOf('\n', end - 1);
    const line = text.slice(start + 1, end).trim();
    if (line && _CODEX_TOKEN_COUNT_PAYLOAD_RE.test(line)) {
      const summary = _codexSummaryFromTokenCountLine(line);
      if (summary) return summary;
    }
    if (start === -1) break;
    end = start;
  }
  return null;
}

function _codexSummaryFromTokenCountLine(line) {
  if (!line) return null;
  try {
    const obj = JSON.parse(line);
    const payload = obj && obj.payload;
    const info = payload && payload.info;
    if (!payload || payload.type !== 'token_count' || !info) return null;
    if (!info.total_token_usage || typeof info.total_token_usage !== 'object') return null;
    return {
      total: info.total_token_usage,
      last: info.last_token_usage && typeof info.last_token_usage === 'object' ? info.last_token_usage : null,
      window: Number(info.model_context_window) || 0,
    };
  } catch {
    return null;
  }
}

function _codexCurrentContextTokens(summary) {
  const last = summary && summary.last;
  if (last && typeof last === 'object') {
    // Codex emits a post-compaction token_count where input/output are zero but
    // total_tokens is the new compacted context size. Treat 0 as valid; do not
    // fall through to cumulative total_token_usage or the badge pins at 100%.
    const lastTotal = Number(last.total_tokens);
    if (Number.isFinite(lastTotal) && lastTotal > 0) return lastTotal;
    const lastInput = Number(last.input_tokens);
    if (Number.isFinite(lastInput) && lastInput > 0) return lastInput;
    const componentTotal = (Number(last.input_tokens) || 0)
      + (Number(last.output_tokens) || 0)
      + (Number(last.reasoning_output_tokens) || 0);
    return Math.max(0, componentTotal);
  }
  const total = summary && summary.total;
  return Number(total && total.input_tokens) || 0;
}

// --- Claude: sum message.usage across assistant entries ------------------------
//
// Extract usage numbers from a single JSONL line WITHOUT JSON.parsing the whole
// (possibly multi-MB) line: the four counters are uniquely-named keys, so a field
// regex is both faster and immune to giant tool-result blobs.
const _CLAUDE_FIELD_RE = {
  input: /"input_tokens":\s*(\d+)/,
  output: /"output_tokens":\s*(\d+)/,
  cacheRead: /"cache_read_input_tokens":\s*(\d+)/,
  cacheCreate: /"cache_creation_input_tokens":\s*(\d+)/,
};

function _claudeUsageFromLine(line) {
  // Only assistant entries carry message.usage; cheap guard before regex work.
  if (line.indexOf('"usage"') === -1) return null;
  const num = (re) => { const m = line.match(re); return m ? parseInt(m[1], 10) || 0 : 0; };
  const u = {
    input: num(_CLAUDE_FIELD_RE.input),
    output: num(_CLAUDE_FIELD_RE.output),
    cacheRead: num(_CLAUDE_FIELD_RE.cacheRead),
    cacheCreate: num(_CLAUDE_FIELD_RE.cacheCreate),
  };
  if (!(u.input || u.output || u.cacheRead || u.cacheCreate)) return null;
  return u;
}

// Scan a chunk of appended JSONL text, accumulating into `acc`. Returns the
// last full assistant turn seen (for ctx), or acc.lastTurn unchanged.
function _scanClaudeChunk(text, acc) {
  let start = 0;
  while (start < text.length) {
    let end = text.indexOf('\n', start);
    if (end === -1) end = text.length;
    const line = text.slice(start, end);
    start = end + 1;
    if (!line) continue;
    // Authoritative window signal: Claude Code's AUTO compact_boundary records the
    // pre-compaction token count near the context limit. Track the largest. Manual
    // /compact is ignored (the user can compact at any size, so it says nothing about
    // the window). These are system lines without usage, so they never double-count.
    if (line.indexOf('compact_boundary') !== -1) {
      if (/"trigger"\s*:\s*"auto"/.test(line)) {
        const pt = line.match(/"preTokens"\s*:\s*(\d+)/);
        if (pt) { const v = Number(pt[1]) || 0; if (v > (acc.autoCompactAt || 0)) acc.autoCompactAt = v; }
      }
      continue;
    }
    const u = _claudeUsageFromLine(line);
    if (!u) continue;
    acc.input += u.input;
    acc.output += u.output;
    acc.cacheRead += u.cacheRead;
    acc.cacheCreate += u.cacheCreate;
    // Current context occupancy ~ the most recent turn's prompt + reply.
    acc.lastTurn = u.input + u.cacheRead + u.cacheCreate + u.output;
    // High-water prompt size (input side only) — a hard lower bound on the real
    // context window, used to self-correct a stale static window estimate.
    const prompt = u.input + u.cacheRead + u.cacheCreate;
    if (prompt > (acc.maxPrompt || 0)) acc.maxPrompt = prompt;
  }
}

// Read the last `readLen` bytes of a file as utf8.
function _readFileTail(jsonlPath, fileSize, readLen) {
  const len = Math.min(readLen, fileSize);
  const fd = fs.openSync(jsonlPath, 'r');
  try {
    const buf = Buffer.alloc(len);
    fs.readSync(fd, buf, 0, len, Math.max(0, fileSize - len));
    return buf.toString('utf8');
  } finally { fs.closeSync(fd); }
}

// Read a small tail first; grow (×8) only when the window didn't contain a
// token_count event. lastIndexOf+forward-scan means a window that includes the
// key includes its (EOF-ward) object, so a partial leading line never corrupts.
function _codexSummaryGrowing(jsonlPath, fileSize) {
  let readLen = CODEX_TAIL_BYTES;
  for (;;) {
    let text;
    try { text = _readFileTail(jsonlPath, fileSize, readLen); } catch { return null; }
    const c = _codexSummaryFromTail(text);
    if (c) return c;
    if (readLen >= fileSize || readLen >= CODEX_TAIL_BYTES_MAX) return null;
    readLen = Math.min(readLen * 8, CODEX_TAIL_BYTES_MAX, fileSize);
  }
}

// --- Codex: read the last cumulative token_count event -------------------------
function _codexSummaryFromTail(text) {
  const structured = _codexSummaryFromTokenCountLines(text);
  if (structured) return structured;

  // Find the last token_count event; its total_token_usage is cumulative.
  const lastIdx = text.lastIndexOf('"total_token_usage"');
  if (lastIdx === -1) return null;
  const total = _matchObjectAfter(text, '"total_token_usage"', lastIdx);
  if (!total) return null;
  const last = _matchObjectAfter(text, '"last_token_usage"', lastIdx);
  const winM = text.slice(lastIdx).match(/"model_context_window":\s*(\d+)/);
  const window = winM ? parseInt(winM[1], 10) || 0 : 0;
  return { total, last, window };
}

// --- Estimate tier -------------------------------------------------------------
let _encoder = null;
let _encoderTried = false;
function _estimateTokens(str) {
  const text = String(str || '');
  if (!text) return 0;
  if (!_encoderTried) {
    _encoderTried = true;
    try { _encoder = require('gpt-tokenizer'); }
    catch { _encoder = null; }
  }
  if (_encoder && typeof _encoder.encode === 'function') {
    try { return _encoder.encode(text).length; } catch { /* fall through */ }
  }
  // Heuristic fallback (~4 chars/token) when the tokenizer dep isn't present.
  return Math.ceil(text.length / 4);
}

function estimateFromMessages(messages, model, opts = {}) {
  const list = Array.isArray(messages) ? messages : [];
  let total = 0;
  for (let k = 0; k < list.length; k++) {
    const m = list[k];
    // Structured-capture rows (tool calls/results, reasoning, patches) carry
    // capped derived text; counting them would shift the estimate tier for
    // sessions imported after the capture upgrade. Estimate from the same
    // user/assistant prose the pre-capture importer produced.
    if (m && m.role === 'system' && m.metadata && m.metadata.kind) continue;
    total += _estimateTokens(m && (m.text || m.content));
  }
  const ctxWindow = ctxWindowForModel(model, opts.contextWindow || opts.ctxWindow);
  return _finalize({ total, ctx: total, ctxWindow, exact: false,
    breakdown: { input: 0, output: 0, cacheRead: 0, cacheCreate: 0, reasoning: 0, estimated: total } });
}

// --- Shaping -------------------------------------------------------------------
function _finalize({ total, ctx, ctxWindow, exact, breakdown, cost }) {
  const window = ctxWindow > 0 ? ctxWindow : DEFAULT_CTX_WINDOW;
  const ctxClamped = Math.max(0, Math.min(Number(ctx) || 0, window));
  const ctxPct = window > 0 ? Math.round((ctxClamped / window) * 100) : 0;
  return {
    total: Math.max(0, Math.round(Number(total) || 0)),
    ctx: Math.max(0, Math.round(Number(ctx) || 0)),
    ctxWindow: window,
    ctxPct,
    exact: !!exact,
    totalLabel: formatTokens(total),
    breakdown: breakdown || null,
    cost: (cost === undefined ? null : cost),
  };
}

// --- Per-file cache (incremental for Claude, tail for Codex) -------------------
// keyed by jsonlPath -> { size, mtimeMs, offset, acc, summary }
// Bounded LRU: a long-lived primary sees many transcripts (incl. rotated/compacted
// ones whose paths never recur), so an unbounded Map would leak. Map preserves
// insertion order; we evict the oldest key and refresh recency on hit.
const _cache = new Map();
const _CACHE_MAX = 500;

function _cacheSet(key, val) {
  if (_cache.has(key)) _cache.delete(key);
  _cache.set(key, val);
  if (_cache.size > _CACHE_MAX) {
    const oldest = _cache.keys().next().value;
    if (oldest !== undefined) _cache.delete(oldest);
  }
}

function _statSafe(p) {
  try { return fs.statSync(p); } catch { return null; }
}

// `statHint` ({size, mtimeMs}) lets a caller that already freshly stat'd the file
// (e.g. _sessionPayload via _activeSessionFileInfo) skip a redundant statSync.
function _computeFromFile(provider, jsonlPath, model, statHint, opts = {}) {
  const st = (statHint && Number.isFinite(statHint.size) && Number.isFinite(statHint.mtimeMs))
    ? statHint
    : _statSafe(jsonlPath);
  if (!st || (typeof st.isFile === 'function' && !st.isFile())) return null;
  const sizeKey = st.size;
  const mtimeKey = st.mtimeMs;
  const cached = _cache.get(jsonlPath);
  if (cached && cached.size === sizeKey && cached.mtimeMs === mtimeKey && cached.summary) {
    _cacheSet(jsonlPath, cached); // refresh LRU recency
    return cached.summary;
  }

  if (_isCodex(provider)) {
    const c = _codexSummaryGrowing(jsonlPath, st.size);
    if (!c) return null;
    const tt = c.total || {};
    // "total" = genuinely-new tokens this session (exclude cached re-reads, which
    // are re-counted every turn and would balloon into the billions). Codex's
    // input_tokens INCLUDES cached_input_tokens, so subtract them out.
    const cachedIn = Number(tt.cached_input_tokens) || 0;
    const nonCachedIn = Math.max(0, (Number(tt.input_tokens) || 0) - cachedIn);
    const total = nonCachedIn + (Number(tt.output_tokens) || 0) + (Number(tt.reasoning_output_tokens) || 0);
    const ctx = _codexCurrentContextTokens(c);
    const summary = _finalize({
      total, ctx, ctxWindow: c.window || ctxWindowForModel(model, opts.contextWindow || opts.ctxWindow), exact: true,
      breakdown: {
        input: Number(tt.input_tokens) || 0,
        output: Number(tt.output_tokens) || 0,
        cacheRead: Number(tt.cached_input_tokens) || 0,
        cacheCreate: 0,
        reasoning: Number(tt.reasoning_output_tokens) || 0,
      },
    });
    _cacheSet(jsonlPath, { size: sizeKey, mtimeMs: mtimeKey, summary });
    return summary;
  }

  // Claude (and any other provider whose transcript matches the Claude usage
  // shape): incremental append-scan.
  let acc;
  let offset;
  if (cached && cached.acc && st.size >= cached.offset) {
    acc = cached.acc;
    offset = cached.offset;
  } else {
    acc = { input: 0, output: 0, cacheRead: 0, cacheCreate: 0, lastTurn: 0, maxPrompt: 0, autoCompactAt: 0 };
    offset = 0;
  }
  if (st.size > offset) {
    try {
      const fd = fs.openSync(jsonlPath, 'r');
      try {
        const len = st.size - offset;
        const buf = Buffer.alloc(len);
        fs.readSync(fd, buf, 0, len, offset);
        _scanClaudeChunk(buf.toString('utf8'), acc);
      } finally { fs.closeSync(fd); }
    } catch { return null; }
  }
  // "total" = genuinely-new tokens (non-cached input + new cache-creation +
  // output). cache_READ is the context re-read every turn; including it balloons
  // the count into the billions, so it lives in the breakdown/tooltip only.
  const total = acc.input + acc.cacheCreate + acc.output;
  if (acc.input + acc.output + acc.cacheRead + acc.cacheCreate === 0) {
    // No usage found (e.g. not actually a Claude transcript). Let caller fall
    // back to estimation rather than showing 0.
    _cacheSet(jsonlPath, { size: sizeKey, mtimeMs: mtimeKey, acc, offset: st.size, summary: null });
    return null;
  }
  const summary = _finalize({
    total,
    ctx: acc.lastTurn || total,
    ctxWindow: _effectiveCtxWindow(ctxWindowForModel(model, opts.contextWindow || opts.ctxWindow), { maxPrompt: acc.maxPrompt, autoCompactAt: acc.autoCompactAt }),
    exact: true,
    breakdown: {
      input: acc.input, output: acc.output,
      cacheRead: acc.cacheRead, cacheCreate: acc.cacheCreate, reasoning: 0,
    },
  });
  _cacheSet(jsonlPath, { size: sizeKey, mtimeMs: mtimeKey, acc, offset: st.size, summary });
  return summary;
}

/**
 * Compute the token-usage summary for a session.
 * @param {object} o
 * @param {string} o.provider  agent type: 'claude' | 'codex' | 'cursor' | ...
 * @param {string} [o.jsonlPath]  transcript path (Claude/Codex authoritative)
 * @param {string} [o.model]  model id, for the ctx-window fallback
 * @param {Array}  [o.messages]  parsed messages, for the estimate tier
 * @param {number} [o.contextWindow] provider/catalog context-window override
 * @param {object} [o.stat]  a fresh {size, mtimeMs} the caller already stat'd,
 *   to skip a redundant statSync (must be current — used as the cache key + read range)
 * @returns {object|null} summary, or null when nothing is computable
 */
function computeSessionTokenUsage({ provider, jsonlPath, model, messages, stat, contextWindow, ctxWindow } = {}) {
  const opts = { contextWindow: contextWindow || ctxWindow };
  if (jsonlPath && (_isClaude(provider) || _isCodex(provider) || !provider)) {
    const fromFile = _computeFromFile(provider, jsonlPath, model, stat, opts);
    if (fromFile) return fromFile;
  }
  if (Array.isArray(messages) && messages.length) {
    return estimateFromMessages(messages, model, opts);
  }
  return null;
}

function _resetCacheForTests() { _cache.clear(); }
function _cacheSizeForTests() { return _cache.size; }

module.exports = {
  computeSessionTokenUsage,
  estimateFromMessages,
  formatTokens,
  ctxWindowForModel,
  _effectiveCtxWindow,
  _resetCacheForTests,
  _cacheSizeForTests,
  _CACHE_MAX,
};
