#!/usr/bin/env bash
# run_iiab_code.sh
#
# Modes:
#   --run (default)              Launch if needed, seed UNITTEST local_vars, run install.txt, then auto-resume (sudo iiab -f)
#   --continue-after-reboot      Only resume: sudo iiab -f (no long wait; just run)
#   --clean                      Stop+delete+purge target VMs (by BASE-<number>, regardless of COUNT)
# Distro selection:
#   --debian-13          Debian 13 only (sets IMAGE=$DEBIAN13_IMAGE_URL and BASE=deb13)
#   --debian-13          Debian 13 only (sets IMAGE=$DEBIAN13_IMAGE_URL and BASE=deb13)
#   --both-distros       Run Ubuntu + Debian 13 in parallel: COUNT=N => 2N VMs (default order: interleaved)
#   --first-ubuntu       (with --both-distros) order: all Ubuntu first, then all Debian
#   --first-debian       (with --both-distros) order: all Debian first, then all Ubuntu
#
# PR selection:
#   --pr 4122            (repeatable) add PR numbers passed to install.txt
#   --run-pr 4122        same as --pr but also forces --run (alias: --test-pr)
#
# Env vars:
#   IIAB_PR="4122 4191"          Space-separated PRs
#   (compat) IIAB_INSTALL_ARGS   If set, used as fallback for IIAB_PR
#
# VM naming:
#   BASE=ubu2404, COUNT=3 => ubu2404-0 ubu2404-1 ubu2404-2
#
# Parallel + stagger:
#   Each phase runs in parallel, but each VM starts its phase offset by STAGGER seconds.
#
# Logs:
#   iiab_multipass_runs_YYYYMMDD/
#     <vm>.<action>.HHMMSS.log / .rc
#     latest.<vm>.<action>.log / .rc (symlinks)

set -euo pipefail

# Debian 13 (Trixie) official cloud image (qcow2). Multipass can launch from URL/file:// on Linux.
# Source: Debian cloud images live under cloud.debian.org/images/cloud/ ('genericcloud' includes cloud-init).
DEBIAN13_IMAGE_URL="${DEBIAN13_IMAGE_URL:-https://cloud.debian.org/images/cloud/trixie/latest/debian-13-genericcloud-amd64.qcow2}"

IMAGE="${IMAGE:-24.04}"
BASE="${BASE:-ubu2404}"
COUNT="${COUNT:-1}"
CPUS="${CPUS:-3}"
MEM="${MEM:-4G}"
DISK="${DISK:-20G}"

# PRs: prefer IIAB_PR, fall back to IIAB_INSTALL_ARGS (back-compat)
IIAB_PR="${IIAB_PR:-${IIAB_INSTALL_ARGS:-}}"

IIAB_FAST="${IIAB_FAST:-1}"
LOCAL_VARS_URL="${LOCAL_VARS_URL:-https://raw.githubusercontent.com/iiab/iiab/refs/heads/master/vars/local_vars_unittest.yml}"

WAIT_TRIES="${WAIT_TRIES:-60}"     # used ONLY for the first auto-resume
WAIT_SLEEP="${WAIT_SLEEP:-5}"
STAGGER="${STAGGER:-20}"

ACTION="run"
modules=()
prs=()
only_clean_vms=()

BOTH_DISTROS=0
DEBIAN13_ONLY=0

FIRST_UBUNTU=0
FIRST_DEBIAN=0

# Save original ubuntu settings so --both-distros can restore them even if --debian-13 was used earlier
UBU_IMAGE_ORIG="$IMAGE"
UBU_BASE_ORIG="$BASE"

usage() {
  cat <<EOF
Usage:
  $0 --run [--module code] [--module X...] [--pr N]...
  $0 --continue-after-reboot
  $0 --clean [--only-vm NAME]...
  $0 --debian-13 [--run|--clean|--continue-after-reboot] ...
  $0 --both-distros [--first-ubuntu|--first-debian|--clean] [...]

Aliases:
  --continue, --resume, --continue-after-upgrade  (same as --continue-after-reboot)
  --test-pr N                                     (same as --run-pr N)
Image shortcuts:
  --debian-13   Use Debian 13 cloud image; sets IMAGE=\$DEBIAN13_IMAGE_URL and BASE=deb13
  --both-distros  Run both Ubuntu + Debian 13 (COUNT=N => N + N VMs)
  --first-ubuntu  With --both-distros: run all Ubuntu first, then Debian
  --first-debian  With --both-distros: run all Debian first, then Ubuntu

PR options:
  --pr N        Add PR number (repeatable)
  --run-pr N    Add PR number and force --run

Env:
  IMAGE BASE COUNT CPUS MEM DISK IIAB_PR IIAB_FAST LOCAL_VARS_URL WAIT_TRIES WAIT_SLEEP STAGGER
EOF
}

# Parse args
while [[ $# -gt 0 ]]; do
  case "$1" in
    --run) ACTION="run"; shift ;;
    --continue-after-reboot|--continue-after-upgrade|--continue|--resume) ACTION="continue"; shift ;;
    --clean) ACTION="clean"; shift ;;

    --both-distros) BOTH_DISTROS=1; shift ;;
    --first-ubuntu) FIRST_UBUNTU=1; shift ;;
    --first-debian) FIRST_DEBIAN=1; shift ;;

    --debian-13)
      DEBIAN13_ONLY=1
      IMAGE="$DEBIAN13_IMAGE_URL"
      BASE="deb13"
      shift ;;

    --module)
      [[ $# -lt 2 ]] && { echo "[ERROR] --module needs a value"; exit 2; }
      modules+=("$2"); shift 2 ;;

    --pr)
      [[ $# -lt 2 ]] && { echo "[ERROR] --pr needs a number"; exit 2; }
      prs+=("$2"); shift 2 ;;

    --run-pr|--test-pr)
      [[ $# -lt 2 ]] && { echo "[ERROR] $1 needs a number"; exit 2; }
      ACTION="run"
      prs+=("$2"); shift 2 ;;

    --only-vm)
      [[ $# -lt 2 ]] && { echo "[ERROR] --only-vm needs a name"; exit 2; }
      only_clean_vms+=("$2"); shift 2 ;;

    -h|--help) usage; exit 0 ;;
    *) echo "[ERROR] Unknown option: $1"; usage; exit 2 ;;
  esac
done

# ---- Incoherency checks (fail fast) ----
if [[ "$FIRST_UBUNTU" == "1" && "$FIRST_DEBIAN" == "1" ]]; then
  echo "[ERROR] Incoherent options: --first-ubuntu and --first-debian cannot be used together."
  exit 2
fi

if [[ ( "$FIRST_UBUNTU" == "1" || "$FIRST_DEBIAN" == "1" ) && "$BOTH_DISTROS" != "1" ]]; then
  echo "[ERROR] --first-ubuntu/--first-debian requires --both-distros."
  exit 2
fi
# ---- 

# Default module
if [[ "${#modules[@]}" -eq 0 ]]; then
  modules=("code")
fi

# If no --pr provided, take from env IIAB_PR (space-separated)
if [[ "${#prs[@]}" -eq 0 && -n "${IIAB_PR:-}" ]]; then
  # shellcheck disable=SC2206
  prs=(${IIAB_PR})
fi

# Uniform VM names for run/continue: BASE-0..BASE-(COUNT-1)
# If both distros is enabled, restore Ubuntu image/base (in case --debian-13 was used too)
if [[ "$BOTH_DISTROS" == "1" ]]; then
  UBU_IMAGE="$UBU_IMAGE_ORIG"
  UBU_BASE="$UBU_BASE_ORIG"
  DEB_IMAGE="$DEBIAN13_IMAGE_URL"
  DEB_BASE="deb13"
else
  UBU_IMAGE="$IMAGE"
  UBU_BASE="$BASE"
  DEB_IMAGE="$DEBIAN13_IMAGE_URL"
  DEB_BASE="deb13"
fi

LOGROOT="${LOGROOT:-iiab_multipass_runs_$(date +%Y%m%d)}"
mkdir -p "$LOGROOT"

stamp() { date +%H%M%S; }

vm_exists() { multipass info "$1" >/dev/null 2>&1; }

wait_for_vm() {
  local vm="$1"
  local i
  for ((i=1; i<=WAIT_TRIES; i++)); do
    if multipass exec "$vm" -- bash -lc 'true' >/dev/null 2>&1; then
      return 0
    fi
    sleep "$WAIT_SLEEP"
  done
  return 1
}

set_latest_links() {
  local vm="$1" action="$2" log="$3" rc="$4"
  ln -sfn "$(basename "$log")" "$LOGROOT/latest.${vm}.${action}.log"
  ln -sfn "$(basename "$rc")"  "$LOGROOT/latest.${vm}.${action}.rc"
}

wait_all() {
  local rc=0
  local pid
  for pid in "$@"; do
    wait "$pid" || rc=1
  done
  return "$rc"
}

# Escape BASE for regex usage
re_escape() { printf '%s' "$1" | sed -e 's/[].[^$*+?(){}|\\]/\\&/g'; }

declare -A VM_IMAGE
names=()

build_vm_lists() {
  names=()
  VM_IMAGE=()

  if [[ "$BOTH_DISTROS" == "1" ]]; then
    if [[ "$FIRST_UBUNTU" == "1" ]]; then
      # all Ubuntu first, then all Debian
      for n in $(seq 0 $((COUNT-1))); do
        local u="${UBU_BASE}-${n}"
        names+=("$u")
        VM_IMAGE["$u"]="$UBU_IMAGE"
      done
      for n in $(seq 0 $((COUNT-1))); do
        local d="${DEB_BASE}-${n}"
        names+=("$d")
        VM_IMAGE["$d"]="$DEB_IMAGE"
      done

    elif [[ "$FIRST_DEBIAN" == "1" ]]; then
      # all Debian first, then all Ubuntu
      for n in $(seq 0 $((COUNT-1))); do
        local d="${DEB_BASE}-${n}"
        names+=("$d")
        VM_IMAGE["$d"]="$DEB_IMAGE"
      done
      for n in $(seq 0 $((COUNT-1))); do
        local u="${UBU_BASE}-${n}"
        names+=("$u")
        VM_IMAGE["$u"]="$UBU_IMAGE"
      done

    else
      # default interleaved
      for n in $(seq 0 $((COUNT-1))); do
        local u="${UBU_BASE}-${n}"
        local d="${DEB_BASE}-${n}"
        names+=("$u" "$d")
        VM_IMAGE["$u"]="$UBU_IMAGE"
        VM_IMAGE["$d"]="$DEB_IMAGE"
      done
    fi
  else
    for n in $(seq 0 $((COUNT-1))); do
      local vm="${BASE}-${n}"
      names+=("$vm")
      VM_IMAGE["$vm"]="$IMAGE"
    done
  fi
}

# Determine clean targets:
#  - If --only-vm is given: use those exact names
#  - Else: delete ALL VMs matching "^BASE-[0-9]+$" found in "multipass list"
clean_targets() {
  if [[ "${#only_clean_vms[@]}" -gt 0 ]]; then
    printf '%s\n' "${only_clean_vms[@]}"
    return 0
  fi

  local list
  list="$(multipass list 2>/dev/null | awk 'NR>1 {print $1}' || true)"

  if [[ "$BOTH_DISTROS" == "1" ]]; then
    local ubu_re deb_re
    ubu_re="$(re_escape "$UBU_BASE")"
    deb_re="$(re_escape "$DEB_BASE")"
    printf '%s\n' "$list" | grep -E "^(${ubu_re}|${deb_re})-[0-9]+$" || true
  else
    local base_re
    base_re="$(re_escape "$BASE")"
    printf '%s\n' "$list" | grep -E "^${base_re}-[0-9]+$" || true
  fi
}

cleanup_vms() {
  local targets=()
  while IFS= read -r line; do
    [[ -n "$line" ]] && targets+=("$line")
  done < <(clean_targets)

  if [[ "${#targets[@]}" -eq 0 ]]; then
    echo "[INFO] No VMs found to clean."
    echo "[INFO] Tip: use --only-vm NAME to force a specific VM."
    return 0
  fi

  echo "[INFO] Cleaning VMs: ${targets[*]}"

  echo "[INFO] Stopping VMs (best-effort)..."
  multipass stop "${targets[@]}" >/dev/null 2>&1 || true

  echo "[INFO] Deleting VMs (best-effort)..."
  multipass delete "${targets[@]}" >/dev/null 2>&1 || true

  echo "[INFO] Purging deleted VMs (best-effort)..."
  multipass purge >/dev/null 2>&1 || true

  echo "[INFO] Done."
}

launch_one() {
  local vm="$1"
  local img="${VM_IMAGE[$vm]:-}"
  [[ -z "$img" ]] && { echo "[ERROR] No image mapping for VM '$vm'"; return 2; }

  if vm_exists "$vm"; then
    echo "[INFO] VM already exists: $vm"
    return 0
  fi
  echo "[INFO] Launching $vm ..."
  multipass launch "$img" -n "$vm" -c "$CPUS" -m "$MEM" -d "$DISK" >/dev/null
}

run_install_txt() {
  local vm="$1"
  local t log rc
  t="$(stamp)"
  log="$LOGROOT/${vm}.install.${t}.log"
  rc="$LOGROOT/${vm}.install.${t}.rc"

  echo "[INFO] Logging files stored at: $LOGROOT"
  echo "[INFO] install.txt in $vm (log $(basename "$log")) ..."

  local modules_str pr_str
  modules_str="$(printf "%s " "${modules[@]}")"
  pr_str="$(printf "%s " "${prs[@]}")"

  set +e
  multipass exec "$vm" -- env \
    IIAB_FAST="$IIAB_FAST" \
    LOCAL_VARS_URL="$LOCAL_VARS_URL" \
    MODULES="$modules_str" \
    IIAB_PR_LIST="$pr_str" \
    bash -lc '
      set -euo pipefail
      export DEBIAN_FRONTEND=noninteractive

      sudo apt-get update
      sudo apt-get install -y curl python3 ca-certificates

      # 1) Seed UNITTEST local_vars (Size 0)
      sudo mkdir -p /etc/iiab
      echo "[INFO] Seeding /etc/iiab/local_vars.yml from: ${LOCAL_VARS_URL}"
      curl -fsSL "${LOCAL_VARS_URL}" | sudo tee /etc/iiab/local_vars.yml >/dev/null

      # 2) Set YAML key to True (edit if exists, else append)
      set_yaml_true() {
        local key="$1"
        if sudo grep -qE "^${key}:" /etc/iiab/local_vars.yml; then
          sudo sed -i -E "s/^${key}:.*/${key}: True/" /etc/iiab/local_vars.yml
        else
          echo "${key}: True" | sudo tee -a /etc/iiab/local_vars.yml >/dev/null
        fi
      }

      # 3) Enable module(s)
      for m in ${MODULES}; do
        set_yaml_true "${m}_install"
        set_yaml_true "${m}_enabled"
      done

      echo "--- local_vars.yml (module keys) ---"
      for m in ${MODULES}; do
        sudo grep -nE "^(${m}_install|${m}_enabled):" /etc/iiab/local_vars.yml || true
      done
      echo "--- end local_vars.yml ---"

      # 4) Run install.txt
      install_args=()
      if [[ "${IIAB_FAST:-1}" == "1" ]]; then
        install_args+=("-f")
      fi

      # IIAB_PR_LIST is space-separated PR numbers
      if [[ -n "${IIAB_PR_LIST:-}" ]]; then
        read -r -a extra <<<"${IIAB_PR_LIST}"
        install_args+=("${extra[@]}")
      fi

      echo "--- install.txt ---"
      echo "curl -fsSL https://iiab.io/install.txt | bash -s -- ${install_args[*]}"
      curl -fsSL https://iiab.io/install.txt | bash -s -- "${install_args[@]}"

      echo "--- install done ---"
    ' >"$log" 2>&1
  echo "$?" >"$rc"
  set -e

  set_latest_links "$vm" "install" "$log" "$rc"
}

resume_iiab() {
  local vm="$1"
  local do_long_wait="$2"   # 1 => wait_for_vm (only for first auto-resume), 0 => no long wait
  local t log rc
  t="$(stamp)"
  log="$LOGROOT/${vm}.resume.${t}.log"
  rc="$LOGROOT/${vm}.resume.${t}.rc"

  echo "[INFO] resume (iiab -f) in $vm (log $(basename "$log")) ..."

  set +e
  {
    multipass start "$vm" >/dev/null 2>&1 || true

    if [[ "$do_long_wait" == "1" ]]; then
      echo "[INFO] Waiting for VM readiness (first auto-resume): $vm"
      if ! wait_for_vm "$vm"; then
        echo "[ERROR] VM did not become ready in time: $vm"
        exit 88
      fi
    fi

    multipass exec "$vm" -- bash -lc '
      set -euo pipefail
      echo "--- resume: sudo iiab -f ---"
      if command -v iiab >/dev/null 2>&1; then
        sudo iiab -f
      else
        echo "[ERROR] iiab command not found; install likely not finished."
        exit 89
      fi
      echo "--- resume done ---"
    '
  } >"$log" 2>&1
  echo "$?" >"$rc"
  set -e

  set_latest_links "$vm" "resume" "$log" "$rc"
}

summary() {
  printf "\n================ SUMMARY (%s) ================\n" "$ACTION"
  printf "%-12s %-8s %-8s %s\n" "VM" "INSTALL" "RESUME" "Latest logs"
  printf "%-12s %-8s %-8s %s\n" "------------" "--------" "--------" "----------------------------------------------"

  for vm in "${names[@]}"; do
    local ir="n/a" rr="n/a"
    [[ -f "$LOGROOT/latest.${vm}.install.rc" ]] && ir="$(cat "$LOGROOT/latest.${vm}.install.rc" 2>/dev/null || echo n/a)"
    [[ -f "$LOGROOT/latest.${vm}.resume.rc"  ]] && rr="$(cat "$LOGROOT/latest.${vm}.resume.rc"  2>/dev/null || echo n/a)"

    printf "%-12s %-8s %-8s %s\n" \
      "$vm" "$ir" "$rr" \
      "latest.${vm}.install.log / latest.${vm}.resume.log"
  done

  echo
  echo "[INFO] Logs are in: $LOGROOT/"
  echo "[INFO] Clean:   $0 --clean"
  echo "[INFO] Resume:  $0 --continue-after-reboot"
}

phase_parallel_stagger() {
  # Runs a phase function in parallel, staggering start by STAGGER seconds.
  # Args: phase_name (launch|install|resume) [resume_wait=0|1]
  local phase="$1"
  local resume_wait="${2:-0}"
  local pids=()

  for i in "${!names[@]}"; do
    local vm="${names[$i]}"
    (
      sleep $((i * STAGGER))
      case "$phase" in
        launch) launch_one "$vm" ;;
        install) run_install_txt "$vm" ;;
        resume) resume_iiab "$vm" "$resume_wait" ;;
        *) echo "[ERROR] Unknown phase: $phase" >&2; exit 2 ;;
      esac
    ) &
    pids+=("$!")
  done

  # Don't abort the whole script if one VM fails; we still want logs + summary.
  set +e
  wait_all "${pids[@]}"
  set -e
}

# ---- Main ----
build_vm_lists

if [[ "$ACTION" == "clean" ]]; then
  cleanup_vms
  exit 0
fi

if [[ "${#prs[@]}" -gt 0 ]]; then
  echo "[INFO] PRs: ${prs[*]}"
fi

case "$ACTION" in
  run)
    pipeline_parallel_stagger
    summary
    ;;
  continue)
    phase_parallel_stagger resume 0
    summary
    ;;
esac