#!/usr/bin/env python3
"""Context-usage monitor hook (UserPromptSubmit + PostToolUse).

Computes context usage from the last main-chain assistant message's API
usage in the session transcript and, when usage crosses a warning band,
emits a rollover warning to the user (systemMessage) and to Claude
(additionalContext) so it can proactively offer /handoff.

Warns once per band per session (state file in the OS temp dir).
Never blocks the session: exits 0 on any error.
"""

from __future__ import annotations

import json
import os
import sys
import tempfile

# Transcripts grow to many MB; the newest usage entry is near the end.
TAIL_BYTES = 512 * 1024


def resolve_window(session: str) -> int:
    # Prefer the real window size persisted by statusline.py (official API data);
    # fall back to the env default. Current Claude models have a 1M window.
    tmp = tempfile.gettempdir()
    for name in (f"claude_ctxwin_{session}", "claude_ctxwin_global"):
        try:
            with open(os.path.join(tmp, name)) as f:
                size = int(f.read().strip())
            if size > 0:
                return size
        except (OSError, ValueError):
            continue
    return int(os.environ.get("CLAUDE_CONTEXT_WINDOW", "1000000"))


def last_context_tokens(path: str) -> int | None:
    size = os.path.getsize(path)
    with open(path, "rb") as f:
        if size > TAIL_BYTES:
            f.seek(size - TAIL_BYTES)
            f.readline()  # drop the partial first line
        lines = f.read().decode("utf-8", "replace").splitlines()
    for line in reversed(lines):
        try:
            entry = json.loads(line)
        except json.JSONDecodeError:
            continue
        # Sidechain entries are subagent traffic; they don't occupy the main window.
        if entry.get("type") != "assistant" or entry.get("isSidechain"):
            continue
        usage = (entry.get("message") or {}).get("usage") or {}
        if "input_tokens" not in usage:
            continue
        return (
            usage.get("input_tokens", 0)
            + usage.get("cache_read_input_tokens", 0)
            + usage.get("cache_creation_input_tokens", 0)
            + usage.get("output_tokens", 0)
        )
    return None


def main() -> None:
    bands = sorted(
        int(b)
        for b in os.environ.get("CLAUDE_ROLLOVER_WARN_BANDS", "75,85,92").split(",")
    )
    data = json.load(sys.stdin)
    path = data.get("transcript_path")
    if not path or not os.path.isfile(path):
        return
    tokens = last_context_tokens(path)
    if tokens is None:
        return
    session = data.get("session_id", "unknown")
    window = resolve_window(session)
    pct = tokens * 100 // window
    band = max((b for b in bands if pct >= b), default=None)
    if band is None:
        return

    state = os.path.join(tempfile.gettempdir(), f"claude_rollover_{session}")
    try:
        with open(state) as f:
            warned = int(f.read().strip())
    except (OSError, ValueError):
        warned = 0
    if band <= warned:
        return
    with open(state, "w") as f:
        f.write(str(band))

    msg = (
        f"Context at ~{pct}% ({tokens:,}/{window:,} tokens). "
        "Rollover recommended: run /handoff, then /clear."
    )
    note = (
        f"[context-rollover] Context usage is at ~{pct}% of the window. "
        "Per this project's rollover protocol: finish the current atomic step "
        "(do not leave the working tree broken), then tell the user it is time "
        "to roll over and offer to run /handoff to write HANDOFF.md. "
        "Do NOT suggest /compact."
    )
    print(
        json.dumps(
            {
                "systemMessage": f"WARNING: {msg}",
                "hookSpecificOutput": {
                    "hookEventName": data.get("hook_event_name", "UserPromptSubmit"),
                    "additionalContext": note,
                },
            }
        )
    )


if __name__ == "__main__":
    try:
        main()
    except Exception:
        pass  # a monitoring hook must never break the session
    sys.exit(0)