#!/usr/bin/env sh
# token-usage.sh — 讀出單一 CLI 這次工作的 token 用量。
# 用法:
#   token-usage.sh <cli> [session-id]   # cli = claude|codex|copilot|antigravity|kiro
# 輸出: 一行「<input>(tab)<output>」。任何來源讀不到就印「N/A(tab)N/A」並 exit 0。
# 護欄: 沒給 cli 或 cli 名稱不認得，回傳 2。
#
# 各 CLI 的取得方式（原本寫在 worklog 的 SKILL.md，現在收在這裡）:
#   claude       加總 transcript JSONL（`$CLAUDE_CONFIG_DIR` 或 `~/.claude` 底下的
#                `projects/**/*.jsonl`）每筆訊息的 `usage.input_tokens` 與
#                `usage.output_tokens`。也可以改讀 `claude -p --output-format json`
#                回應裡的 `usage` 欄位。
#   codex        讀 `~/.codex/sessions/**` 的 token 計數欄位；session 內打 `/status`
#                也看得到同一組數字。
#   copilot      只有 session 內的 `/usage` 看得到，沒有可讀檔案 → N/A。
#   antigravity  只顯示在 UI，沒有可讀檔案 → N/A。
#   kiro         沒有公開來源 → N/A。
#
# 有給 session id 就優先取檔名含該 id 的 session 檔，因為日誌的花費時間也是按同一個
# session id 取的（session-timer.sh report {sid}），兩個數字必須描述同一個工作階段。
# 沒給 session id 才退回「最後修改的那一份」——那份可能是別的工作階段，或同一階段的
# sub agent 紀錄，數字只能當粗估。
set -u

na() { printf 'N/A\tN/A\n'; exit 0; }

cli="${1:-}"
sid="${2:-}"
if [ -z "$cli" ]; then
  echo 'usage: token-usage.sh <cli> [session-id]' >&2
  exit 2
fi

newest() { # <dir> <name-pattern> -> 最後修改的檔案路徑；找不到回非零
  [ -d "$1" ] || return 1
  f=$(find "$1" -type f -name "$2" -exec ls -t {} + 2>/dev/null | head -1)
  [ -n "$f" ] || return 1
  printf '%s\n' "$f"
}

pick() { # <dir> -> 先找 session id 相符的檔案，找不到才退回最後修改的
  if [ -n "$sid" ]; then
    if f=$(newest "$1" "*$sid*.jsonl"); then printf '%s\n' "$f"; return 0; fi
    echo "warn: 找不到 session $sid 的 transcript，改用最後修改的那一份，數字可能不是本階段的。" >&2
  fi
  newest "$1" '*.jsonl'
}

sum_field() { # <file> <欄位名> -> 該欄位所有出現值的總和；沒有值就印空字串
  grep -o "\"$2\"[[:space:]]*:[[:space:]]*[0-9][0-9]*" "$1" 2>/dev/null \
    | grep -o '[0-9][0-9]*$' \
    | awk '{ s += $1 } END { if (NR > 0) print s }'
}

case "$cli" in
  claude)
    src=$(pick "${CLAUDE_CONFIG_DIR:-$HOME/.claude}/projects") || na ;;
  codex)
    src=$(pick "$HOME/.codex/sessions") || na ;;
  copilot|antigravity|kiro)
    na ;;
  *)
    echo "unknown cli: $cli" >&2
    exit 2 ;;
esac

[ -r "$src" ] || na
# input 要把快取的部分加回來：Claude 的 usage 只在 input_tokens 記未快取的那幾個 token，
# 真正的輸入量落在 cache_creation_input_tokens 與 cache_read_input_tokens。只加第一欄，
# 報表會出現輸入幾百、輸出幾十萬的假數字。
input=$(sum_field "$src" input_tokens)
cache_new=$(sum_field "$src" cache_creation_input_tokens)
cache_hit=$(sum_field "$src" cache_read_input_tokens)
output=$(sum_field "$src" output_tokens)
[ -n "$input" ] && [ -n "$output" ] || na
input=$((input + ${cache_new:-0} + ${cache_hit:-0}))
printf '%s\t%s\n' "$input" "$output"