#!/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)"

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

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

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 "${REPO_ROOT}/.env" ]]; then
    cp "${REPO_ROOT}/.env.example" "${REPO_ROOT}/.env"
    log "Created .env from .env.example."
    err "IMPORTANT: edit .env and set HF_TOKEN (and optionally VLLM_API_KEY) before production use."
  fi

  if [[ ! -f "${REPO_ROOT}/backend/config/model.env" ]]; then
    cp "${REPO_ROOT}/backend/config/model.env.example" "${REPO_ROOT}/backend/config/model.env"
    log "Created backend/config/model.env from example."
  fi

  if [[ ! -f "${REPO_ROOT}/frontend/config/frontend.env" ]]; then
    cp "${REPO_ROOT}/frontend/config/frontend.env.example" "${REPO_ROOT}/frontend/config/frontend.env"
    log "Created frontend/config/frontend.env from example."
  fi

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

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 "${REPO_ROOT}/.env" pull

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

show_status_and_urls() {
  local backend_port frontend_port
  backend_port="$(grep -E '^BACKEND_PORT=' "${REPO_ROOT}/.env" | tail -n1 | cut -d'=' -f2 || true)"
  frontend_port="$(grep -E '^FRONTEND_PORT=' "${REPO_ROOT}/.env" | tail -n1 | cut -d'=' -f2 || true)"
  backend_port="${backend_port:-8000}"
  frontend_port="${frontend_port:-3000}"

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

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

  printf '
'
  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 "$@"