#!/usr/bin/env python3
"""
Context Budget Awareness — UserPromptSubmit hook + reporter.

Hook mode (default, no args):
  Reads $TRANSCRIPT and warns at 50% / 75% / 90% usage. Helps avoid surprise
  context compression. Exit 0 always.

Report mode (--report):
  Prints a compact status line plus the most recent sub-skill load history
  from .claude/sub-skill-log.jsonl. Used by the router before deciding
  whether to load a sub-skill inline or delegate to a subagent.

Thresholds (chars of transcript):
  - 80_000 (~50%): INFO — recommend delegating heavy tasks
  - 140_000 (~75%): WARN — strongly recommend agents, avoid large reads
  - 180_000 (~90%): CRITICAL — commit work NOW; compression imminent
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

THRESHOLD_INFO = 80_000
THRESHOLD_WARNING = 140_000
THRESHOLD_CRITICAL = 180_000

_DEFAULT_ROOT = Path(__file__).resolve().parents[2] if len(Path(__file__).resolve().parents) >= 3 else Path(".").resolve()
SUB_SKILL_LOG = Path(os.environ.get("PROJECT_ROOT", _DEFAULT_ROOT)) / ".claude" / "sub-skill-log.jsonl"


def _size_to_state(size: int) -> tuple[str, int]:
    if size >= THRESHOLD_CRITICAL:
        return "CRITICAL", min(99, int(size / 200_000 * 100))
    if size >= THRESHOLD_WARNING:
        return "WARNING", int(size / 200_000 * 100)
    if size >= THRESHOLD_INFO:
        return "INFO", int(size / 200_000 * 100)
    return "OK", int(size / 200_000 * 100)


def hook_mode() -> int:
    transcript = os.environ.get("TRANSCRIPT", "")
    if not transcript:
        return 0
    state, pct = _size_to_state(len(transcript))
    if state == "CRITICAL":
        print(
            f"CONTEXT BUDGET CRITICAL (~{pct}%): Compression is imminent. "
            f"COMMIT any uncommitted work NOW. Delegate remaining tasks to agents. "
            f"Avoid reading large files.",
            file=sys.stderr,
        )
    elif state == "WARNING":
        print(
            f"CONTEXT BUDGET WARNING (~{pct}%): Context window is filling up. "
            f"Use agents (Task tool) for remaining CODE+TEST work. "
            f"Keep main session for pipeline operations only.",
            file=sys.stderr,
        )
    elif state == "INFO":
        print(
            f"CONTEXT BUDGET INFO (~{pct}%): Consider delegating heavy tasks to agents "
            f"to preserve context for pipeline operations.",
            file=sys.stderr,
        )
    return 0


def report_mode() -> int:
    transcript = os.environ.get("TRANSCRIPT", "")
    size = len(transcript)
    state, pct = _size_to_state(size)
    print(f"context: ~{pct}% ({size} chars transcript) — {state}")

    if SUB_SKILL_LOG.exists():
        lines = [l for l in SUB_SKILL_LOG.read_text(encoding="utf-8").splitlines() if l.strip()]
        recent = lines[-5:]
        if recent:
            print(f"recent sub-skill loads ({len(lines)} total):")
            for line in recent:
                try:
                    e = json.loads(line)
                    print(f"  {e['ts']}  {e['sub_skill']:<10}  {e.get('sr') or '-':<18}  {e.get('reason') or ''}")
                except json.JSONDecodeError:
                    continue
    return 0


def main() -> int:
    p = argparse.ArgumentParser(description="Context budget hook and reporter.")
    p.add_argument("--report", action="store_true",
                   help="print status + recent sub-skill load history")
    args = p.parse_args()
    if args.report:
        return report_mode()
    return hook_mode()


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