#!/usr/bin/env python3
"""Claude Code UserPromptSubmit 훅: 대화가 예산(기본 30만 토큰)을 넘으면 인계를 권하게 한다.

설치: 이 파일을 ~/.claude/hooks/context-budget-hook.py 에 두고 실행 권한(chmod +x)을 준 뒤
~/.claude/settings.json 의 "hooks" 에 아래를 넣는다.
  "UserPromptSubmit": [{"hooks": [{"type": "command",
      "command": "~/.claude/hooks/context-budget-hook.py", "timeout": 3}]}]
예산 바꾸기: command 앞에 CONTEXT_BUDGET_TOKENS=150000 처럼 붙인다.
대화 창이 20만 토큰인 환경은 그 근처에서 자동 요약이 먼저 돌므로 15만 정도로 낮춘다(/context로 확인).
실패하면 아무것도 출력하지 않는다(대화를 막지 않음).
"""
import json
import os
import sys
from pathlib import Path

TAIL = 2 * 1024 * 1024  # 대화 기록 끝 2MiB만 읽는다


def last_context(path: Path):
    """마지막 assistant 응답의 입력 토큰(캐시 포함) 합계."""
    with open(path, "rb") as handle:
        size = handle.seek(0, 2)
        handle.seek(max(0, size - TAIL))
        lines = handle.read().splitlines()
    if size > TAIL:
        lines = lines[1:]
    for line in reversed(lines):
        try:
            row = json.loads(line)
        except ValueError:
            continue
        usage = (row.get("message") or {}).get("usage") if row.get("type") == "assistant" else None
        if usage:
            return sum(int(usage.get(key) or 0) for key in
                       ("input_tokens", "cache_read_input_tokens", "cache_creation_input_tokens"))
    return None


def main():
    payload = json.load(sys.stdin)  # 붙여 넣은 긴 입력도 그대로 받는다
    session = str(payload.get("session_id", ""))
    if not session.replace("-", "").isalnum():
        return
    tokens = last_context(Path(payload["transcript_path"]))
    budget = int(os.environ.get("CONTEXT_BUDGET_TOKENS", "300000"))
    if tokens is None or tokens <= budget:
        return
    bucket = (tokens - budget - 1) // 100000  # 30만, 40만, 50만…을 넘을 때 한 번씩만
    state = Path.home() / ".cache" / "context-budget" / session
    state.parent.mkdir(parents=True, exist_ok=True)
    if state.exists() and state.read_text().strip().isdigit() and int(state.read_text()) >= bucket:
        return
    state.write_text(str(bucket))
    message = (
        f"컨텍스트가 약 {tokens / 10000:.0f}만 토큰입니다(예산 {budget / 10000:g}만). "
        "이번 답의 끝에 새 세션으로 넘어갈지 한 줄로 권하세요. 사용자가 원하면 HANDOFF.md에 "
        "'- 정한 것:', '- 맡은 일:'(남은 일, 돌던 작업의 식별자·상태·로그 위치), "
        "'- 상시 작업:'(이 창을 닫으면 멈추는 프로세스만, 없으면 정확히 '없음'), "
        "'- 참고:'(설명·다른 사람 몫·승인 대기, 실행 금지)를 쓰고 그 파일만 커밋하세요. "
        "도구를 쓰는 도중에는 끼워 넣지 마세요."
    )
    print(json.dumps({"hookSpecificOutput": {"hookEventName": "UserPromptSubmit",
                                             "additionalContext": message}}, ensure_ascii=False))


if __name__ == "__main__":
    try:
        main()
    except Exception:
        pass