#!/usr/bin/env python3
"""Estimate token usage inside the rolling 5-hour limit window.

The real quota is server-side and nothing local states it. What IS local is
every assistant message Claude Code has written, with a timestamp and a usage
block, in ~/.claude/projects/**/*.jsonl. Summing those over a rolling window
gives a usable proxy: not the quota, but a consistent measure of how hard the
window has been worked, plus when the oldest tokens in it age out.

A RAW TOKEN COUNT IS NOT A SIGNAL. It is a numerator. "4.8 million" sounds
alarming and means nothing without the denominator — on a large plan that can
be 30% of the window. Decisions about backing off MUST be made on percentage,
which means the quota has to be calibrated at least once. See --calibrate.

SCOPE — read this before trusting the number:

  Counted:     every Claude Code session on THIS machine, all projects, not
               just the session asking. Parallel sessions in other repos land
               in the same window and this picks them up.
  NOT counted: Claude Code on any other machine, claude.ai web usage, direct
               API calls, and anything else that never writes a transcript here.

So the figure is a floor, never a ceiling. Real usage is this or higher, and
the gap is however much the account is being used elsewhere. Size your margin
accordingly: if the user works on several machines, treat a "comfortable"
reading with suspicion.

Treat the number as a trend, not a truth. It is here so a long autonomous run
can back off *before* an agent dies mid-task, rather than discovering the limit
by being killed by it. The reactive backstop still matters: an agent that dies
on a limit error is telling you the truth this script only estimates.

Usage:
    usage-window.py                 # human summary
    usage-window.py --json          # machine readable
    usage-window.py --hours 5       # window size (default 5)
    usage-window.py --quota 16e6    # one-off quota override
    usage-window.py --calibrate 30  # "Claude Code says I'm at 30%" -> store implied quota

Calibration:
    Ask the user what percentage their /status (or the Claude Code UI) reports,
    then run --calibrate with it. The implied quota is written to
    ~/.claude/nightshift-quota.json and used by every later run. Re-calibrate
    if the plan changes. One data point beats zero; two beats one.
"""

from __future__ import annotations

import argparse
import json
import os
from datetime import datetime, timedelta, timezone
from pathlib import Path

PROJECTS = Path.home() / ".claude" / "projects"
QUOTA_FILE = Path.home() / ".claude" / "nightshift-quota.json"

# Percentage bands. These are about RISK TO IN-FLIGHT WORK, not thrift — the
# window refills continuously, and an unused window is wasted capacity, not
# saved money. The only thing being avoided is agents dying mid-task.
BAND_CLEAR = 60.0   # below this: dispatch freely
BAND_WRAP = 80.0    # below this: keep going, prefer shorter tasks
# at or above BAND_WRAP: finish what's running, commit, push, park


def parse_ts(value: str | None) -> datetime | None:
    if not value:
        return None
    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError:
        return None


def billable(usage: dict) -> int:
    """Tokens that plausibly count against a quota.

    Cache reads are excluded: they are the cheap path and counting them would
    make a long cached session look far heavier than it is. Cache *creation* is
    included, because writing the cache is real work.
    """
    return (
        int(usage.get("input_tokens") or 0)
        + int(usage.get("cache_creation_input_tokens") or 0)
        + int(usage.get("output_tokens") or 0)
    )


def load_quota() -> tuple[int | None, str | None]:
    """Calibrated window quota, if the user has ever supplied one.

    Precedence: --quota (handled by caller) > env > stored file > unknown.
    Returns (quota, source) so the output can say where the number came from —
    an uncalibrated guess must never masquerade as a measurement.
    """
    env = os.environ.get("NIGHTSHIFT_WINDOW_QUOTA")
    if env:
        try:
            return int(float(env)), "env NIGHTSHIFT_WINDOW_QUOTA"
        except ValueError:
            pass
    try:
        data = json.loads(QUOTA_FILE.read_text(encoding="utf-8"))
        quota = int(data["windowTokenQuota"])
        at = data.get("calibratedAt", "unknown date")
        return quota, f"calibrated {at[:10]}"
    except (OSError, KeyError, ValueError, TypeError):
        return None, None


def save_quota(quota: int, observed_tokens: int, observed_pct: float) -> None:
    QUOTA_FILE.parent.mkdir(parents=True, exist_ok=True)
    QUOTA_FILE.write_text(
        json.dumps(
            {
                "windowTokenQuota": quota,
                "calibratedAt": datetime.now(timezone.utc).isoformat(),
                "observedTokens": observed_tokens,
                "observedPercent": observed_pct,
                "note": (
                    "Implied quota = observedTokens / (observedPercent/100). "
                    "Derived from what Claude Code reported at one moment, so it "
                    "inherits this script's blind spots (other machines, web, API). "
                    "Re-run --calibrate after a plan change or if parking decisions "
                    "start feeling wrong."
                ),
            },
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )


def collect(hours: float) -> dict:
    now = datetime.now(timezone.utc)
    cutoff = now - timedelta(hours=hours)

    total = 0
    output = 0
    messages = 0
    oldest: datetime | None = None
    newest: datetime | None = None
    per_project: dict[str, int] = {}

    if not PROJECTS.is_dir():
        return {
            "error": f"no transcript directory at {PROJECTS}",
            "windowHours": hours,
        }

    for path in PROJECTS.glob("*/*.jsonl"):
        # Cheap skip: a file untouched since the cutoff cannot contribute.
        try:
            if datetime.fromtimestamp(path.stat().st_mtime, timezone.utc) < cutoff:
                continue
        except OSError:
            continue

        project = path.parent.name
        try:
            with path.open(encoding="utf-8", errors="replace") as handle:
                for line in handle:
                    if '"usage"' not in line:
                        continue
                    try:
                        entry = json.loads(line)
                    except json.JSONDecodeError:
                        continue
                    stamp = parse_ts(entry.get("timestamp"))
                    if stamp is None or stamp < cutoff:
                        continue
                    usage = (entry.get("message") or {}).get("usage")
                    if not isinstance(usage, dict):
                        continue

                    tokens = billable(usage)
                    if tokens <= 0:
                        continue

                    total += tokens
                    output += int(usage.get("output_tokens") or 0)
                    messages += 1
                    per_project[project] = per_project.get(project, 0) + tokens
                    oldest = stamp if oldest is None or stamp < oldest else oldest
                    newest = stamp if newest is None or stamp > newest else newest
        except OSError:
            continue

    # When the oldest tokens in the window age out. This is the soonest the
    # window measurably loosens; it is NOT a quota reset time, which only the
    # server knows.
    ages_out_in = None
    if oldest is not None:
        ages_out_in = max(0, int(((oldest + timedelta(hours=hours)) - now).total_seconds()))

    return {
        "windowHours": hours,
        "tokens": total,
        "outputTokens": output,
        "messages": messages,
        "projects": dict(sorted(per_project.items(), key=lambda kv: -kv[1])),
        "oldestInWindow": oldest.isoformat() if oldest else None,
        "newestInWindow": newest.isoformat() if newest else None,
        "oldestAgesOutInSeconds": ages_out_in,
        "now": now.isoformat(),
    }


def add_verdict(report: dict, quota: int | None, source: str | None) -> dict:
    """Attach percent-of-quota and an explicit act-on-this verdict.

    Without a calibrated quota there is NO verdict — deliberately. An earlier
    version of this tool emitted only a raw count, and the reading agent
    invented a magnitude threshold to compare it against, then parked a run at
    what turned out to be 30% of the window. Saying "unknown" is strictly better
    than handing back a number that invites a made-up denominator.
    """
    report["quota"] = quota
    report["quotaSource"] = source
    if not quota or quota <= 0:
        report["percentUsed"] = None
        report["verdict"] = "unknown"
        report["verdictReason"] = (
            "No calibrated quota. Percentage is unknown, so the raw token count says "
            "NOTHING about how close the window is. Do not park on magnitude. "
            "Ask the user what percent Claude Code reports, then run --calibrate."
        )
        return report

    pct = 100.0 * report["tokens"] / quota
    report["percentUsed"] = round(pct, 1)
    if pct < BAND_CLEAR:
        report["verdict"] = "clear"
        report["verdictReason"] = "Dispatch freely. The window refills; unused capacity is wasted."
    elif pct < BAND_WRAP:
        report["verdict"] = "wrap-up"
        report["verdictReason"] = (
            "Keep working, but prefer shorter tasks over long fan-outs so nothing "
            "large is in flight if the window tightens."
        )
    else:
        report["verdict"] = "park"
        report["verdictReason"] = (
            "Finish what is running, commit, push, write the next step down, and park."
        )
    return report


def human(report: dict) -> str:
    if "error" in report:
        return f"usage: unavailable ({report['error']})"

    lines = [
        f"rolling {report['windowHours']}h window — all local Claude Code sessions",
    ]

    quota = report.get("quota")
    pct = report.get("percentUsed")
    if quota and pct is not None:
        lines.append(
            f"  usage    {report['tokens']:,} / {quota:,} tokens  =  {pct:.1f}%"
            f"   [{report['quotaSource']}]"
        )
    else:
        lines.append(f"  tokens   {report['tokens']:,} (no calibrated quota — percentage UNKNOWN)")

    lines.append(
        f"  detail   {report['outputTokens']:,} output over {report['messages']:,} messages"
    )

    if report["oldestAgesOutInSeconds"] is not None:
        secs = report["oldestAgesOutInSeconds"]
        lines.append(f"  oldest entry ages out in {secs // 60}m{secs % 60:02d}s")

    top = list(report["projects"].items())[:3]
    if top:
        lines.append("  busiest: " + ", ".join(f"{k} {v:,}" for k, v in top))
    lines.append(f"  sessions: {len(report['projects'])} project(s) contributing")

    lines.append(f"  VERDICT: {report['verdict']} — {report['verdictReason']}")
    lines.append(
        "  FLOOR, not a ceiling: misses other machines, claude.ai and direct API use"
    )
    if not quota:
        lines.append(
            "  calibrate: ask the user their /status percentage, then "
            "usage-window.py --calibrate <pct>"
        )
    return "\n".join(lines)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--hours", type=float, default=5.0)
    parser.add_argument("--json", action="store_true")
    parser.add_argument(
        "--quota",
        type=float,
        default=None,
        help="Window token quota for this run only (accepts 16e6). Overrides stored calibration.",
    )
    parser.add_argument(
        "--calibrate",
        type=float,
        default=None,
        metavar="PCT",
        help="Percent-used that Claude Code currently reports. Stores the implied quota.",
    )
    args = parser.parse_args()

    report = collect(args.hours)

    if args.calibrate is not None:
        if "error" in report:
            print(human(report))
            raise SystemExit(1)
        if not 0 < args.calibrate <= 100:
            print("calibrate: percentage must be in (0, 100]")
            raise SystemExit(2)
        implied = int(report["tokens"] / (args.calibrate / 100.0))
        save_quota(implied, report["tokens"], args.calibrate)
        print(
            f"calibrated: {report['tokens']:,} tokens reported as {args.calibrate}% "
            f"=> window quota ~{implied:,}\n"
            f"written to {QUOTA_FILE}\n"
            "Note this inherits the script's blind spots (other machines, web, direct API),\n"
            "so the true quota is this or LARGER. Re-calibrate after a plan change."
        )
        return

    if args.quota is not None:
        quota, source = int(args.quota), "--quota flag"
    else:
        quota, source = load_quota()

    report = add_verdict(report, quota, source)

    if args.json:
        print(json.dumps(report, indent=2))
    else:
        print(human(report))


if __name__ == "__main__":
    main()