#!/usr/bin/env bash
# Installs prerequisites (if needed), prepares config files, and starts Gemma 3 + vLLM + chat UI stack.
set -euo pipefail

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

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

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

load_env_file() {
  if [[ ! -f "${ENV_FILE}" ]]; then
    return
  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() {
  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."
}

require_linux() {
  if [[ "$(uname -s)" != "Linux" ]]; then
    err "This script supports Linux only."
    exit 1
  fi
}

install_docker_ubuntu() {
  log "Installing Docker Engine and Compose plugin using official Docker apt repository."
  sudo apt-get update
  sudo apt-get install -y ca-certificates curl gnupg
  sudo install -m 0755 -d /etc/apt/keyrings

  if [[ ! -f /etc/apt/keyrings/docker.gpg ]]; then
    curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg
    sudo chmod a+r /etc/apt/keyrings/docker.gpg
  fi

  source /etc/os-release
  local arch
  arch="$(dpkg --print-architecture)"
  local codename
  codename="${VERSION_CODENAME:-jammy}"

  echo "deb [arch=${arch} signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu ${codename} stable" \
    | sudo tee /etc/apt/sources.list.d/docker.list >/dev/null

  sudo apt-get update
  sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

  if ! sudo systemctl is-active --quiet docker; then
    sudo systemctl enable --now docker
  fi

  if ! getent group docker >/dev/null; then
    sudo groupadd docker
  fi

  if ! id -nG "${USER}" | grep -qw docker; then
    sudo usermod -aG docker "${USER}"
    log "Added ${USER} to docker group. You may need to log out and back in for group membership to apply."
  fi
}

check_or_install_docker() {
  local have_docker=1
  local have_compose=1

  if ! command -v docker >/dev/null 2>&1; then
    have_docker=0
  fi

  if ! docker compose version >/dev/null 2>&1; then
    have_compose=0
  fi

  if [[ ${have_docker} -eq 1 && ${have_compose} -eq 1 ]]; then
    log "Docker and Compose plugin are already available."
    return
  fi

  if [[ -f /etc/os-release ]]; then
    source /etc/os-release
    if [[ "${ID:-}" == "ubuntu" ]]; then
      install_docker_ubuntu
      return
    fi
  fi

  err "Docker/Compose missing and automatic installation is implemented for Ubuntu only."
  err "See docs/TROUBLESHOOTING.md#docker-and-compose-not-available"
  exit 1
}

prepare_env_files() {
  if [[ ! -f "${ENV_FILE}" ]]; then
    cp "${REPO_ROOT}/.env.example" "${ENV_FILE}"
    log "Created .env from .env.example."
    err "IMPORTANT: edit .env and set HF_TOKEN, VLLM_API_KEY, GEMMA_MODEL_ID, and other values before production use."
  fi

  mkdir -p     "${REPO_ROOT}/backend/config"     "${REPO_ROOT}/frontend/config"     "${REPO_ROOT}/models"     "${REPO_ROOT}/frontend/data/open-webui"

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

warn_if_rocm_devices_missing() {
  if [[ ! -e /dev/kfd || ! -d /dev/dri ]]; then
    err "ROCm device files /dev/kfd or /dev/dri are not available."
    err "See docs/TROUBLESHOOTING.md#rocm-devices-not-visible-in-host"
  fi
}

start_stack() {
  log "Pulling container images."
  docker compose -f "${REPO_ROOT}/docker-compose.yml" --env-file "${ENV_FILE}" pull

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

show_status_and_urls() {
  local backend_port frontend_port
  backend_port="${BACKEND_PORT:-8000}"
  frontend_port="${FRONTEND_PORT:-3000}"

  log "Backend status:"
  docker compose -f "${REPO_ROOT}/docker-compose.yml" --env-file "${ENV_FILE}" ps gemma3-vllm || true

  log "Frontend status:"
  docker compose -f "${REPO_ROOT}/docker-compose.yml" --env-file "${ENV_FILE}" ps chat-ui || true

  printf '\n'
  log "API endpoint: http://localhost:${backend_port}/v1"
  log "Chat UI endpoint: http://localhost:${frontend_port}"
  log "If startup fails, inspect logs with: docker compose logs --tail=200 gemma3-vllm chat-ui"
}

main() {
  require_linux
  check_or_install_docker
  prepare_env_files
  warn_if_rocm_devices_missing
  start_stack
  show_status_and_urls
}

main "$@"