#!/usr/bin/env bash
# token-tracker.sh -- PostToolUse hook (matcher: Agent)
# Records token usage per Agent tool call for anomaly detection baselines.
# Never blocks (always exits 0).
#
# Part of @connexum/ai-governance

set -uo pipefail

INPUT=$(cat)

# AIAG-018: enforce 1MB input size limit
if [ "${#INPUT}" -gt 1048576 ]; then
  echo "BLOCKED: Hook input exceeds 1MB. Possible injection attempt." >&2
  exit 0  # PostToolUse: skip tracking but don't block the tool.
fi

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

# Only track Agent calls (which consume API tokens)
if [ "$TOOL_NAME" != "Agent" ]; then
  exit 0
fi

# F-NEW-CRIT-38b: env var fallback when payload .agent_name absent.
AGENT_NAME=$(echo "$INPUT" | jq -r '.agent_name // empty' 2>/dev/null)
if [ -z "$AGENT_NAME" ]; then
  AGENT_NAME="${CXNI_AGENT_NAME:-unknown}"
fi
TIMESTAMP=$(date -u +"%Y-%m-%dT%H:%M:%SZ")

# Extract usage from tool result if available
TOOL_RESULT=$(echo "$INPUT" | jq -r '.tool_result // ""' 2>/dev/null)
TOTAL_TOKENS=$(echo "$TOOL_RESULT" | grep -oP 'total_tokens:\s*\K\d+' | head -1)
DURATION_MS=$(echo "$TOOL_RESULT" | grep -oP 'duration_ms:\s*\K\d+' | head -1)

TOTAL_TOKENS="${TOTAL_TOKENS:-0}"
DURATION_MS="${DURATION_MS:-0}"

# Token tracking directory
TRACK_DIR=""
for candidate in \
  ".claude/governance/token-tracking" \
  ".claude/cost-tracking" \
  "GOVERNANCE/token-tracking"; do
  if [ -d "$candidate" ] || mkdir -p "$candidate" 2>/dev/null; then
    TRACK_DIR="$candidate"
    break
  fi
done

if [ -z "$TRACK_DIR" ]; then
  exit 0
fi

TRACK_FILE="$TRACK_DIR/tokens-$(date -u +%Y-%m-%d).jsonl"

# Write token usage entry
# F-NEW-CRIT-38b: jq -c so each .jsonl entry is one line.
ENTRY=$(jq -nc \
  --arg ts "$TIMESTAMP" \
  --arg agent "$AGENT_NAME" \
  --argjson tokens "$TOTAL_TOKENS" \
  --argjson duration "$DURATION_MS" \
  '{timestamp: $ts, agent: $agent, tokens: $tokens, duration_ms: $duration}' 2>/dev/null)

if [ -n "$ENTRY" ]; then
  echo "$ENTRY" >> "$TRACK_FILE" 2>/dev/null
fi

# Update running session total for budget gate
BUDGET_FILE=""
for candidate in \
  ".claude/governance/token-budget.json" \
  ".claude/cost-tracking/budget.json"; do
  if [ -f "$candidate" ]; then
    BUDGET_FILE="$candidate"
    break
  fi
done

if [ -n "$BUDGET_FILE" ] && [ "$TOTAL_TOKENS" -gt 0 ] 2>/dev/null; then
  CURRENT=$(jq -r '.current_session_tokens // 0' "$BUDGET_FILE" 2>/dev/null)
  NEW_TOTAL=$((CURRENT + TOTAL_TOKENS))
  jq ".current_session_tokens = $NEW_TOTAL" "$BUDGET_FILE" > "${BUDGET_FILE}.tmp" 2>/dev/null && \
    mv "${BUDGET_FILE}.tmp" "$BUDGET_FILE" 2>/dev/null
fi

exit 0
