#!/usr/bin/env bash
# Switch between stable / dev(inference) / training modes
# Usage: bash switch_mode.sh [stable|dev|inference|training]

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
LOG_DIR="${PROJECT_ROOT}/llm_service/logs"
MODE="${1:-}"

REDIS_HOST="${REDIS_HOST:-localhost}"
REDIS_PORT="${REDIS_PORT:-6379}"
REDIS_DB="${REDIS_DB:-0}"
DB_URL="postgresql://user:kAvtYfAUMmwQgezaXz6Scw@localhost:5433/base_platform"

if [[ -z "$MODE" ]]; then
  echo "Usage: $0 [stable|dev|inference|training]"
  exit 1
fi

# ─── Redis ────────────────────────────────────────────────────────────────────

set_redis_mode() {
  if command -v redis-cli &> /dev/null; then
    redis-cli -h "$REDIS_HOST" -p "$REDIS_PORT" -n "$REDIS_DB" SET "llm:service:status" "$1" > /dev/null 2>&1 || true
    echo "   Redis status: $1"
  fi
}

# ─── DB ───────────────────────────────────────────────────────────────────────

update_db_models() {
  local mode="$1"
  local PYTHON=/lsinfo/ai/hellotax_ai/base_platform/venv/bin/python
  if [[ "$mode" == "stable" ]]; then
    $PYTHON - <<PYEOF
import psycopg2
conn = psycopg2.connect("$DB_URL")
cur = conn.cursor()
cur.execute("UPDATE models SET name = 'Qwen3.5-27B-AWQ (本地)', remote_model_id = 'Qwen3.5-27B-AWQ' WHERE code = 'local-qwen-chat'")
cur.execute("UPDATE models SET name = 'bge-m3 (本地)', remote_model_id = 'bge-m3', dimension = 1024 WHERE code = 'local-bge-m3'")
cur.execute("UPDATE models SET name = 'bge-reranker-v2-m3 (本地)', remote_model_id = 'bge-reranker-v2-m3' WHERE code = 'local-bge-reranker'")
conn.commit(); conn.close()
PYEOF
  elif [[ "$mode" == "dev" ]]; then
    $PYTHON - <<PYEOF
import psycopg2
conn = psycopg2.connect("$DB_URL")
cur = conn.cursor()
cur.execute("UPDATE models SET name = 'Qwen3.6-27B (本地)', remote_model_id = 'Qwen3.6-27B' WHERE code = 'local-qwen-chat'")
cur.execute("UPDATE models SET name = 'Qwen3-Embedding-8B (本地)', remote_model_id = 'Qwen3-Embedding-8B', dimension = 4096 WHERE code = 'local-bge-m3'")
cur.execute("UPDATE models SET name = 'Qwen3-Reranker-8B (本地)', remote_model_id = 'Qwen3-Reranker-8B' WHERE code = 'local-bge-reranker'")
conn.commit(); conn.close()
PYEOF
  fi
  echo "   DB model names updated for mode: $mode"
}

# ─── Stop helpers ─────────────────────────────────────────────────────────────

stop_service() {
  local name="$1" pid_file="$2" pattern="$3"
  echo ">> Stopping ${name}..."
  if [[ -f "$pid_file" ]]; then
    local pid; pid=$(cat "$pid_file")
    if kill -0 "$pid" 2>/dev/null; then
      kill "$pid"
      echo "   Stopped ${name} (PID $pid)"
    fi
    rm -f "$pid_file"
  else
    pkill -f "$pattern" 2>/dev/null || true
    echo "   ${name} stopped (by name match)"
  fi
  sleep 2
}

stop_vllm()      { stop_service "vLLM (8100)"      "${LOG_DIR}/vllm-8100.pid"        "vllm.*8100"; }
stop_embedding() { stop_service "Embedding (8200)"  "${LOG_DIR}/embedding-8200.pid"   "embedding_server.*8200"; }
stop_reranker()  { stop_service "Reranker (8300)"   "${LOG_DIR}/reranker-8300.pid"    "rerank_server.*8300"; }
stop_ocr()       { stop_service "OCR (8400)"        "${LOG_DIR}/ocr-8400.pid"         "vllm.*8400"; }
stop_llamafactory() { stop_service "LLaMA-Factory (7860)" "${LOG_DIR}/llamafactory-7860.pid" "llamafactory.*webui\|train_web_ui"; }

# ─── Modes ────────────────────────────────────────────────────────────────────

if [[ "$MODE" == "stable" ]]; then
  echo "=========================================="
  echo "Switching to STABLE MODE"
  echo "  LLM:       Qwen3.5-27B-AWQ"
  echo "  Embedding: bge-m3"
  echo "  Reranker:  bge-reranker-v2-m3"
  echo "=========================================="
  set_redis_mode "switching"
  stop_ocr
  stop_vllm
  stop_embedding
  stop_reranker
  stop_llamafactory
  echo ""
  bash "${SCRIPT_DIR}/start_vllm_stable.sh"
  bash "${SCRIPT_DIR}/start_embedding_stable.sh"
  bash "${SCRIPT_DIR}/start_reranker_stable.sh"
  update_db_models "stable"
  set_redis_mode "stable"
  echo ""
  echo "=========================================="
  echo "Stable mode active:"
  echo "  Chat:      http://localhost:8100/v1/chat/completions"
  echo "  Embedding: http://localhost:8200/v1/embeddings"
  echo "  Rerank:    http://localhost:8300/v1/rerank"
  echo "=========================================="

elif [[ "$MODE" == "dev" || "$MODE" == "inference" ]]; then
  echo "=========================================="
  echo "Switching to DEV MODE"
  echo "  LLM:       Qwen3.6-27B"
  echo "  Embedding: Qwen3-Embedding-8B"
  echo "  Reranker:  Qwen3-Reranker-8B"
  echo "  OCR:       Qwen3-VL-8B-Instruct (常驻)"
  echo "=========================================="
  set_redis_mode "switching"
  stop_ocr
  stop_vllm
  stop_embedding
  stop_reranker
  stop_llamafactory
  echo ""
  bash "${SCRIPT_DIR}/start_vllm.sh"
  bash "${SCRIPT_DIR}/start_embedding_vllm.sh"
  bash "${SCRIPT_DIR}/start_reranker_vllm.sh"
  bash "${SCRIPT_DIR}/start_ocr.sh"
  update_db_models "dev"
  set_redis_mode "dev"
  echo ""
  echo "=========================================="
  echo "Dev mode active:"
  echo "  Chat:      http://localhost:8100/v1/chat/completions"
  echo "  Embedding: http://localhost:8200/v1/embeddings"
  echo "  Rerank:    http://localhost:8300/v1/rerank"
  echo "  OCR:       http://localhost:8400/v1/chat/completions"
  echo "=========================================="

elif [[ "$MODE" == "training" ]]; then
  echo "=========================================="
  echo "Switching to TRAINING MODE"
  echo "=========================================="
  set_redis_mode "switching"
  stop_ocr
  stop_vllm
  stop_embedding
  stop_reranker
  echo ""
  echo ">> GPU memory freed. Starting LLaMA-Factory WebUI..."
  bash "${SCRIPT_DIR}/start_llamafactory_webui.sh"
  set_redis_mode "training"
  echo ""
  echo "=========================================="
  echo "Training mode active:"
  echo "  LLaMA-Factory WebUI: http://localhost:7860"
  echo "=========================================="

else
  echo "Unknown mode: $MODE"
  echo "Usage: $0 [stable|dev|inference|training]"
  exit 1
fi