#!/usr/bin/env bash
set -euo pipefail

# Count LLM/agentic coding costs from local session logs.
# Usage: ./agentic-coding-spend.sh [help] [--breakdown=day|week|month|year]
# Optional: COST_RATES_JSON='{"model":{"input":3,"output":15,"cache_write":3.75,"cache_read":0.30}}' FORCE_COLOR=1

show_help() {
  cat <<'EOF'
agentic-coding-spend — estimate local LLM/agentic coding spend

Usage:
  ./agentic-coding-spend.sh [options]
  ./agentic-coding-spend.sh help

Options:
  --breakdown=day|week|month|year
      Group spend by day, ISO week, month, or year. Default: month.

  -h, --help, help
      Show this help message.

Environment:
  COST_RATES_JSON
      JSON object for adding or overriding USD-per-1M-token rates.
      Example:
        COST_RATES_JSON='{"claude-sonnet-4-6":{"input":3,"output":15,"cache_write":3.75,"cache_read":0.30}}'

Scanned paths:
  ~/.pi/agent/sessions/**/*.jsonl
  ~/.claude/projects/**/*.jsonl
  ~/.codex/sessions/**/*.jsonl
  ~/.codex/**/*.jsonl
  ~/.config/codex/**/*.jsonl

Notes:
  Exact embedded costs are used when present. Otherwise costs are estimated
  from token usage and the built-in or overridden model rate table.
  This is API-equivalent spend, not necessarily your actual bill if you use
  a flat-rate subscription plan.
  Claude Code fast mode is priced only when logs include usage.speed=fast;
  transcripts without that field cannot be distinguished from standard mode.
EOF
}

breakdown="month"
for arg in "$@"; do
  case "$arg" in
    --breakdown=*) breakdown="${arg#*=}" ;;
    help|-h|--help)
      show_help
      exit 0
      ;;
    *) echo "unknown argument: $arg" >&2; echo "run './agentic-coding-spend.sh help' for usage" >&2; exit 2 ;;
  esac
done

case "$breakdown" in day|week|month|year) ;; *) echo "--breakdown must be day, week, month, or year" >&2; exit 2;; esac

python3 - "$breakdown" <<'PY'
import json, os, sys, glob, datetime, re
from collections import defaultdict

breakdown=sys.argv[1]
home=os.path.expanduser('~')

# USD per 1M tokens. Embedded costs in logs are preferred over these estimates.
RATES={
  # Anthropic (recent common CLI models; adjust with COST_RATES_JSON as needed)
  'claude-opus-4-7': {'input':5,'output':25,'cache_write_5m':6.25,'cache_write_1h':10,'cache_read':0.50},
  'claude-opus-4-6': {'input':5,'output':25,'cache_write_5m':6.25,'cache_write_1h':10,'cache_read':0.50},
  'claude-opus-4-5-20251101': {'input':5,'output':25,'cache_write_5m':6.25,'cache_write_1h':10,'cache_read':0.50},
  'claude-sonnet-4-6': {'input':3,'output':15,'cache_write_5m':3.75,'cache_write_1h':6,'cache_read':0.30},
  'claude-sonnet-4-5': {'input':3,'output':15,'cache_write_5m':3.75,'cache_write_1h':6,'cache_read':0.30},
  'claude-sonnet-4-5-20250929': {'input':3,'output':15,'cache_write_5m':3.75,'cache_write_1h':6,'cache_read':0.30},
  'claude-haiku-4-5-20251001': {'input':1,'output':5,'cache_write_5m':1.25,'cache_write_1h':2,'cache_read':0.10},
  # OpenAI / Codex examples
  'gpt-5.5': {'input':5,'output':30,'cache_write':0,'cache_read':0.50},
  'gpt-5': {'input':1.25,'output':10,'cache_write':0,'cache_read':0.125},
  'gpt-5.2-codex': {'input':1.75,'output':14,'cache_write':0,'cache_read':0.175},
  'gpt-5-codex': {'input':1.25,'output':10,'cache_write':0,'cache_read':0.125},
  'gpt-4.1': {'input':2,'output':8,'cache_write':0,'cache_read':0.50},
  'gpt-4.1-mini': {'input':0.4,'output':1.6,'cache_write':0,'cache_read':0.10},
  'o3': {'input':2,'output':8,'cache_write':0,'cache_read':0.50},
  'o4-mini': {'input':1.1,'output':4.4,'cache_write':0,'cache_read':0.275},
}
try:
    RATES.update(json.loads(os.environ.get('COST_RATES_JSON','{}')))
except Exception as e:
    print(f"warning: invalid COST_RATES_JSON: {e}", file=sys.stderr)

def parse_time(s, fallback_file=None):
    if isinstance(s, (int,float)):
        if s > 10_000_000_000: s/=1000
        return datetime.datetime.fromtimestamp(s, datetime.timezone.utc)
    if isinstance(s, str):
        t=s.replace('Z','+00:00')
        try: return datetime.datetime.fromisoformat(t)
        except Exception: pass
        m=re.search(r'(20\d\d-\d\d-\d\d)T(\d\d)[-:](\d\d)[-:](\d\d)', s)
        if m: return datetime.datetime.fromisoformat(f"{m.group(1)}T{m.group(2)}:{m.group(3)}:{m.group(4)}+00:00")
    if fallback_file:
        try: return datetime.datetime.fromtimestamp(os.path.getmtime(fallback_file), datetime.timezone.utc)
        except Exception: pass
    return None

def bucket(dt):
    if breakdown=='day': return dt.strftime('%Y-%m-%d')
    if breakdown=='week':
        y,w,_=dt.isocalendar(); return f'{y}-W{w:02d}'
    if breakdown=='month': return dt.strftime('%Y-%m')
    return dt.strftime('%Y')

def rate_for(model, speed=None):
    if not model: return None
    # Claude Code fast mode for Opus 4.6 is priced at 6x standard rates.
    # Anthropic lists this as $30 input / $150 output per MTok; prompt
    # caching multipliers apply on top of fast-mode input pricing.
    if speed == 'fast' and model.startswith('claude-opus-4-6'):
        return {'input':30,'output':150,'cache_write_5m':37.50,'cache_write_1h':60,'cache_read':3.00}
    if model in RATES: return RATES[model]
    # Claude Opus 4.5+ uses reduced $5/$25 pricing. Opus 4.5 launched
    # 2025-11-24, so treat dated 20251124+ snapshots as Opus 4.5 pricing.
    m=re.match(r'claude-opus-4-5-(\d{8})$', model)
    if m and m.group(1) >= '20251124':
        return {'input':5,'output':25,'cache_write_5m':6.25,'cache_write_1h':10,'cache_read':0.50}
    for k,v in sorted(RATES.items(), key=lambda kv: -len(kv[0])):
        if model.startswith(k): return v
    return None

def calc_cost(model, usage):
    # Prefer tool-provided precise cost.
    c=usage.get('cost') if isinstance(usage, dict) else None
    if isinstance(c, dict) and isinstance(c.get('total'), (int,float)): return float(c['total']), False
    if isinstance(c, (int,float)): return float(c), False
    speed=usage.get('speed') if isinstance(usage, dict) else None
    r=rate_for(model, speed)
    if not r: return 0.0, True
    inp=usage.get('input', usage.get('input_tokens',0)) or 0
    out=usage.get('output', usage.get('output_tokens',0)) or 0
    cr=usage.get('cacheRead', usage.get('cache_read_input_tokens', usage.get('cached_input_tokens',0))) or 0
    cw=usage.get('cacheWrite', usage.get('cache_creation_input_tokens',0)) or 0
    # Codex/OpenAI token_count events report input_tokens including cached_input_tokens.
    # Price those cached tokens at the read rate, not again at full input rate.
    if cr and usage.get('cached_input_tokens') is not None:
        inp=max(0, inp-cr)
    cc=usage.get('cache_creation') or {}
    cw5=0; cw1h=0
    if isinstance(cc, dict):
        cw5=cc.get('ephemeral_5m_input_tokens',0) or 0
        cw1h=cc.get('ephemeral_1h_input_tokens',0) or 0
    # If only aggregate cache_creation_input_tokens is present, assume 5m.
    if not cw5 and not cw1h and cw:
        cw5=cw
    cache_write_5m_rate=r.get('cache_write_5m', r.get('cache_write', r.get('input',0)))
    cache_write_1h_rate=r.get('cache_write_1h', r.get('cache_write', r.get('input',0)*2))
    return (inp*r.get('input',0)+out*r.get('output',0)+cr*r.get('cache_read',0)+cw5*cache_write_5m_rate+cw1h*cache_write_1h_rate)/1_000_000, True

rows=[]; seen=set(); warnings=set(); claude_opus46_without_speed=0
paths=[]
patterns=[
    '~/.pi/agent/sessions/**/*.jsonl',
    '~/.claude/projects/**/*.jsonl',
    '~/.codex/sessions/**/*.jsonl', '~/.codex/**/*.jsonl',
    '~/.config/codex/**/*.jsonl',
]
for p in patterns: paths.extend(glob.glob(os.path.expanduser(p), recursive=True))

for path in sorted(set(paths)):
    path_tool='pi' if '/.pi/' in path else 'claude' if '/.claude/' in path else 'codex' if 'codex' in path.lower() else 'unknown'
    session=os.path.basename(path).replace('.jsonl','')
    try: lines=open(path, errors='ignore').read().splitlines()
    except Exception: continue
    current_model='unknown'
    for line in lines:
        try: obj=json.loads(line)
        except Exception: continue
        payload=obj.get('payload') if isinstance(obj.get('payload'), dict) else {}
        if isinstance(payload, dict):
            current_model=payload.get('model') or payload.get('model_id') or current_model
            settings=(payload.get('collaboration_mode') or {}).get('settings') if isinstance(payload.get('collaboration_mode'), dict) else None
            if isinstance(settings, dict): current_model=settings.get('model') or current_model
        msg=obj.get('message') if isinstance(obj.get('message'), dict) else obj
        usage=msg.get('usage') if isinstance(msg, dict) else None
        model=msg.get('model') or obj.get('model') or obj.get('modelId') or current_model
        tool=obj.get('agent') or obj.get('tool') or msg.get('agent') or msg.get('tool') or path_tool
        rid=obj.get('requestId') or (msg.get('id') if isinstance(msg, dict) else None) or obj.get('uuid')

        # Native Codex CLI logs token usage as event_msg.payload.type=token_count.
        # Use last_token_usage because total_token_usage is cumulative within the session.
        if path_tool == 'codex' and obj.get('type') == 'event_msg' and payload.get('type') == 'token_count':
            info=payload.get('info') if isinstance(payload.get('info'), dict) else {}
            usage=info.get('last_token_usage') if isinstance(info.get('last_token_usage'), dict) else None
            tool='codex'
            rid=rid or f"{path}:{obj.get('timestamp')}:{json.dumps(usage, sort_keys=True)}"

        if not isinstance(usage, dict): continue
        if model in ('<synthetic>', 'synthetic'):
            continue
        rid=rid or f'{path}:{hash(line)}'
        key=(tool, rid, json.dumps(usage, sort_keys=True))
        if key in seen: continue
        seen.add(key)
        dt=parse_time(obj.get('timestamp') or (msg.get('timestamp') if isinstance(msg, dict) else None), path)
        if not dt: continue
        cost, estimated=calc_cost(model, usage)
        if estimated and not rate_for(model, usage.get('speed') if isinstance(usage, dict) else None): warnings.add(model)
        speed=usage.get('speed') if isinstance(usage, dict) else None
        if tool == 'claude' and model.startswith('claude-opus-4-6') and not speed:
            claude_opus46_without_speed+=1
        display_model=f'{model} ({speed})' if speed and speed != 'standard' else model
        rows.append({'bucket':bucket(dt), 'tool':tool, 'model':display_model, 'cost':cost, 'estimated':estimated, 'session':session})

agg=defaultdict(float); bytool=defaultdict(float); bymodel=defaultdict(float)
exact_total=0.0; estimated_total=0.0
for r in rows:
    agg[r['bucket']]+=r['cost']; bytool[(r['bucket'],r['tool'])]+=r['cost']; bymodel[(r['bucket'],r['tool'],r['model'])]+=r['cost']
    if r['estimated']:
        estimated_total+=r['cost']
    else:
        exact_total+=r['cost']

def money(v):
    return f' $ {v:,.2f}'

if os.environ.get('FORCE_COLOR'):
    color_enabled=True
elif os.environ.get('NO_COLOR'):
    color_enabled=False
else:
    color_enabled=sys.stdout.isatty()

def style(text, code):
    return f'\033[{code}m{text}\033[0m' if color_enabled else text

def dim(text):
    return style(text, '2')

def bold(text):
    return style(text, '1')

def green(text):
    return style(text, '32')

def cyan(text):
    return style(text, '36')

def magenta(text):
    return style(text, '35')

def yellow(text):
    return style(text, '33')

total=sum(r['cost'] for r in rows)
agent_w=max([5, len('TOTAL')] + [len(t) for _,t in bytool.keys()])
model_w=max([5] + [len(m) for _,_,m in bymodel.keys()])
period_w=max(10, len(breakdown))
cost_w=max(12, max([len(money(v)) for v in list(agg.values()) + list(bytool.values()) + list(bymodel.values())], default=0))

print()
print(f'💸 {bold("Total")}:  {green(money(total))} {dim(f"(exact: {money(exact_total)}, estimated: {money(estimated_total)})")}')
print(f'   across {bold(str(len(rows)))} billable log entries {dim(f"({breakdown} breakdown)")}')
print()
for b in sorted(agg):
    period_cell=f'{b:<{period_w}}'
    total_cell=f'{"TOTAL":<{agent_w}}'
    model_cell=f'{"":<{model_w}}'
    cost_cell=f'{money(agg[b]):>{cost_w}}'
    print(f'{cyan(period_cell)}  {bold(yellow(total_cell))}  {model_cell}  {bold(green(cost_cell))}  💰')
    for (bb,t),v in sorted(bytool.items()):
        if bb==b:
            period_cell=f'{"":<{period_w}}'
            agent_cell=f'{t:<{agent_w}}'
            model_cell=f'{"":<{model_w}}'
            cost_cell=f'{money(v):>{cost_w}}'
            print(f'{period_cell}  {magenta(agent_cell)}  {model_cell}  {green(cost_cell)}  🤖')
    for (bb,t,m),v in sorted(bymodel.items()):
        if bb==b:
            period_cell=f'{"":<{period_w}}'
            agent_cell=f'{t:<{agent_w}}'
            model_cell=f'{m:<{model_w}}'
            cost_cell=f'{money(v):>{cost_w}}'
            print(f'{period_cell}  {dim(agent_cell)}  {cyan(model_cell)}  {green(cost_cell)}  🧠')
    print()
if claude_opus46_without_speed:
    print(f'warning: {claude_opus46_without_speed} Claude Opus 4.6 log entries had no usage.speed field; any /fast usage in those entries is indistinguishable from standard mode and may be undercounted.', file=sys.stderr)
if warnings:
    print('Models without rates (counted as $0 unless logs had embedded costs):', file=sys.stderr)
    for m in sorted(warnings): print(f'  {m}', file=sys.stderr)
    print('Set COST_RATES_JSON to add/override USD-per-1M-token rates.', file=sys.stderr)
PY