#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
ANSIBLE_VENV="$ROOT_DIR/.venv-ansible"
ANSIBLE_PYTHON="$ANSIBLE_VENV/bin/python"
ANSIBLE_PLAYBOOK="$ANSIBLE_VENV/bin/ansible-playbook"
export ANSIBLE_CONFIG="$ROOT_DIR/ansible/ansible.cfg"

run_project_script() {
  local script_path=$1
  shift || true
  bash "$script_path" "$@"
}

usage() {
  cat <<'EOF'
Usage:
  ./labctl up
  ./labctl down
  ./labctl update_wiki
  ./labctl ollama_models
  ./labctl preflight
  ./labctl versions
  ./labctl assets lab2 [--refresh]
  ./labctl start [core|all|service...]
  ./labctl stop [all|service...]
  ./labctl status [all|service...]
  ./labctl urls
  ./labctl open kiln
  ./labctl logs <service>
EOF
}

netron_version() {
  local version_file=$ROOT_DIR/ansible/group_vars/all.yml

  if [ ! -f "$version_file" ]; then
    printf '%s\n' "unknown"
    return
  fi

  sed -nE 's/^courseware_netron_version:[[:space:]]*"([^"]+)".*/\1/p' "$version_file" | head -n 1
}

minimum_ollama_version() {
  local version_file=$ROOT_DIR/ansible/group_vars/all.yml

  if [ ! -f "$version_file" ]; then
    printf '%s\n' "unknown"
    return
  fi

  sed -nE 's/^courseware_ollama_min_version:[[:space:]]*"([^"]+)".*/\1/p' "$version_file" | head -n 1
}

print_versions() {
  cat <<EOF
Pinned component versions:
  Netron: $(netron_version)
  Minimum Ollama: $(minimum_ollama_version)
  Ansible Core: 2.18.6
EOF
}

confirm_installation() {
  local response
  local pinned_netron
  local min_ollama
  pinned_netron=$(netron_version)
  min_ollama=$(minimum_ollama_version)

  if [ ! -t 0 ]; then
    cat <<EOF >&2
WARNING: THIS SCRIPT WILL CONFIGURE YOUR ENVIRONMENT WILL THE FOLLOWING SOFTWARE:

- Ollama
- llama.cpp
- Netron (${pinned_netron})
- Open WebUI
- ChunkViz
- Embedding Atlas
- Promptfoo
- Unsloth Studio
- Kiln Desktop
- Course-specific support assets for lab 1, lab 2, and lab 4
- Pre-pulled Gemma 4 E2B Ollama models for Lab 1 and Lab 2
- Lab 1 confidence support through Gemma 4 E2B Q4 (requires Ollama ${min_ollama}+)

IT IS RECOMMENDED TO RUN THIS IN AN ISLOATED ENVIRONMENT (Dedicated WSL, VM, etc.)

This process may take a long time.

This command requires interactive confirmation. Re-run it from a terminal and answer the prompt.
EOF
    exit 1
  fi

  cat <<EOF
WARNING: THIS SCRIPT WILL CONFIGURE YOUR ENVIRONMENT WILL THE FOLLOWING SOFTWARE:

- Ollama
- llama.cpp
- Netron (${pinned_netron})
- Open WebUI
- ChunkViz
- Embedding Atlas
- Promptfoo
- Unsloth Studio
- Kiln Desktop
- Course-specific support assets for lab 1, lab 2, and lab 4
- Pre-pulled Gemma 4 E2B Ollama models for Lab 1 and Lab 2
- Lab 1 confidence support through Gemma 4 E2B Q4 (requires Ollama ${min_ollama}+)

IT IS RECOMMENDED TO RUN THIS IN AN ISLOATED ENVIRONMENT (Dedicated WSL, VM, etc.)

This process may take a long time.
EOF

  read -r -p "CONFIRM INSTALLATION (y/N): " response
  case "$response" in
    y|Y)
      ;;
    *)
      echo "Installation cancelled."
      exit 1
      ;;
  esac
}

host_is_wsl() {
  [ "$(uname -s)" = "Linux" ] && uname -r | grep -qiE 'microsoft|wsl'
}

host_is_macos() {
  [ "$(uname -s)" = "Darwin" ]
}

host_is_linux() {
  [ "$(uname -s)" = "Linux" ]
}

host_is_arm_mac() {
  host_is_macos && [ "$(uname -m)" = "arm64" ]
}

host_profile() {
  if host_is_wsl; then
    printf '%s\n' "wsl"
    return
  fi

  if host_is_macos; then
    printf '%s\n' "macos"
    return
  fi

  if host_is_linux && host_is_debian_family; then
    printf '%s\n' "native-debian-ubuntu"
    return
  fi

  printf '%s\n' "unsupported"
}

check_wsl_gpu_readiness() {
  if [ "$(host_profile)" != "wsl" ]; then
    return
  fi

  if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi >/dev/null 2>&1; then
    return
  fi

  cat <<'EOF' >&2
WSL GPU support is not ready yet, so the installer is stopping before Ansible runs.

This courseware expects WSL to already see your NVIDIA GPU.

Please do this on the Windows side first:
1. Install or update the current NVIDIA Windows driver with WSL/CUDA support.
2. Open Windows PowerShell and run: wsl --update
3. Fully restart WSL: wsl --shutdown
4. Reopen your Linux distro and confirm this works: nvidia-smi

If `nvidia-smi` still fails inside WSL:
- Reboot Windows
- Confirm WSL 2 is in use
- Confirm the GPU is NVIDIA and the Windows driver is recent

After `nvidia-smi` works inside WSL, rerun:
  bash deploy-courseware.sh
EOF
  exit 1
}

linux_cuda_toolkit_is_ready() {
  command -v nvcc >/dev/null 2>&1 && return 0
  [ -x /usr/local/cuda/bin/nvcc ] && return 0
  for candidate in /usr/local/cuda-*/bin/nvcc; do
    if [ -x "$candidate" ]; then
      return 0
    fi
  done
  [ -f /usr/local/cuda/include/cuda_runtime.h ] && return 0
  [ -f /usr/include/cuda_runtime.h ] && return 0
  return 1
}

linux_cuda_toolkit_package_is_available() {
  command -v apt-cache >/dev/null 2>&1 || return 1
  apt-cache show nvidia-cuda-toolkit >/dev/null 2>&1
}

print_linux_cuda_toolkit_guidance() {
  local package_hint
  package_hint="sudo apt update && sudo apt install -y nvidia-cuda-toolkit"

  if command -v apt-cache >/dev/null 2>&1; then
    if ! apt-cache show nvidia-cuda-toolkit >/dev/null 2>&1; then
      package_hint="Your distro does not expose nvidia-cuda-toolkit in its default apt sources, so add NVIDIA's CUDA repository for your Debian/Ubuntu release and install the toolkit from there."
    fi
  fi

  if host_is_wsl; then
    cat <<EOF >&2
CUDA Toolkit is still missing inside this WSL Linux environment, so the installer is stopping before Ansible runs.

WSL splits GPU support into two layers:
- Windows side: the NVIDIA driver makes the GPU visible to WSL
- Linux side: llama.cpp still needs the CUDA toolkit headers/compiler inside the distro

What to do:
1. Confirm the driver side works in WSL: nvidia-smi
2. Install the Linux-side CUDA toolkit
   $package_hint
3. Verify the toolkit:
   nvcc --version
   ls /usr/local/cuda/include/cuda_runtime.h

After that, rerun:
  bash deploy-courseware.sh
EOF
    return
  fi

  cat <<EOF >&2
CUDA Toolkit is not installed inside this Linux environment, so the installer is stopping before Ansible runs.

This courseware can only build CUDA-enabled llama.cpp when the Linux-side toolkit is present.

What to do:
1. Install the CUDA toolkit
   $package_hint
2. Verify the toolkit:
   nvcc --version
   ls /usr/local/cuda/include/cuda_runtime.h

After that, rerun:
  bash deploy-courseware.sh
EOF
}

check_linux_cuda_toolkit() {
  if [ "$(host_profile)" != "native-debian-ubuntu" ]; then
    return
  fi

  if linux_cuda_toolkit_is_ready; then
    return
  fi

  if linux_cuda_toolkit_package_is_available; then
    return
  fi

  print_linux_cuda_toolkit_guidance
  exit 1
}

resolve_python() {
  if [ -n "${PYTHON_BIN:-}" ] && command -v "${PYTHON_BIN}" >/dev/null 2>&1; then
    command -v "${PYTHON_BIN}"
    return
  fi

  if command -v python3 >/dev/null 2>&1; then
    command -v python3
    return
  fi

  if command -v python >/dev/null 2>&1; then
    command -v python
    return
  fi

  cat <<'EOF' >&2
Python 3 was not found.

Install it first, then rerun this command:
- Debian/Ubuntu/WSL: sudo apt update && sudo apt install -y python3 python3-venv
- macOS: brew install python@3.11
EOF
  exit 1
}

host_is_debian_family() {
  [ -f /etc/os-release ] || return 1
  grep -qiE '^(ID|ID_LIKE)=(.*debian|debian.*)$' /etc/os-release
}

python_has_venv_support() {
  local python_bin=$1
  local probe_parent
  local probe_venv

  probe_parent=$(mktemp -d 2>/dev/null || mktemp -d -t labctl-venv-probe)
  probe_venv="$probe_parent/venv"

  if "$python_bin" -m venv "$probe_venv" >/dev/null 2>&1; then
    rm -rf "$probe_parent"
    return 0
  fi

  rm -rf "$probe_parent"
  return 1
}

install_debian_python_venv_support() {
  local python_bin=$1
  local installer=(apt-get)
  local versioned_venv_pkg
  local packages=(python3-venv python3-pip)

  if ! command -v apt-get >/dev/null 2>&1; then
    return 1
  fi

  if [ "$(id -u)" -ne 0 ]; then
    if ! command -v sudo >/dev/null 2>&1; then
      cat <<'EOF' >&2
Python can be found, but this Debian/Ubuntu system is missing the venv bootstrap packages that Ansible needs.

Install them, then rerun this command:
  apt-get update && apt-get install -y python3-venv python3-pip
EOF
      return 1
    fi

    if [ ! -t 0 ]; then
      cat <<'EOF' >&2
Python can be found, but this Debian/Ubuntu system still needs sudo access to install python3-venv and python3-pip.

Run this command from an interactive terminal so sudo can prompt for your password, or install the packages manually:
  sudo apt-get update && sudo apt-get install -y python3-venv python3-pip
EOF
      return 1
    fi

    echo "This step needs sudo to install missing Python packages. You may be prompted for your password."
    installer=(sudo apt-get)
  fi

  versioned_venv_pkg=$("$python_bin" -c "import sys; print(f'python{sys.version_info.major}.{sys.version_info.minor}-venv')")
  if command -v apt-cache >/dev/null 2>&1 && apt-cache show "$versioned_venv_pkg" >/dev/null 2>&1; then
    packages=("$versioned_venv_pkg" "${packages[@]}")
  fi

  echo "Installing missing Python venv support for this Debian/Ubuntu system..."
  "${installer[@]}" update
  "${installer[@]}" install -y "${packages[@]}"
}

ansible_venv_is_usable() {
  if [ ! -x "$ANSIBLE_PYTHON" ]; then
    return 1
  fi

  if [ ! -x "$ANSIBLE_PLAYBOOK" ]; then
    return 1
  fi

  "$ANSIBLE_PYTHON" -c "import ansible" >/dev/null 2>&1 || return 1
  "$ANSIBLE_PLAYBOOK" --version >/dev/null 2>&1
}

rebuild_ansible() {
  local python_bin

  python_bin=$(resolve_python)

  if ! python_has_venv_support "$python_bin"; then
    if host_is_debian_family; then
      install_debian_python_venv_support "$python_bin"
    fi
  fi

  if ! python_has_venv_support "$python_bin"; then
    cat <<'EOF' >&2
Python 3 is installed, but its virtual environment support is still unavailable.

Install the missing venv package for your platform, then rerun this command:
- Debian/Ubuntu/WSL: sudo apt update && sudo apt install -y python3-venv python3-pip
- macOS: brew reinstall python@3.11
EOF
    exit 1
  fi

  rm -rf "$ANSIBLE_VENV"
  echo "Preparing local installer runtime..."
  "$python_bin" -m venv "$ANSIBLE_VENV"
  "$ANSIBLE_PYTHON" -m pip install --upgrade pip
  "$ANSIBLE_PYTHON" -m pip install "ansible-core==2.18.6"
}

ensure_ansible() {
  if ansible_venv_is_usable; then
    return
  fi

  if [ -e "$ANSIBLE_VENV" ]; then
    echo "Refreshing installer runtime for this machine..."
  fi

  rebuild_ansible
}

sudo_keepalive_pid=""

ensure_sudo_session() {
  if [ "$(uname -s)" != "Linux" ]; then
    return
  fi

  if [ "$(id -u)" -eq 0 ]; then
    return
  fi

  if ! command -v sudo >/dev/null 2>&1; then
    cat <<'EOF' >&2
This installer needs sudo for Linux package setup, but `sudo` is not installed.

Install sudo or rerun this command as root, then try again.
EOF
    exit 1
  fi

  if sudo -n true >/dev/null 2>&1; then
    return
  fi

  if [ ! -t 0 ]; then
    cat <<'EOF' >&2
This installer needs sudo for Linux package setup.

Run this command from an interactive terminal so sudo can prompt for your password, then rerun the same `./labctl` command.
EOF
    exit 1
  fi

  echo "This step needs sudo for Linux package setup. You may be prompted for your password."
  sudo -v
}

start_sudo_keepalive() {
  if [ "$(uname -s)" != "Linux" ]; then
    return
  fi

  if [ "$(id -u)" -eq 0 ]; then
    return
  fi

  if ! command -v sudo >/dev/null 2>&1; then
    return
  fi

  (
    while true; do
      sudo -n true >/dev/null 2>&1 || exit 0
      sleep 60
    done
  ) &
  sudo_keepalive_pid=$!
}

stop_sudo_keepalive() {
  if [ -n "${sudo_keepalive_pid:-}" ]; then
    kill "$sudo_keepalive_pid" >/dev/null 2>&1 || true
    wait "$sudo_keepalive_pid" >/dev/null 2>&1 || true
    sudo_keepalive_pid=""
  fi
}

run_playbook() {
  local playbook=$1
  local escaped_root
  shift

  check_wsl_gpu_readiness
  check_linux_cuda_toolkit
  ensure_ansible
  ensure_sudo_session

  if [ ! -x "$ANSIBLE_PLAYBOOK" ]; then
    echo "Installer runtime is incomplete. Rebuilding it..."
    rebuild_ansible
  fi

  escaped_root=${ROOT_DIR//\\/\\\\}
  escaped_root=${escaped_root//\"/\\\"}

  start_sudo_keepalive
  trap stop_sudo_keepalive RETURN
  "$ANSIBLE_PLAYBOOK" \
    -e "{\"courseware_root\":\"$escaped_root\"}" \
    "$ROOT_DIR/ansible/playbooks/$playbook" \
    "$@"
}

require_arg() {
  if [ $# -lt 1 ]; then
    usage
    exit 1
  fi
}

require_managed_runtime() {
  if [ ! -f "$ROOT_DIR/state/runtime.env" ]; then
    cat <<'EOF' >&2
Missing state/runtime.env. Run ./labctl up first so the managed environment exists before using this command.
EOF
    exit 1
  fi
}

handle_assets_command() {
  local asset_group=${1:-}
  shift || true

  case "$asset_group" in
    lab2)
      exec bash "$ROOT_DIR/scripts/bootstrap_lab2_assets.sh" "$@"
      ;;
    *)
      usage
      exit 1
      ;;
  esac
}

refresh_ollama_models() {
  require_managed_runtime
  run_playbook up.yml --tags ollama_models
}

update_wiki() {
  require_managed_runtime
  run_playbook up.yml --tags wiki -e "courseware_wiki_force_update=true"
  run_project_script "$ROOT_DIR/scripts/service_manager.sh" restart-wiki
}

main() {
  local cmd=${1:-}
  shift || true

  case "$cmd" in
    up)
      confirm_installation
      run_playbook up.yml
      run_project_script "$ROOT_DIR/scripts/service_manager.sh" start all
      ;;
    ollama_models)
      refresh_ollama_models
      ;;
    update_wiki)
      update_wiki
      ;;
    down)
      run_project_script "$ROOT_DIR/scripts/service_manager.sh" stop all || true
      run_playbook down.yml
      ;;
    preflight)
      confirm_installation
      run_playbook up.yml --tags preflight
      ;;
    versions)
      print_versions
      ;;
    assets)
      require_arg "$@"
      handle_assets_command "$@"
      ;;
    start|stop|status|urls|open|logs)
      exec bash "$ROOT_DIR/scripts/service_manager.sh" "$cmd" "$@"
      ;;
    ""|-h|--help|help)
      usage
      ;;
    *)
      usage
      exit 1
      ;;
  esac
}

main "$@"