import { readdirSync, readFileSync, writeFileSync, mkdirSync } from 'node:fs'
import { join, dirname } from 'node:path'

const BENCH = process.env.BENCH_DIR || process.cwd()
const capturesRoot = process.env.CAPTURE_DIR || join(BENCH, 'captures')
const RESULTS_PATH = process.env.RESULTS_PATH || join(BENCH, 'results', 'results.json')
const GATEWAY_ENVELOPE_TOKENS = Number(process.env.GATEWAY_ENVELOPE_TOKENS ?? 0)
const MAIN_MODEL_PREFIXES = (process.env.MAIN_MODEL_PREFIXES ?? 'claude-sonnet,claude-fable,claude-opus')
  .split(',')
  .map((p) => p.trim())
  .filter(Boolean)

const isMainModel = (model) => model != null && MAIN_MODEL_PREFIXES.some((p) => model.startsWith(p))

const lanes = {}
for (const label of readdirSync(capturesRoot).sort()) {
  const dir = join(capturesRoot, label)
  let files
  try {
    files = readdirSync(dir).filter((f) => f.endsWith('.json')).sort()
  } catch {
    continue
  }
  const requests = []
  for (const file of files) {
    const c = JSON.parse(readFileSync(join(dir, file), 'utf8'))
    const s = c.requestSummary
    const u = c.usage && Object.keys(c.usage).length > 0 ? c.usage : null
    requests.push({
      id: c.id,
      path: c.path,
      status: c.statusCode,
      reqModel: s ? s.model : null,
      servedModel: c.responseModel ?? null,
      systemChars: s ? s.systemChars : 0,
      systemBlocks: s ? s.systemBlocks : 0,
      toolCount: s ? s.toolCount : 0,
      toolsChars: s ? s.toolsChars : 0,
      msgCount: s ? s.messageCount : 0,
      msgChars: s ? s.messagesChars : 0,
      payloadChars: s ? s.systemChars + s.toolsChars + s.msgChars : 0,
      usage: u
        ? {
            in: u.input_tokens ?? 0,
            cacheWrite: u.cache_creation_input_tokens ?? 0,
            cacheRead: u.cache_read_input_tokens ?? 0,
            out: u.output_tokens ?? 0,
            thinking: u.output_tokens_details?.thinking_tokens ?? 0,
          }
        : null,
    })
  }
  lanes[label] = requests
}

const totalInputTokens = (u) => (u ? u.in + u.cacheWrite + u.cacheRead : null)

const report = {}
for (const [label, requests] of Object.entries(lanes)) {
  const main = requests.filter((r) => isMainModel(r.reqModel))
  const side = requests.filter((r) => r.reqModel != null && !isMainModel(r.reqModel))
  const probes = requests.filter((r) => r.reqModel == null)
  const first = main[0] ?? null
  report[label] = {
    requests: requests.length,
    mainModelRequests: main.length,
    sideCallRequests: side.length,
    bodylessProbes: probes.length,
    firstMainRequest: first
      ? {
          systemChars: first.systemChars,
          systemBlocks: first.systemBlocks,
          toolCount: first.toolCount,
          toolsChars: first.toolsChars,
          msgChars: first.msgChars,
          payloadChars: first.payloadChars,
          meteredInputTokens: totalInputTokens(first.usage),
          harnessTokensCalibrated:
            totalInputTokens(first.usage) != null
              ? totalInputTokens(first.usage) - GATEWAY_ENVELOPE_TOKENS
              : null,
        }
      : null,
    totals: {
      mainInput: main.reduce((acc, r) => acc + (totalInputTokens(r.usage) ?? 0), 0),
      mainCacheWrite: main.reduce((acc, r) => acc + (r.usage ? r.usage.cacheWrite : 0), 0),
      mainOutput: main.reduce((acc, r) => acc + (r.usage ? r.usage.out : 0), 0),
      sideInput: side.reduce((acc, r) => acc + (totalInputTokens(r.usage) ?? 0), 0),
    },
    perRequest: requests.map((r) => ({
      id: r.id,
      reqModel: r.reqModel,
      servedModel: r.servedModel,
      systemChars: r.systemChars,
      toolCount: r.toolCount,
      toolsChars: r.toolsChars,
      msgCount: r.msgCount,
      msgChars: r.msgChars,
      usage: r.usage,
    })),
  }
}

mkdirSync(dirname(RESULTS_PATH), { recursive: true })
writeFileSync(RESULTS_PATH, JSON.stringify(report, null, 2))

for (const [label, r] of Object.entries(report)) {
  const f = r.firstMainRequest
  console.log(
    `${label}: reqs=${r.requests} (main=${r.mainModelRequests} side=${r.sideCallRequests} probes=${r.bodylessProbes})` +
      (f
        ? ` | first main: sys=${f.systemChars}ch tools=${f.toolCount} (${f.toolsChars}ch) metered=${f.meteredInputTokens} calibrated=${f.harnessTokensCalibrated}`
        : ''),
  )
  console.log(
    `  totals: mainIn=${r.totals.mainInput} mainCacheWrite=${r.totals.mainCacheWrite} mainOut=${r.totals.mainOutput} sideIn=${r.totals.sideInput}`,
  )
}
console.log(`\nwritten: ${RESULTS_PATH}`)