#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
CONFIG_PATH="${MODEL_GUARD_CONFIG:-$ROOT_DIR/config/model-budget-guard.config.json}"

if ! command -v jq >/dev/null 2>&1; then
  echo "[model-guard] jq is required" >&2
  exit 1
fi
if ! command -v openclaw >/dev/null 2>&1; then
  echo "[model-guard] openclaw CLI is required" >&2
  exit 1
fi

expand_tilde() {
  local p="$1"
  if [[ "$p" == "~" ]]; then
    echo "$HOME"
  elif [[ "$p" == "~/"* ]]; then
    echo "$HOME/${p#~/}"
  else
    echo "$p"
  fi
}

safe_now_ms() {
  python3 - <<'PY'
import time
print(int(time.time()*1000))
PY
}

send_notice() {
  local text="$1"
  local last_channel="$2"
  local last_to="$3"

  if [[ -z "$last_channel" || -z "$last_to" ]]; then
    echo "[model-guard] notice skipped (missing channel/target): $text" >&2
    return 0
  fi

  local target="$last_to"
  if [[ "$target" == *:* ]]; then
    target="${target#*:}"
  fi

  openclaw message send \
    --channel "$last_channel" \
    --target "$target" \
    --message "$text" \
    >/dev/null 2>&1 || true
}

if [[ ! -f "$CONFIG_PATH" ]]; then
  echo "[model-guard] missing config: $CONFIG_PATH" >&2
  exit 1
fi

enabled="$(jq -r '.enabled // true' "$CONFIG_PATH")"
if [[ "$enabled" != "true" ]]; then
  exit 0
fi

session_key="$(jq -r '.sessionKey // "agent:main:main"' "$CONFIG_PATH")"
low_model="$(jq -r '.lowModel // empty' "$CONFIG_PATH")"
warn_after_min="$(jq -r '.warnAfterMinutes // 2' "$CONFIG_PATH")"
revert_after_min="$(jq -r '.revertAfterMinutes // 45' "$CONFIG_PATH")"
min_warn_interval_min="$(jq -r '.minWarnIntervalMinutes // 20' "$CONFIG_PATH")"
state_file_raw="$(jq -r '.stateFile // "~/.openclaw/model-budget-guard-state.json"' "$CONFIG_PATH")"
state_file="$(expand_tilde "$state_file_raw")"

if [[ -z "$low_model" ]]; then
  echo "[model-guard] lowModel must be set" >&2
  exit 1
fi

state_dir="${OPENCLAW_STATE_DIR:-$HOME/.openclaw}"
sessions_file="$state_dir/agents/main/sessions/sessions.json"

if [[ ! -r "$sessions_file" ]]; then
  echo "[model-guard] sessions file not found: $sessions_file" >&2
  exit 0
fi

session_json="$(jq -c --arg key "$session_key" '.[$key] // {}' "$sessions_file" 2>/dev/null || true)"
if [[ -z "$session_json" || "$session_json" == "{}" ]]; then
  exit 0
fi

session_id="$(jq -r '.sessionId // empty' <<<"$session_json")"
updated_at="$(jq -r '.updatedAt // 0' <<<"$session_json")"
model_raw="$(jq -r '.model // empty' <<<"$session_json")"
provider_raw="$(jq -r '.provider // empty' <<<"$session_json")"
last_channel="$(jq -r '.lastChannel // empty' <<<"$session_json")"
last_to="$(jq -r '.lastTo // empty' <<<"$session_json")"

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

model_full="$model_raw"
if [[ "$model_raw" != */* && -n "$provider_raw" && "$provider_raw" != "null" ]]; then
  model_full="$provider_raw/$model_raw"
fi

high_match="$(jq -r --arg m "$model_full" '.highModels // [] | index($m) != null' "$CONFIG_PATH")"
if [[ "$high_match" != "true" ]]; then
  exit 0
fi

if [[ ! "$updated_at" =~ ^[0-9]+$ ]]; then
  exit 0
fi

now_ms="$(safe_now_ms)"
age_ms=$((now_ms - updated_at))
if (( age_ms < 0 )); then age_ms=0; fi
warn_after_ms=$((warn_after_min * 60 * 1000))
revert_after_ms=$((revert_after_min * 60 * 1000))
min_warn_interval_ms=$((min_warn_interval_min * 60 * 1000))

mkdir -p "$(dirname "$state_file")"
if [[ ! -f "$state_file" ]]; then
  cat > "$state_file" <<'JSON'
{}
JSON
fi

state_json="$(cat "$state_file")"
last_warned_model="$(jq -r --arg key "$session_key" '.[$key].lastWarnedModel // empty' <<<"$state_json")"
last_warned_session="$(jq -r --arg key "$session_key" '.[$key].lastWarnedSessionId // empty' <<<"$state_json")"
last_warn_at="$(jq -r --arg key "$session_key" '.[$key].lastWarnAt // 0' <<<"$state_json")"
last_revert_at="$(jq -r --arg key "$session_key" '.[$key].lastRevertAt // 0' <<<"$state_json")"

should_warn="false"
if (( age_ms >= warn_after_ms )); then
  if [[ "$last_warned_model" != "$model_full" || "$last_warned_session" != "$session_id" ]]; then
    should_warn="true"
  elif [[ "$last_warn_at" =~ ^[0-9]+$ ]] && (( now_ms - last_warn_at >= min_warn_interval_ms )); then
    should_warn="true"
  fi
fi

if [[ "$should_warn" == "true" ]]; then
  send_notice "Heads up: high-cost model active ($model_full). Switch back after heavy work: /model $low_model" "$last_channel" "$last_to"
  state_json="$(jq \
    --arg key "$session_key" \
    --arg model "$model_full" \
    --arg sid "$session_id" \
    --argjson now "$now_ms" \
    '.[$key].lastWarnedModel=$model | .[$key].lastWarnedSessionId=$sid | .[$key].lastWarnAt=$now' \
    <<<"$state_json")"
fi

if (( age_ms >= revert_after_ms )); then
  can_revert="true"
  if [[ "$last_revert_at" =~ ^[0-9]+$ ]] && (( now_ms - last_revert_at < min_warn_interval_ms )); then
    can_revert="false"
  fi

  if [[ "$can_revert" == "true" && -n "$session_id" ]]; then
    python3 - "$session_id" "$low_model" <<'PY'
import subprocess
import sys

session_id = sys.argv[1]
low_model = sys.argv[2]
cmd = [
    "openclaw",
    "agent",
    "--session-id", session_id,
    "--channel", "last",
    "--message", f"/model {low_model}",
]
try:
    subprocess.run(
        cmd,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        check=False,
        timeout=20,
    )
except subprocess.TimeoutExpired:
    pass
PY

    send_notice "Auto-switched back to $low_model to avoid high-model quota burn. Use high model only for deep tasks." "$last_channel" "$last_to"

    state_json="$(jq \
      --arg key "$session_key" \
      --arg model "$low_model" \
      --arg sid "$session_id" \
      --argjson now "$now_ms" \
      '.[$key].lastRevertAt=$now | .[$key].lastWarnedModel=$model | .[$key].lastWarnedSessionId=$sid | .[$key].lastWarnAt=$now' \
      <<<"$state_json")"
  fi
fi

printf '%s\n' "$state_json" > "$state_file"