#!/usr/bin/env bash
# Restarts the Gemma 3 vLLM stack and shows service status.
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
ENV_FILE="${REPO_ROOT}/.env"

log() {
  printf '[restart] %s\n' "$*"
}

err() {
  printf '[restart][error] %s\n' "$*" >&2
}

load_env_file() {
  if [[ ! -f "${ENV_FILE}" ]]; then
    err ".env file not found. Copy .env.example to .env first."
    exit 1
  fi

  set -a
  # shellcheck disable=SC1090
  source "${ENV_FILE}"
  set +a

  : "${GEMMA_MODEL_ID:?GEMMA_MODEL_ID must be set in .env}"
}

sync_service_env_files() {
  mkdir -p "${REPO_ROOT}/backend/config" "${REPO_ROOT}/frontend/config"

  cat > "${REPO_ROOT}/backend/config/model.env" <<EOF
HF_TOKEN=${HF_TOKEN:-YOUR_HF_TOKEN_HERE}
VLLM_API_KEY=${VLLM_API_KEY:-YOUR_LOCAL_API_KEY_HERE}
GEMMA_MODEL_ID=${GEMMA_MODEL_ID}
BACKEND_PORT=${BACKEND_PORT:-8000}
HUGGINGFACE_CACHE_DIR=${HUGGINGFACE_CACHE_DIR:-/home/${USER}/.cache/huggingface}
VLLM_MAX_MODEL_LEN=${VLLM_MAX_MODEL_LEN:-512}
VLLM_GPU_MEMORY_UTILIZATION=${VLLM_GPU_MEMORY_UTILIZATION:-0.7}
EOF

  cat > "${REPO_ROOT}/frontend/config/frontend.env" <<EOF
FRONTEND_PORT=${FRONTEND_PORT:-3000}
OPENAI_API_BASE_URL=${OPENAI_API_BASE_URL:-http://gemma3-vllm:8000/v1}
VLLM_API_KEY=${VLLM_API_KEY:-YOUR_LOCAL_API_KEY_HERE}
GEMMA_MODEL_ID=${GEMMA_MODEL_ID}
OPEN_WEBUI_DATA_DIR=${OPEN_WEBUI_DATA_DIR:-./frontend/data/open-webui}
EOF

  log "Synced backend/config/model.env and frontend/config/frontend.env from .env."
}

load_env_file
log "Using GEMMA_MODEL_ID=${GEMMA_MODEL_ID}"
sync_service_env_files

log "Stopping stack."
docker compose -f "${REPO_ROOT}/docker-compose.yml" --env-file "${ENV_FILE}" down

log "Starting stack."
docker compose -f "${REPO_ROOT}/docker-compose.yml" --env-file "${ENV_FILE}" up -d --force-recreate

log "Current status:"
docker compose -f "${REPO_ROOT}/docker-compose.yml" --env-file "${ENV_FILE}" ps