#!/usr/bin/env bash
#
# pre-compact-snapshot.sh — Snapshot impact metrics before context compaction.
#
# Runs as a PreCompact hook. Reads the conversation transcript, extracts
# actual token counts when available (falls back to heuristic estimates),
# and appends a timestamped entry to the impact log.
#
# Input: JSON on stdin with fields: trigger, session_id, transcript_path, cwd
# Output: nothing on stdout (hook succeeds silently). Logs to impact-log.jsonl.

set -euo pipefail

HOOK_INPUT=$(cat)
PROJECT_DIR="${CLAUDE_PROJECT_DIR:-$(echo "$HOOK_INPUT" | jq -r '.cwd')}"
TRANSCRIPT_PATH=$(echo "$HOOK_INPUT" | jq -r '.transcript_path')
SESSION_ID=$(echo "$HOOK_INPUT" | jq -r '.session_id')
TRIGGER=$(echo "$HOOK_INPUT" | jq -r '.trigger')
TIMESTAMP=$(date -u +"%Y-%m-%dT%H:%M:%SZ")

LOG_DIR="$PROJECT_DIR/.claude/impact"
LOG_FILE="$LOG_DIR/impact-log.jsonl"
mkdir -p "$LOG_DIR"

# --- Extract or estimate metrics from transcript ---

if [ -f "$TRANSCRIPT_PATH" ]; then
  TRANSCRIPT_BYTES=$(wc -c < "$TRANSCRIPT_PATH")
  TRANSCRIPT_LINES=$(wc -l < "$TRANSCRIPT_PATH")

  # Count tool uses
  TOOL_USES=$(grep -c '"tool_use"' "$TRANSCRIPT_PATH" 2>/dev/null || echo 0)

  # Try to extract actual token counts from usage fields in the transcript.
  # The transcript contains .message.usage with input_tokens,
  # cache_creation_input_tokens, cache_read_input_tokens, output_tokens.
  USAGE_DATA=$(python3 -c "
import json, sys
input_tokens = 0
cache_creation = 0
cache_read = 0
output_tokens = 0
turns = 0
with open(sys.argv[1]) as f:
    for line in f:
        try:
            d = json.loads(line.strip())
            u = d.get('message', {}).get('usage')
            if u and 'input_tokens' in u:
                turns += 1
                input_tokens += u.get('input_tokens', 0)
                cache_creation += u.get('cache_creation_input_tokens', 0)
                cache_read += u.get('cache_read_input_tokens', 0)
                output_tokens += u.get('output_tokens', 0)
        except Exception:
            pass
# Print as tab-separated for easy shell parsing
print(f'{turns}\t{input_tokens}\t{cache_creation}\t{cache_read}\t{output_tokens}')
" "$TRANSCRIPT_PATH" 2>/dev/null || echo "")

  if [ -n "$USAGE_DATA" ] && [ "$(echo "$USAGE_DATA" | cut -f1)" -gt 0 ] 2>/dev/null; then
    # Actual token counts available
    TOKEN_SOURCE="actual"
    ASSISTANT_TURNS=$(echo "$USAGE_DATA" | cut -f1)
    INPUT_TOKENS=$(echo "$USAGE_DATA" | cut -f2)
    CACHE_CREATION=$(echo "$USAGE_DATA" | cut -f3)
    CACHE_READ=$(echo "$USAGE_DATA" | cut -f4)
    OUTPUT_TOKENS=$(echo "$USAGE_DATA" | cut -f5)

    # Cumulative input = all tokens that went through the model.
    # Cache reads are cheaper (~10-20% of full compute), so we weight them.
    # Full-cost tokens: input_tokens + cache_creation_input_tokens
    # Reduced-cost tokens: cache_read_input_tokens (weight at 0.1x for energy)
    FULL_COST_INPUT=$(( INPUT_TOKENS + CACHE_CREATION ))
    CACHE_READ_EFFECTIVE=$(( CACHE_READ / 10 ))
    CUMULATIVE_INPUT=$(( FULL_COST_INPUT + CACHE_READ_EFFECTIVE ))
    # Also track raw total for the log
    CUMULATIVE_INPUT_RAW=$(( INPUT_TOKENS + CACHE_CREATION + CACHE_READ ))
  else
    # Fallback: heuristic estimation
    TOKEN_SOURCE="heuristic"
    ESTIMATED_TOKENS=$((TRANSCRIPT_BYTES / 4))
    ASSISTANT_TURNS=$(grep -c '"role":\s*"assistant"' "$TRANSCRIPT_PATH" 2>/dev/null || echo 0)

    if [ "$ASSISTANT_TURNS" -gt 0 ]; then
      AVG_CONTEXT=$((ESTIMATED_TOKENS / 2))
      CUMULATIVE_INPUT=$((AVG_CONTEXT * ASSISTANT_TURNS))
    else
      CUMULATIVE_INPUT=$ESTIMATED_TOKENS
    fi
    CUMULATIVE_INPUT_RAW=$CUMULATIVE_INPUT
    OUTPUT_TOKENS=$((ESTIMATED_TOKENS / 20))
    CACHE_CREATION=0
    CACHE_READ=0
    INPUT_TOKENS=0
  fi

  # --- Cost estimates ---
  # Energy: 0.1 Wh per 1K input tokens, 0.5 Wh per 1K output tokens, PUE 1.2
  # Calibrated against Google (Patterson et al., Aug 2025) and Jegham et al. (May 2025)
  # Using integer arithmetic in centiwatt-hours to avoid bc dependency
  INPUT_CWH=$(( CUMULATIVE_INPUT * 100 / 10000 ))   # 0.1 Wh/1K = 100 cWh/10K
  OUTPUT_CWH=$(( OUTPUT_TOKENS * 500 / 10000 ))      # 0.5 Wh/1K = 500 cWh/10K
  ENERGY_CWH=$(( (INPUT_CWH + OUTPUT_CWH) * 12 / 10 ))  # PUE 1.2
  ENERGY_WH=$(( ENERGY_CWH / 100 ))

  # CO2: 325g/kWh -> 0.325g/Wh -> 325 mg/Wh
  CO2_MG=$(( ENERGY_WH * 325 ))
  CO2_G=$(( CO2_MG / 1000 ))

  # Financial: $15/M input, $75/M output (in cents)
  # Use effective cumulative input (cache-weighted) for cost too
  COST_INPUT_CENTS=$(( CUMULATIVE_INPUT * 15 / 10000 ))  # $15/M = 1.5c/100K
  COST_OUTPUT_CENTS=$(( OUTPUT_TOKENS * 75 / 10000 ))
  COST_CENTS=$(( COST_INPUT_CENTS + COST_OUTPUT_CENTS ))
else
  TRANSCRIPT_BYTES=0
  TRANSCRIPT_LINES=0
  ASSISTANT_TURNS=0
  TOOL_USES=0
  CUMULATIVE_INPUT=0
  CUMULATIVE_INPUT_RAW=0
  OUTPUT_TOKENS=0
  CACHE_CREATION=0
  CACHE_READ=0
  ENERGY_WH=0
  CO2_G=0
  COST_CENTS=0
  TOKEN_SOURCE="none"
fi

# --- Write log entry ---

cat >> "$LOG_FILE" <<EOF
{"timestamp":"$TIMESTAMP","session_id":"$SESSION_ID","trigger":"$TRIGGER","token_source":"$TOKEN_SOURCE","transcript_bytes":$TRANSCRIPT_BYTES,"transcript_lines":$TRANSCRIPT_LINES,"assistant_turns":$ASSISTANT_TURNS,"tool_uses":$TOOL_USES,"cumulative_input_tokens":$CUMULATIVE_INPUT,"cumulative_input_raw":$CUMULATIVE_INPUT_RAW,"cache_creation_tokens":$CACHE_CREATION,"cache_read_tokens":$CACHE_READ,"output_tokens":$OUTPUT_TOKENS,"energy_wh":$ENERGY_WH,"co2_g":$CO2_G,"cost_cents":$COST_CENTS}
EOF

exit 0