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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
PROFILES_CONFIG="${MODEL_PROFILES_CONFIG:-$ROOT_DIR/config/model-profiles.config.json}"
BUDGET_CONFIG="${MODEL_GUARD_CONFIG:-$ROOT_DIR/config/model-budget-guard.config.json}"
AUTH_STORE="${OPENCLAW_AUTH_STORE:-$HOME/.openclaw/agents/main/agent/auth-profiles.json}"

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

usage() {
  cat <<'USAGE'
Usage:
  bash ./scripts/model_hygiene_workflow.sh <profile> [--no-live] [--no-guards]

Behavior:
  1) Applies profile via model_profile_switch.sh
  2) Reinstalls guard LaunchAgents (unless --no-guards)
  3) Prunes legacy models/providers/auth not referenced by config
  4) Prints models status

Examples:
  bash ./scripts/model_hygiene_workflow.sh paid
  bash ./scripts/model_hygiene_workflow.sh free --no-live
USAGE
}

in_array() {
  local needle="$1"
  shift
  local item
  for item in "$@"; do
    if [[ "$item" == "$needle" ]]; then
      return 0
    fi
  done
  return 1
}

profile="${1:-}"
if [[ -z "$profile" ]]; then
  usage
  exit 1
fi
if [[ "$profile" == "-h" || "$profile" == "--help" ]]; then
  usage
  exit 0
fi
shift || true

apply_live="true"
install_guards="true"
while (( "$#" )); do
  case "$1" in
    --no-live) apply_live="false" ;;
    --no-guards) install_guards="false" ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "[model-hygiene] unknown argument: $1" >&2
      usage
      exit 1
      ;;
  esac
  shift
done

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

profile_exists="$(jq -r --arg p "$profile" '.profiles[$p] != null' "$PROFILES_CONFIG")"
if [[ "$profile_exists" != "true" ]]; then
  echo "[model-hygiene] unknown profile '$profile'" >&2
  jq -r '.profiles | keys[]' "$PROFILES_CONFIG" | sed 's/^/  - /'
  exit 1
fi

echo "[model-hygiene] collecting managed models/providers from config"
keep_models=()
while IFS= read -r model; do
  [[ -z "$model" ]] && continue
  if ! in_array "$model" "${keep_models[@]-}"; then
    keep_models+=("$model")
  fi
done < <(jq -r '.profiles | to_entries[] | .value.primary, (.value.fallbacks[]? // empty)' "$PROFILES_CONFIG")

if [[ -f "$BUDGET_CONFIG" ]]; then
  while IFS= read -r model; do
    [[ -z "$model" ]] && continue
    if ! in_array "$model" "${keep_models[@]-}"; then
      keep_models+=("$model")
    fi
  done < <(jq -r '.lowModel // empty, (.highModels[]? // empty)' "$BUDGET_CONFIG")
fi

if [[ ${#keep_models[@]} -eq 0 ]]; then
  echo "[model-hygiene] no models discovered in config; aborting" >&2
  exit 1
fi

keep_providers=()
for model in "${keep_models[@]}"; do
  provider="${model%%/*}"
  [[ -z "$provider" || "$provider" == "$model" ]] && continue
  if ! in_array "$provider" "${keep_providers[@]-}"; then
    keep_providers+=("$provider")
  fi
done

echo "[model-hygiene] applying profile '$profile'"
switch_args=("$profile" "--no-status")
if [[ "$apply_live" != "true" ]]; then
  switch_args+=("--no-live")
fi
bash "$SCRIPT_DIR/model_profile_switch.sh" "${switch_args[@]}"

if [[ "$install_guards" == "true" ]]; then
  echo "[model-hygiene] reinstalling guard LaunchAgents"
  bash "$SCRIPT_DIR/install_local_model_guardrails.sh" >/dev/null
fi

echo "[model-hygiene] pruning legacy configured models"
current_models_json="$(openclaw config get agents.defaults.models 2>/dev/null || echo '{}')"
while IFS= read -r model; do
  [[ -z "$model" ]] && continue
  if ! in_array "$model" "${keep_models[@]-}"; then
    openclaw config unset "agents.defaults.models[$model]" >/dev/null || true
  fi
done < <(jq -r 'keys[]? // empty' <<<"$current_models_json")

for model in "${keep_models[@]}"; do
  openclaw config set --json "agents.defaults.models[$model]" '{}' >/dev/null || true
done

echo "[model-hygiene] pruning legacy aliases"
while IFS=$'\t' read -r alias target; do
  [[ -z "$alias" || -z "$target" ]] && continue
  if ! in_array "$target" "${keep_models[@]-}"; then
    openclaw models aliases remove "$alias" >/dev/null || true
  fi
done < <(openclaw models status --json 2>/dev/null | jq -r '.aliases // {} | to_entries[] | "\(.key)\t\(.value)"')

echo "[model-hygiene] pruning legacy provider blocks"
providers_json="$(openclaw config get models.providers 2>/dev/null || echo '{}')"
while IFS= read -r provider; do
  [[ -z "$provider" ]] && continue
  if ! in_array "$provider" "${keep_providers[@]-}"; then
    openclaw config unset "models.providers.$provider" >/dev/null || true
  fi
done < <(jq -r 'keys[]? // empty' <<<"$providers_json")

echo "[model-hygiene] pruning legacy auth profiles in openclaw config"
auth_profiles_json="$(openclaw config get auth.profiles 2>/dev/null || echo '{}')"
while IFS=$'\t' read -r profile_id provider; do
  [[ -z "$profile_id" || -z "$provider" ]] && continue
  if ! in_array "$provider" "${keep_providers[@]-}"; then
    openclaw config unset "auth.profiles[$profile_id]" >/dev/null || true
  fi
done < <(jq -r 'to_entries[]? | "\(.key)\t\(.value.provider // empty)"' <<<"$auth_profiles_json")

if [[ -f "$AUTH_STORE" ]]; then
  echo "[model-hygiene] pruning legacy auth profiles in auth store"
  keep_providers_json="$(printf '%s\n' "${keep_providers[@]-}" | jq -R . | jq -s .)"
  tmp_auth="$(mktemp)"
  jq --argjson keep "$keep_providers_json" '
    .profiles |= with_entries(select((.value.provider // "") as $p | ($keep | index($p)) != null))
    | .lastGood |= with_entries(select(.key as $p | ($keep | index($p)) != null))
    | .usageStats |= with_entries(select((.key | split(":")[0]) as $p | ($keep | index($p)) != null))
  ' "$AUTH_STORE" > "$tmp_auth"
  mv "$tmp_auth" "$AUTH_STORE"
fi

echo "[model-hygiene] complete"
echo "[model-hygiene] providers kept: ${keep_providers[*]}"
echo "[model-hygiene] models kept: ${keep_models[*]}"
echo ""
openclaw models status