#!/usr/bin/env python3
"""
cost-tracker — aggregate Claude Code session token usage + estimated cost.

Reads ~/.claude/projects/<dir>/<sessionId>.jsonl files, sums tokens per
day per model, multiplies by published Anthropic pricing.

Output:
  - HQ-Vault/_System/CostLog.ndjson  (one row per day per model, append-only)
  - HQ-Vault/_System/CostLog-summary.md  (last-30d table for human read)

Designed to run daily on Rocky's local machine where Claude Code lives.
Schedule via Windows Task Scheduler:
  schtasks /Create /TN "VaultCostTracker" /SC DAILY /ST 22:30 \
    /TR "python C:\\Users\\mclou\\rateright-growth-deploy\\scripts\\cost-tracker.py"

Usage:
  python scripts/cost-tracker.py                  # daily incremental update
  python scripts/cost-tracker.py --backfill 30    # backfill last N days
  python scripts/cost-tracker.py --dry-run        # preview without writing
"""
import argparse
import json
import os
import sys
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path

# Anthropic published pricing (USD per 1M tokens) as of 2026-01.
# Update when prices change.
PRICING = {
    "claude-opus-4-7":   {"input": 15.0, "output": 75.0, "cache_write": 18.75, "cache_read": 1.50},
    "claude-opus-4-6":   {"input": 15.0, "output": 75.0, "cache_write": 18.75, "cache_read": 1.50},
    "claude-sonnet-4-6": {"input":  3.0, "output": 15.0, "cache_write":  3.75, "cache_read": 0.30},
    "claude-haiku-4-5":  {"input":  1.0, "output":  5.0, "cache_write":  1.25, "cache_read": 0.10},
    # Older models retained for backfill
    "claude-3-5-sonnet": {"input":  3.0, "output": 15.0, "cache_write":  3.75, "cache_read": 0.30},
    "claude-3-5-haiku":  {"input":  0.80,"output":  4.0, "cache_write":  1.0, "cache_read": 0.08},
    "claude-3-opus":     {"input": 15.0, "output": 75.0, "cache_write": 18.75, "cache_read": 1.50},
}

DEFAULT_PRICING = PRICING["claude-sonnet-4-6"]  # fallback

CLAUDE_PROJECTS = Path.home() / ".claude" / "projects"
VAULT_ROOT = Path("HQ-Vault")
COST_LOG = VAULT_ROOT / "_System" / "CostLog.ndjson"
COST_SUMMARY = VAULT_ROOT / "_System" / "CostLog-summary.md"


def model_pricing(model: str) -> dict:
    if not model:
        return DEFAULT_PRICING
    # match longest prefix
    for k in sorted(PRICING.keys(), key=len, reverse=True):
        if model.startswith(k):
            return PRICING[k]
    return DEFAULT_PRICING


def parse_session_files(roots: list[Path]) -> dict:
    """Return {(date_iso, project, model): {tokens..., msgs}}."""
    agg: dict = defaultdict(lambda: {
        "input_tokens": 0,
        "output_tokens": 0,
        "cache_creation_input_tokens": 0,
        "cache_read_input_tokens": 0,
        "messages": 0,
    })

    for project_dir in roots:
        if not project_dir.is_dir():
            continue
        project = project_dir.name
        for jsonl in project_dir.rglob("*.jsonl"):
            try:
                with open(jsonl, encoding="utf-8", errors="ignore") as f:
                    for line in f:
                        try:
                            d = json.loads(line)
                        except json.JSONDecodeError:
                            continue
                        msg = d.get("message")
                        if not isinstance(msg, dict):
                            continue
                        usage = msg.get("usage")
                        if not isinstance(usage, dict):
                            continue
                        ts = d.get("timestamp") or d.get("ts") or d.get("createdAt")
                        if not ts:
                            continue
                        try:
                            day = ts[:10]  # YYYY-MM-DD
                        except (TypeError, IndexError):
                            continue
                        model = msg.get("model", "unknown")
                        key = (day, project, model)
                        agg[key]["input_tokens"] += usage.get("input_tokens", 0) or 0
                        agg[key]["output_tokens"] += usage.get("output_tokens", 0) or 0
                        agg[key]["cache_creation_input_tokens"] += usage.get("cache_creation_input_tokens", 0) or 0
                        agg[key]["cache_read_input_tokens"] += usage.get("cache_read_input_tokens", 0) or 0
                        agg[key]["messages"] += 1
            except OSError:
                continue
    return agg


def cost_usd(tokens: dict, model: str) -> float:
    p = model_pricing(model)
    return (
        tokens["input_tokens"] / 1_000_000 * p["input"]
        + tokens["output_tokens"] / 1_000_000 * p["output"]
        + tokens["cache_creation_input_tokens"] / 1_000_000 * p["cache_write"]
        + tokens["cache_read_input_tokens"] / 1_000_000 * p["cache_read"]
    )


def existing_keys(path: Path) -> set:
    if not path.exists():
        return set()
    seen = set()
    for line in path.read_text(encoding="utf-8", errors="ignore").splitlines():
        if not line.strip():
            continue
        try:
            r = json.loads(line)
            seen.add((r.get("date"), r.get("project"), r.get("model")))
        except json.JSONDecodeError:
            continue
    return seen


def write_summary(rows: list[dict], path: Path, days: int = 30) -> None:
    rows_sorted = sorted(rows, key=lambda r: r["date"], reverse=True)[: days * 5]
    by_date: dict = defaultdict(lambda: {"cost": 0.0, "msgs": 0, "models": set()})
    for r in rows_sorted:
        b = by_date[r["date"]]
        b["cost"] += r["cost_usd"]
        b["msgs"] += r["messages"]
        b["models"].add(r["model"])

    lines = [
        "---",
        "title: CostLog Summary",
        "type: system",
        f"updated: {datetime.now(timezone.utc).strftime('%Y-%m-%d')}",
        "---",
        "",
        "# CostLog Summary (last 30 days)",
        "",
        "Daily aggregate USD cost from Claude Code session logs. Source: `CostLog.ndjson`.",
        "",
        "| Date | $ USD | Messages | Models |",
        "|---|---:|---:|---|",
    ]
    total = 0.0
    for date in sorted(by_date.keys(), reverse=True)[:days]:
        b = by_date[date]
        total += b["cost"]
        models_str = ", ".join(sorted(b["models"]))
        lines.append(f"| {date} | ${b['cost']:.2f} | {b['msgs']} | {models_str} |")
    lines.append(f"| **Total** | **${total:.2f}** | | |")
    lines.append("")
    path.write_text("\n".join(lines), encoding="utf-8")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--dry-run", action="store_true", help="don't write files")
    ap.add_argument("--backfill", action="store_true", help="rewrite all (not just new days)")
    args = ap.parse_args()

    if not CLAUDE_PROJECTS.is_dir():
        print(f"cost-tracker: {CLAUDE_PROJECTS} not found — run on Rocky's local machine", file=sys.stderr)
        return 1

    roots = [d for d in CLAUDE_PROJECTS.iterdir() if d.is_dir()]
    print(f"cost-tracker: scanning {len(roots)} project(s) under {CLAUDE_PROJECTS}")

    agg = parse_session_files(roots)
    print(f"cost-tracker: aggregated {len(agg)} (date, project, model) rows")

    seen = set() if args.backfill else existing_keys(COST_LOG)
    new_rows = []
    for (day, project, model), tokens in sorted(agg.items()):
        if (day, project, model) in seen:
            continue
        row = {
            "date": day,
            "project": project,
            "model": model,
            **tokens,
            "cost_usd": round(cost_usd(tokens, model), 4),
        }
        new_rows.append(row)

    if args.dry_run:
        for r in new_rows[:10]:
            print(json.dumps(r))
        print(f"... ({len(new_rows)} new rows total) [dry run, not written]")
        return 0

    COST_LOG.parent.mkdir(parents=True, exist_ok=True)
    if args.backfill and COST_LOG.exists():
        COST_LOG.write_text("", encoding="utf-8")  # clear
    with COST_LOG.open("a", encoding="utf-8") as f:
        for row in new_rows:
            f.write(json.dumps(row) + "\n")

    # Re-read all rows for summary
    all_rows = []
    for line in COST_LOG.read_text(encoding="utf-8", errors="ignore").splitlines():
        if not line.strip():
            continue
        try:
            all_rows.append(json.loads(line))
        except json.JSONDecodeError:
            pass
    write_summary(all_rows, COST_SUMMARY)

    total = sum(r["cost_usd"] for r in new_rows)
    print(f"cost-tracker: wrote {len(new_rows)} new rows (${total:.2f}); summary at {COST_SUMMARY}")
    return 0


if __name__ == "__main__":
    sys.exit(main())