#!/usr/bin/env bash
# cost-budget-gate.sh -- PreToolUse hook
# Blocks LLM calls when agent has exceeded its per-session token budget.
# Exit code 2 = block. Exit code 0 = allow.
#
# Part of @connexum/ai-governance

set -euo pipefail

INPUT=""; IFS= read -r -d '' -t 5 INPUT || true

# AIAG-018: enforce 1MB input size limit
if [ "${#INPUT}" -gt 1048576 ]; then
  echo "BLOCKED: Hook input exceeds 1MB. Possible injection attempt." >&2
  exit 2
fi

TOOL_NAME=$(echo "$INPUT" | jq -r '.tool_name // empty' 2>/dev/null)

# Only gate Agent tool (which spawns LLM calls)
case "$TOOL_NAME" in
  Agent) ;;
  *) exit 0 ;;
esac

# Check if budget tracking is enabled
BUDGET_FILE=""
for candidate in \
  ".claude/governance/token-budget.json" \
  ".claude/cost-tracking/budget.json" \
  "GOVERNANCE/token-budget.json"; do
  if [ -f "$candidate" ]; then
    BUDGET_FILE="$candidate"
    break
  fi
done

# F-NEW-MED-18: emit a one-time runtime advisory so the customer SEES
# that cost-budget-gate is dormant. Without this, an install ships
# cost-budget-gate.sh wired into Claude Code, the hook returns 0 on
# every PreToolUse, and the customer believes per-session token / cost
# limits are enforced when they are not. Same operator-knob pattern as
# F-NEW-MED-17 (file-boundary-guard advisory). Backward compat: still
# exit 0.
if [ -z "$BUDGET_FILE" ]; then
  if [ -z "${CXNI_BUDGET_NOOP_WARNED:-}" ]; then
    echo "[CXNI_BUDGET_NOOP] cost-budget-gate is installed but DORMANT: no token-budget.json found. Create one at .claude/governance/token-budget.json with {\"max_session_tokens\": <int>, \"max_session_cost_usd\": <number>} to enable enforcement. Set CXNI_BUDGET_NOOP_WARNED=true in your shell to silence this advisory after acknowledgement." >&2
  fi
  exit 0
fi

# Read budget limits
MAX_TOKENS=$(jq -r '.max_session_tokens // 0' "$BUDGET_FILE" 2>/dev/null)
CURRENT_TOKENS=$(jq -r '.current_session_tokens // 0' "$BUDGET_FILE" 2>/dev/null)
MAX_COST_USD=$(jq -r '.max_session_cost_usd // 0' "$BUDGET_FILE" 2>/dev/null)
CURRENT_COST_USD=$(jq -r '.current_session_cost_usd // 0' "$BUDGET_FILE" 2>/dev/null)

# Check token budget
if [ "$MAX_TOKENS" -gt 0 ] 2>/dev/null && [ "$CURRENT_TOKENS" -ge "$MAX_TOKENS" ] 2>/dev/null; then
  echo "BLOCKED: Token budget exceeded. Used $CURRENT_TOKENS / $MAX_TOKENS tokens this session." >&2
  echo "Increase budget in $BUDGET_FILE or start a new session." >&2
  exit 2
fi

# Check cost budget (compare as strings since bash doesn't do float math well)
if [ "$MAX_COST_USD" != "0" ] 2>/dev/null; then
  OVER_BUDGET=$(echo "$CURRENT_COST_USD >= $MAX_COST_USD" | bc -l 2>/dev/null || echo "0")
  if [ "$OVER_BUDGET" = "1" ]; then
    echo "BLOCKED: Cost budget exceeded. Spent \$$CURRENT_COST_USD / \$$MAX_COST_USD this session." >&2
    exit 2
  fi
fi

exit 0
