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

action=${1:-up}
script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
state_root=${XDG_STATE_HOME:-"$HOME/.local/state"}/nereus-lab
image_root=/var/lib/libvirt/images/nereus-lab
inventory="$state_root/inventory.yml"
key="$state_root/id_ed25519"
known_hosts="$state_root/known_hosts"
token_file="$state_root/k3s-token"
fedora_release=44
nodes=(nereus-node1 nereus-node2)
created_nodes=()

log() { printf '[nereus-lab] %s\n' "$*"; }
fail() { printf '[nereus-lab] ERROR: %s\n' "$*" >&2; exit 1; }

install_dependencies() {
  local missing=()
  for tool in qemu-img virsh virt-install virt-customize cloud-localds ansible-playbook curl ssh-keygen openssl ping; do
    command -v "$tool" >/dev/null 2>&1 || missing+=("$tool")
  done
  ((${#missing[@]} == 0)) && return
  [[ -r /etc/os-release ]] || fail "missing tools: ${missing[*]}"
  # shellcheck disable=SC1091
  . /etc/os-release
  log "installing virtualization, cloud-init, and Ansible tooling"
  case ${ID:-} in
    fedora)
      sudo dnf install -y @virtualization cloud-utils guestfs-tools ansible-core openssl iputils
      ;;
    arch|cachyos)
      sudo pacman -S --needed --noconfirm qemu-full libvirt virt-install cloud-image-utils guestfs-tools ansible-core openssl iputils dnsmasq
      ;;
    *)
      fail "automatic dependency installation does not support ${ID:-this host}; missing: ${missing[*]}"
      ;;
  esac
}

ensure_host() {
  local uplink
  install_dependencies
  [[ $(uname -m) == x86_64 ]] || fail "the bundled Fedora image workflow currently supports x86_64 hosts"
  [[ -e /dev/kvm ]] || fail "/dev/kvm is unavailable; enable CPU virtualization in firmware"
  sudo systemctl enable --now libvirtd
  if ! sudo virsh net-info default >/dev/null 2>&1; then
    fail "libvirt's default network is missing"
  fi
  sudo virsh net-start default >/dev/null 2>&1 || true
  sudo virsh net-autostart default >/dev/null
  if command -v ufw >/dev/null 2>&1 && sudo ufw status | grep -q '^Status: active'; then
    uplink=$(ip route show default | awk '/default/ {print $5; exit}')
    [[ -n "$uplink" ]] || fail "could not determine the host's default network interface"
    sudo ufw allow in on virbr0 to any port 53 proto udp >/dev/null
    sudo ufw allow in on virbr0 to any port 53 proto tcp >/dev/null
    sudo ufw allow in on virbr0 to any port 67 proto udp >/dev/null
    sudo ufw route allow in on virbr0 out on virbr0 from 192.168.122.0/24 to 192.168.122.0/24 >/dev/null
    sudo ufw route allow in on virbr0 out on "$uplink" from 192.168.122.0/24 >/dev/null
  fi
  mkdir -p "$state_root"
  chmod 0700 "$state_root"
  if [[ ! -f "$key" ]]; then
    ssh-keygen -q -t ed25519 -N '' -f "$key"
  fi
  touch "$known_hosts"
  chmod 0600 "$key" "$known_hosts"
  if [[ ! -f "$token_file" ]]; then
    umask 077
    openssl rand -hex 32 >"$token_file"
  fi
  chmod 0600 "$token_file"
  sudo install -d -m 0755 "$image_root"
}

image_url() {
  local arch=x86_64 listing filename
  listing="https://download.fedoraproject.org/pub/fedora/linux/releases/${fedora_release}/Cloud/${arch}/images/"
  filename=$(curl -fsSL "$listing" | sed -nE "s/.*href=\"(Fedora-Cloud-Base-Generic-${fedora_release}-[0-9.]+\.${arch}\.qcow2)\".*/\1/p" | sort -V | tail -n1)
  [[ -n "$filename" ]] || fail "could not discover the Fedora ${fedora_release} cloud image"
  printf '%s%s\n' "$listing" "$filename"
}

ensure_base_image() {
  local cached="$state_root/fedora-${fedora_release}-cloud-base.qcow2" source=${FEDORA_IMAGE:-}
  if [[ ! -s "$cached" ]]; then
    if [[ -n "$source" ]]; then
      [[ -f "$source" ]] || fail "FEDORA_IMAGE does not exist: $source"
      [[ $(basename "$source") == Fedora-Cloud-Base-Generic-* ]] || fail "FEDORA_IMAGE must be a Fedora Cloud Base Generic QCOW2 image"
      qemu-img check -q "$source" || fail "FEDORA_IMAGE is not a valid QCOW2 image"
      log "using local Fedora ${fedora_release} image: $source"
      install -m 0644 "$source" "$cached"
    else
      log "downloading Fedora ${fedora_release} cloud image"
      curl -fL --retry 3 --continue-at - -o "$cached.partial" "$(image_url)"
      if ! qemu-img check -q "$cached.partial"; then
        rm -f "$cached.partial"
        fail "downloaded Fedora image is not a valid QCOW2 image"
      fi
      mv "$cached.partial" "$cached"
    fi
  fi
  sudo install -m 0644 "$cached" "$image_root/fedora-${fedora_release}-cloud-base.qcow2"
}

node_mac() {
  case $1 in
    nereus-node1) printf '52:54:00:6e:01:01\n' ;;
    nereus-node2) printf '52:54:00:6e:01:02\n' ;;
    *) fail "no MAC address assigned for $1" ;;
  esac
}

create_seed() {
  local node=$1 seed="$state_root/${node}-seed.iso" user_data="$state_root/${node}-user-data"
  {
    printf '#cloud-config\nhostname: %s\nmanage_etc_hosts: true\nusers:\n' "$node"
    printf '  - name: fedora\n    groups: [wheel]\n    sudo: ALL=(ALL) NOPASSWD:ALL\n    shell: /bin/bash\n    ssh_authorized_keys:\n'
    printf '      - %s\n' "$(<"$key.pub")"
    printf 'ssh_pwauth: false\npackages: [qemu-guest-agent]\nruncmd:\n  - [systemctl, enable, --now, qemu-guest-agent]\n'
  } >"$user_data"
  chmod 0600 "$user_data"
  cloud-localds --dsmode local --hostname "$node" "$seed" "$user_data"
  sudo install -m 0644 "$seed" "$image_root/${node}-seed.iso"
}

install_network_service() {
  local node=$1 disk=$2
  sudo virt-customize -q -a "$disk" \
    --copy-in "$script_dir/nereus-network.service:/etc/systemd/system" \
    --copy-in "$script_dir/${node}.nmconnection:/etc/NetworkManager/system-connections" \
    --chmod "0600:/etc/NetworkManager/system-connections/${node}.nmconnection" \
    --run-command "chown root:root /etc/NetworkManager/system-connections/${node}.nmconnection" \
    --run-command 'rm -f /etc/NetworkManager/system-connections/cloud-init-ens2.nmconnection' \
    --run-command 'systemctl enable nereus-network.service'
}

create_vm() {
  local node=$1 disk="$image_root/${node}.qcow2" mac
  mac=$(node_mac "$node")
  if sudo virsh dominfo "$node" >/dev/null 2>&1; then
    if [[ $(sudo virsh domstate "$node") == "shut off" ]]; then
      install_network_service "$node" "$disk"
    fi
    sudo virsh start "$node" >/dev/null 2>&1 || true
    return
  fi
  log "creating $node"
  created_nodes+=("$node")
  ssh-keygen -q -f "$known_hosts" -R "$(node_ip "$node")" >/dev/null 2>&1 || true
  create_seed "$node"
  sudo qemu-img create -q -f qcow2 -F qcow2 -b "$image_root/fedora-${fedora_release}-cloud-base.qcow2" "$disk" 30G
  install_network_service "$node" "$disk"
  sudo virt-install --name "$node" --memory 3072 --vcpus 2 --import \
    --disk "path=$disk,format=qcow2,bus=virtio" \
    --disk "path=$image_root/${node}-seed.iso,format=raw,bus=virtio,readonly=on" \
    --network "network=default,model=virtio,mac=$mac" --graphics none --noautoconsole \
    --boot uefi --osinfo detect=on,require=off
}

node_ip() {
  case $1 in
    nereus-node1) printf '192.168.122.10\n' ;;
    nereus-node2) printf '192.168.122.11\n' ;;
    *) fail "no address assigned for $1" ;;
  esac
}

wait_ssh() {
  local ip=$1
  for _ in $(seq 1 90); do
    if ssh -i "$key" -o BatchMode=yes -o ConnectTimeout=2 -o StrictHostKeyChecking=accept-new -o UserKnownHostsFile="$known_hosts" "fedora@$ip" true 2>/dev/null; then return; fi
    sleep 2
  done
  fail "SSH did not become ready at $ip"
}

wait_ping() {
  local node=$1 ip=$2
  for _ in $(seq 1 60); do
    ping -c 1 -W 1 "$ip" >/dev/null 2>&1 && return
    sleep 2
  done
  fail "$node does not answer ping at $ip"
}

write_inventory() {
  local ip1=$1 ip2=$2
  umask 077
  printf 'all:\n  vars:\n    ansible_user: fedora\n    ansible_ssh_private_key_file: %s\n    ansible_ssh_common_args: "-o StrictHostKeyChecking=accept-new -o UserKnownHostsFile=%s"\n    k3s_version: v1.33.4+k3s1\n    k3s_cluster_cidr: 10.42.0.0/16\n    k3s_service_cidr: 10.43.0.0/16\n    k3s_node_cidr: 192.168.122.0/24\n    k3s_operator_cidrs:\n      - 192.168.122.1/32\n  children:\n    k3s_server:\n      hosts:\n        node1:\n          ansible_host: %s\n          k3s_node_ip: %s\n    k3s_agent:\n      hosts:\n        node2:\n          ansible_host: %s\n          k3s_node_ip: %s\n' "$key" "$known_hosts" "$ip1" "$ip1" "$ip2" "$ip2" >"$inventory"
}

check_lab() {
  [[ -f "$inventory" ]] || fail "run '$0 up' first"
  log "checking reachability, services, ports, packages, and directories"
  ansible all -i "$inventory" -m ansible.builtin.ping
  ansible all -i "$inventory" -b -m ansible.builtin.shell -a 'for i in $(seq 1 60); do test -d /var/lib/rancher/k3s && command -v k3s >/dev/null && systemctl is-active --quiet firewalld && ss -lnt | grep -q ":10250 " && exit 0; sleep 2; done; exit 1'
  ansible k3s_server -i "$inventory" -b -m ansible.builtin.shell -a 'for i in $(seq 1 60); do systemctl is-active --quiet k3s && ss -lnt | grep -q ":6443 " && exit 0; sleep 2; done; exit 1'
  ansible k3s_agent -i "$inventory" -b -m ansible.builtin.shell -a 'for i in $(seq 1 60); do systemctl is-active --quiet k3s-agent && exit 0; sleep 2; done; exit 1'
  ansible k3s_server -i "$inventory" -b -m ansible.builtin.command -a 'k3s kubectl get nodes -o wide'
  ansible k3s_server -i "$inventory" -b -m ansible.builtin.command -a 'k3s kubectl wait --for=condition=Ready nodes --all --timeout=60s'
  log "all checks passed"
}

case "$action" in
  up)
    ensure_host
    ensure_base_image
    for node in "${nodes[@]}"; do create_vm "$node"; done
    ip1=$(node_ip "${nodes[0]}"); ip2=$(node_ip "${nodes[1]}")
    log "node1=$ip1 node2=$ip2"
    wait_ping node1 "$ip1"
    wait_ping node2 "$ip2"
    wait_ssh "$ip1"; wait_ssh "$ip2"
    for node in "${created_nodes[@]}"; do
      if [[ "$node" != "${nodes[0]}" ]]; then
        ssh -i "$key" -o BatchMode=yes -o StrictHostKeyChecking=accept-new -o UserKnownHostsFile="$known_hosts" \
          "fedora@$ip1" "sudo -n k3s kubectl delete node '$node' --ignore-not-found=true"
      fi
    done
    write_inventory "$ip1" "$ip2"
    K3S_TOKEN=$(<"$token_file") "$script_dir/bootstrap.sh" "$inventory"
    check_lab
    ;;
  check) check_lab ;;
  stop) for node in "${nodes[@]}"; do sudo virsh shutdown "$node" >/dev/null 2>&1 || true; done ;;
  destroy)
    for node in "${nodes[@]}"; do sudo virsh destroy "$node" >/dev/null 2>&1 || true; sudo virsh undefine "$node" --remove-all-storage; done
    ;;
  *) fail "usage: $0 [up|check|stop|destroy]" ;;
esac