#!/usr/bin/env zsh
# llm-rotate-model.sh — Download (if needed) and swap the primary model.
#
# Usage:
#   llm-rotate-model.sh                          # interactive prompt
#   llm-rotate-model.sh <hf-url-or-repo-id>      # e.g. https://huggingface.co/mlx-community/gemma-4-27b-it-4bit
#                                                 #      or mlx-community/gemma-4-27b-it-4bit
#   llm-rotate-model.sh <local-model-basename>   # model already in $MODEL_DIR
#
# The script:
#   1. Parses the HF URL/repo-id to derive a local basename
#   2. Downloads the model if not already in $MODEL_DIR
#   3. Updates llm.conf PRIMARY_MODEL
#   4. Restarts oMLX and verifies the new model is serving

set -euo pipefail
source "$(dirname "$0")/llm.conf"

# ── Helpers ───────────────────────────────────────────────────────────────────
_info() { echo "  $*"; }
_ok()   { echo "  ✓ $*"; }
_err()  { echo "  ✗ $*" >&2; }

# Parse a HF URL or repo-id into "owner/repo" form.
# Input:  https://huggingface.co/mlx-community/gemma-4-27b-it-4bit
#         mlx-community/gemma-4-27b-it-4bit
# Output: mlx-community/gemma-4-27b-it-4bit
_parse_hf_input() {
  local raw="$1"
  # Strip trailing slashes and any ?... or /tree/... suffixes
  raw="${raw%%\?*}"
  raw="${raw%%/tree/*}"
  raw="${raw%/}"
  # Strip scheme + host
  raw="${raw#https://huggingface.co/}"
  raw="${raw#http://huggingface.co/}"
  raw="${raw#huggingface.co/}"
  echo "$raw"
}

# ── Resolve input ─────────────────────────────────────────────────────────────
INPUT="${1:-}"

if [[ -z "$INPUT" ]]; then
  echo ""
  echo "Enter a Hugging Face URL or repo-id (e.g. mlx-community/gemma-4-27b-it-4bit),"
  echo "or a local model name already in $MODEL_DIR:"
  echo ""
  printf "  HF URL / repo-id / local name: "
  read -r INPUT
  echo ""
fi

[[ -z "$INPUT" ]] && { _err "No input provided."; exit 1; }

# Check if it looks like a local-only name (no slash → basename in MODEL_DIR)
if [[ "$INPUT" != *"/"* ]] || [[ -d "$MODEL_DIR/$INPUT" ]] || [[ -L "$MODEL_DIR/$INPUT" ]]; then
  REPO_ID=""
  NEW_MODEL_BASENAME="${INPUT##*/}"  # handle accidental full paths
else
  REPO_ID="$(_parse_hf_input "$INPUT")"
  # Validate it looks like owner/repo
  if [[ "$REPO_ID" != *"/"* ]]; then
    _err "Cannot parse '$INPUT' as a HF repo-id (expected owner/repo or full URL)."
    exit 1
  fi
  NEW_MODEL_BASENAME="${REPO_ID##*/}"
fi

NEW_MODEL_PATH="$MODEL_DIR/$NEW_MODEL_BASENAME"
OLD_MODEL_BASENAME="$(basename "$PRIMARY_MODEL")"

echo "╔══════════════════════════════════════════════╗"
echo "║  Model Rotation"
echo "╠══════════════════════════════════════════════╣"
printf "║  OLD: %-38s║\n" "$OLD_MODEL_BASENAME"
printf "║  NEW: %-38s║\n" "$NEW_MODEL_BASENAME"
[[ -n "$REPO_ID" ]] && printf "║  SRC: %-38s║\n" "hf.co/$REPO_ID"
echo "╚══════════════════════════════════════════════╝"
echo ""

[[ "$OLD_MODEL_BASENAME" == "$NEW_MODEL_BASENAME" ]] && {
  _ok "primary is already $NEW_MODEL_BASENAME — nothing to do."
  exit 0
}

# ── 1. Download if needed ─────────────────────────────────────────────────────
if [[ -d "$NEW_MODEL_PATH" ]] || [[ -L "$NEW_MODEL_PATH" ]]; then
  _ok "Model already in $MODEL_DIR — skipping download."
else
  if [[ -z "$REPO_ID" ]]; then
    _err "Model not found locally: $NEW_MODEL_PATH"
    _err "Provide a HF URL or repo-id to download it."
    exit 1
  fi

  echo "[1/4] Downloading $REPO_ID → $NEW_MODEL_PATH"

  if command -v huggingface-cli >/dev/null 2>&1; then
    mkdir -p "$NEW_MODEL_PATH"
    huggingface-cli download "$REPO_ID" --local-dir "$NEW_MODEL_PATH" --local-dir-use-symlinks False \
      && _ok "Download complete." \
      || { _err "huggingface-cli download failed."; rm -rf "$NEW_MODEL_PATH"; exit 1; }
  else
    _err "huggingface-cli not found. Install with: pip install huggingface-hub"
    _err "Then re-run this script."
    exit 1
  fi
fi

# ── 2. Update llm.conf ────────────────────────────────────────────────────────
STEP=2
[[ -z "$REPO_ID" ]] && STEP=1  # no download step occurred
echo "[$STEP/$(( STEP == 1 ? 3 : 4 ))] Updating llm.conf..."
CONF_FILE="$(dirname "$0")/llm.conf"
sed -i '' "s|PRIMARY_MODEL=.*|PRIMARY_MODEL=\$MODEL_DIR/${NEW_MODEL_BASENAME}|" "$CONF_FILE"
_ok "PRIMARY_MODEL → \$MODEL_DIR/$NEW_MODEL_BASENAME"

# ── 3. Restart oMLX ──────────────────────────────────────────────────────────
STEP=$(( STEP + 1 ))
echo "[$STEP/$(( STEP == 2 ? 3 : 4 ))] Restarting oMLX on :$PRIMARY_PORT..."

if [[ -f "$PRIMARY_PID" ]] && kill -0 "$(cat "$PRIMARY_PID")" 2>/dev/null; then
  kill "$(cat "$PRIMARY_PID")" 2>/dev/null
  sleep 2
  kill -0 "$(cat "$PRIMARY_PID")" 2>/dev/null && kill -9 "$(cat "$PRIMARY_PID")" 2>/dev/null || true
elif lsof -ti tcp:"$PRIMARY_PORT" >/dev/null 2>&1; then
  kill $(lsof -ti tcp:"$PRIMARY_PORT") 2>/dev/null || true
  sleep 2
fi

source "$CONF_FILE"

nohup "$RMLX_BIN" serve "$PRIMARY_MODEL" \
  --port "$PRIMARY_PORT" \
  --memory-guard-gb 25 \
  --hot-cache-max-size 6GB \
  --initial-cache-blocks 1024 \
  --max-concurrent-requests 2 \
  >> "$PRIMARY_LOG" 2>&1 &

echo $! > "$PRIMARY_PID"
_info "PID $(cat "$PRIMARY_PID")"

printf "  Waiting for :$PRIMARY_PORT "
for i in $(seq 1 60); do
  if curl -sf "http://localhost:$PRIMARY_PORT/v1/models" >/dev/null 2>&1; then
    echo " ready"
    break
  fi
  printf "."
  sleep 2
done

# ── 4. Verify ─────────────────────────────────────────────────────────────────
STEP=$(( STEP + 1 ))
echo "[$STEP/$STEP] Verifying..."

SERVED=$(curl -sf "http://localhost:$PRIMARY_PORT/v1/models" \
  | python3 -c "import sys,json;print(json.load(sys.stdin)['data'][0]['id'])" 2>/dev/null || echo "")

if [[ -n "$SERVED" ]]; then
  _ok "Server reports: $SERVED"
else
  _err "Could not read model ID from /v1/models"
  exit 1
fi

RESPONSE=$(curl -sf "http://localhost:$PRIMARY_PORT/v1/chat/completions" \
  -H "Content-Type: application/json" \
  -d "{\"model\":\"$SERVED\",\"messages\":[{\"role\":\"user\",\"content\":\"Say OK\"}],\"max_tokens\":3,\"stream\":false}" \
  | python3 -c "import sys,json;print(json.load(sys.stdin)['choices'][0]['message']['content'][:20])" 2>/dev/null || echo "")

[[ -n "$RESPONSE" ]] && _ok "Smoke test: \"$RESPONSE\"" || _err "Smoke test got no response — check $PRIMARY_LOG"

echo ""
echo "Done. Primary rotated to $NEW_MODEL_BASENAME"