#!/usr/bin/env bash
# Bedrock quota monitor — SessionStart hook (synchronous).
# 1. Refreshes quota cache from CloudWatch (if stale)
# 2. Auto-configures statusline on first run (one-time, with backup)
# 3. Returns systemMessage warning if any quota window is above threshold
#
# Runs synchronously so the systemMessage is available before Claude responds.
# CloudWatch query takes ~1-2s on cache miss, instant on cache hit.

set -euo pipefail

# Only run when on Bedrock
if [ "${CLAUDE_CODE_USE_BEDROCK:-}" != "1" ]; then
  exit 0
fi

PLUGIN_ROOT="${CLAUDE_PLUGIN_ROOT:?CLAUDE_PLUGIN_ROOT not set}"
SETTINGS_FILE="$HOME/.claude/settings.json"
CACHE_FILE="/tmp/bedrock-quota-cache-$(whoami).json"
CACHE_MAX_AGE=${BEDROCK_QUOTA_CACHE_TTL:-300}
USERNAME_CACHE="/tmp/bedrock-iam-username-$(whoami)"
STATUSLINE_SCRIPT="$PLUGIN_ROOT/scripts/statusline.sh"

REGION="${AWS_REGION:-eu-central-1}"
PROFILE="${AWS_PROFILE:-bedrock}"

# ── 1. Auto-configure statusline (one-time) ──────────────────

if [ -f "$SETTINGS_FILE" ] && command -v jq &>/dev/null; then
  current_cmd=$(jq -r '.statusLine.command // ""' "$SETTINGS_FILE" 2>/dev/null)

  if [[ "$current_cmd" != *"gemma-bedrock-monitor"* ]]; then
    cp "$SETTINGS_FILE" "$SETTINGS_FILE.pre-bedrock-monitor.bak" 2>/dev/null || true
    tmp=$(mktemp)
    jq --arg cmd "bash \"$STATUSLINE_SCRIPT\"" \
      '.statusLine = {"type": "command", "command": $cmd}' \
      "$SETTINGS_FILE" > "$tmp" && mv "$tmp" "$SETTINGS_FILE"
  fi
fi

# ── 2. Refresh quota cache (if stale) ────────────────────────

needs_refresh=true
if [ -f "$CACHE_FILE" ]; then
  cache_age=$(( $(date +%s) - $(stat -c %Y "$CACHE_FILE" 2>/dev/null || stat -f %m "$CACHE_FILE" 2>/dev/null || echo 0) ))
  if [ "$cache_age" -lt "$CACHE_MAX_AGE" ]; then
    needs_refresh=false
  fi
fi

if [ "$needs_refresh" = true ]; then
  # Resolve IAM username
  if [ -f "$USERNAME_CACHE" ]; then
    IAM_USER=$(cat "$USERNAME_CACHE")
  else
    IAM_USER=$(aws sts get-caller-identity --profile "$PROFILE" --region "$REGION" \
      --query 'Arn' --output text 2>/dev/null | sed 's|.*/||') || exit 0
    [ -z "$IAM_USER" ] && exit 0
    echo "$IAM_USER" > "$USERNAME_CACHE"
  fi

  now=$(date -u +%Y-%m-%dT%H:%M:%SZ)
  start=$(date -u -d '30 minutes ago' +%Y-%m-%dT%H:%M:%SZ 2>/dev/null || \
          date -u -v-30M +%Y-%m-%dT%H:%M:%SZ)

  mk_query() {
    jq -n --arg id "$1" --arg metric "$2" --arg user "$IAM_USER" \
      '{Id: $id, MetricStat: {Metric: {Namespace: "Custom/Bedrock", MetricName: $metric, Dimensions: [{Name: "User", Value: $user}]}, Period: 900, Stat: "Maximum"}}'
  }

  queries=$(jq -n \
    --argjson q0 "$(mk_query cost3h   QuotaWindowCost3h)" \
    --argjson q1 "$(mk_query costday  QuotaWindowCostDay)" \
    --argjson q2 "$(mk_query costweek QuotaWindowCostWeek)" \
    --argjson q3 "$(mk_query lim3h    QuotaWindowCost3hLimit)" \
    --argjson q4 "$(mk_query limday   QuotaWindowCostDayLimit)" \
    --argjson q5 "$(mk_query limweek  QuotaWindowCostWeekLimit)" \
    '[$q0, $q1, $q2, $q3, $q4, $q5]')

  aws cloudwatch get-metric-data \
    --profile "$PROFILE" \
    --region "$REGION" \
    --start-time "$start" \
    --end-time "$now" \
    --metric-data-queries "$queries" \
    --output json 2>/dev/null | jq --arg iam_user "$IAM_USER" '{
      cost_3h:    (.MetricDataResults[] | select(.Id == "cost3h")   | .Values[0] // 0),
      cost_day:   (.MetricDataResults[] | select(.Id == "costday")  | .Values[0] // 0),
      cost_week:  (.MetricDataResults[] | select(.Id == "costweek") | .Values[0] // 0),
      limit_3h:   (.MetricDataResults[] | select(.Id == "lim3h")    | .Values[0] // null),
      limit_day:  (.MetricDataResults[] | select(.Id == "limday")   | .Values[0] // null),
      limit_week: (.MetricDataResults[] | select(.Id == "limweek")  | .Values[0] // null),
      user: $iam_user,
      updated_at: now | todate
    }' > "$CACHE_FILE" 2>/dev/null || true
fi

# ── 3. Check quota and return systemMessage ───────────────────

if [ ! -f "$CACHE_FILE" ]; then
  exit 0
fi

cost_3h=$(jq -r '.cost_3h // 0' "$CACHE_FILE" 2>/dev/null)   || exit 0
cost_day=$(jq -r '.cost_day // 0' "$CACHE_FILE" 2>/dev/null)  || exit 0
cost_week=$(jq -r '.cost_week // 0' "$CACHE_FILE" 2>/dev/null) || exit 0

limit_3h=$(jq -r '.limit_3h // empty' "$CACHE_FILE" 2>/dev/null)
limit_day=$(jq -r '.limit_day // empty' "$CACHE_FILE" 2>/dev/null)
limit_week=$(jq -r '.limit_week // empty' "$CACHE_FILE" 2>/dev/null)

limit_3h=${limit_3h:-${BEDROCK_QUOTA_3H:-6.00}}
limit_day=${limit_day:-${BEDROCK_QUOTA_DAY:-15.00}}
limit_week=${limit_week:-${BEDROCK_QUOTA_WEEK:-50.00}}

pct_3h=$(awk "BEGIN { l=$limit_3h; printf \"%.0f\", (l > 0 ? ($cost_3h / l) * 100 : 0) }")
pct_day=$(awk "BEGIN { l=$limit_day; printf \"%.0f\", (l > 0 ? ($cost_day / l) * 100 : 0) }")
pct_week=$(awk "BEGIN { l=$limit_week; printf \"%.0f\", (l > 0 ? ($cost_week / l) * 100 : 0) }")

max_pct=$pct_3h
[ "$pct_day" -gt "$max_pct" ] 2>/dev/null && max_pct=$pct_day
[ "$pct_week" -gt "$max_pct" ] 2>/dev/null && max_pct=$pct_week

# Only warn if any window is at 50%+
if [ "$max_pct" -lt 50 ]; then
  exit 0
fi

# Build warning
parts=""
[ "$pct_3h" -ge 20 ] 2>/dev/null && parts="${parts}3h: ${pct_3h}% (\$${cost_3h}/\$${limit_3h})  "
[ "$pct_day" -ge 20 ] 2>/dev/null && parts="${parts}day: ${pct_day}% (\$${cost_day}/\$${limit_day})  "
[ "$pct_week" -ge 20 ] 2>/dev/null && parts="${parts}week: ${pct_week}% (\$${cost_week}/\$${limit_week})"

if [ "$max_pct" -ge 100 ]; then
  level="QUOTA EXCEEDED — Bedrock access is denied until the quota window resets. Do NOT attempt API-heavy operations."
elif [ "$max_pct" -ge 80 ]; then
  level="APPROACHING LIMIT — be mindful of token usage, avoid unnecessary agent spawning or large context operations"
else
  level="MODERATE usage — monitor your spending"
fi

jq -n \
  --arg msg "IMPORTANT — Bedrock quota warning ($level). Current usage: $parts. You MUST inform the user about this quota status in your first response." \
  '{"systemMessage": $msg}'