#!/bin/bash
# Written by Paul Clevett
# (C)Copyright Wolf Software Systems Ltd
# https://wolf.uk.com
#
#
# WolfStack Quick Install Script
# Installs WolfStack server management dashboard
# Supported: Ubuntu/Debian, Fedora/RHEL/CentOS, SLES/openSUSE, Arch Linux, IBM Power (ppc64le)
#
# Usage: curl -sSL https://raw.githubusercontent.com/wolfsoftwaresystemsltd/WolfStack/master/setup.sh | sudo bash
#        curl -sSL https://raw.githubusercontent.com/wolfsoftwaresystemsltd/WolfStack/beta/setup.sh | sudo bash -s -- --beta
#        sudo bash setup.sh --install-dir /mnt/usb      # build & install from external drive
#

set -e

# Helper: read from /dev/tty if available, otherwise return empty (use defaults)
prompt_read() {
    if [ -e /dev/tty ] && : < /dev/tty 2>/dev/null; then
        read "$1" < /dev/tty 2>/dev/null || eval "$1="
    else
        eval "$1="
    fi
}

# ─── Parse arguments ─────────────────────────────────────────────────────────
BRANCH="master"
CUSTOM_INSTALL_DIR=""
while [ $# -gt 0 ]; do
    case "$1" in
        --beta) BRANCH="beta" ;;
        --install-dir|--install)
            if [ -n "$2" ]; then
                shift
                CUSTOM_INSTALL_DIR="$1"
            else
                echo "✗ --install-dir requires a path argument"
                exit 1
            fi
            ;;
    esac
    shift
done

# Allow git to operate on repos owned by other users (setup.sh runs as root
# but repos may have been cloned by a regular user)
export GIT_CONFIG_COUNT=1
export GIT_CONFIG_KEY_0=safe.directory
export GIT_CONFIG_VALUE_0="*"

# ─── Architecture detection for prebuilt binaries ──────────────────────────
HOST_ARCH=$(uname -m)
case "$HOST_ARCH" in
    x86_64)  BINARY_ARCH="x86_64" ;;
    aarch64) BINARY_ARCH="aarch64" ;;
    *)       BINARY_ARCH="" ;;  # unsupported — will build from source
esac

# Download a prebuilt binary from GitHub Releases.
# Usage: download_prebuilt <repo> <binary_name> <dest_path>
# Returns 0 on success, 1 on failure (caller should fall back to source build)
download_prebuilt() {
    local repo="$1" binary="$2" dest="$3"
    if [ -z "$BINARY_ARCH" ]; then
        return 1
    fi
    local url="https://github.com/${repo}/releases/latest/download/${binary}-${BINARY_ARCH}"
    echo "  Downloading prebuilt ${binary} for ${BINARY_ARCH}..."
    local tmpfile="${dest}.download"
    if curl -fSL --connect-timeout 15 --max-time 300 --retry 2 -o "$tmpfile" "$url" 2>&1; then
        mv "$tmpfile" "$dest"
        chmod +x "$dest"
        echo "  ✓ Downloaded prebuilt ${binary} (${BINARY_ARCH})"
        return 0
    else
        echo "  ⚠ Prebuilt binary not available — will build from source"
        rm -f "$tmpfile"
        return 1
    fi
}

# ─── Custom install directory (for low-disk devices like Raspberry Pi) ───────
if [ -n "$CUSTOM_INSTALL_DIR" ]; then
    # If given a block device, mount it
    if [ -b "$CUSTOM_INSTALL_DIR" ]; then
        MOUNT_DEV="$CUSTOM_INSTALL_DIR"
        CUSTOM_INSTALL_DIR="/mnt/wolfstack-build"
        mkdir -p "$CUSTOM_INSTALL_DIR"
        if ! mountpoint -q "$CUSTOM_INSTALL_DIR" 2>/dev/null; then
            echo "Mounting $MOUNT_DEV at $CUSTOM_INSTALL_DIR..."
            mount "$MOUNT_DEV" "$CUSTOM_INSTALL_DIR"
        fi
    fi
    mkdir -p "$CUSTOM_INSTALL_DIR"

    # Redirect EVERYTHING to external drive: Rust toolchain, build cache, temp files
    export RUSTUP_HOME="$CUSTOM_INSTALL_DIR/.rustup"
    export CARGO_HOME="$CUSTOM_INSTALL_DIR/.cargo"
    export TMPDIR="$CUSTOM_INSTALL_DIR/tmp"
    export PATH="$CARGO_HOME/bin:$PATH"
    mkdir -p "$TMPDIR"
fi

echo ""
echo "  🐺 WolfStack Installer"
echo "  ─────────────────────────────────────"
echo "  Server Management Platform"
if [ "$BRANCH" != "master" ]; then
    echo "  Branch: $BRANCH"
fi
if [ -n "$CUSTOM_INSTALL_DIR" ]; then
    echo "  Install dir: $CUSTOM_INSTALL_DIR"
fi
echo ""

# ─── Must run as root ────────────────────────────────────────────────────────
if [ "$(id -u)" -ne 0 ]; then
    echo "✗ This script must be run as root."
    echo "  Usage: sudo bash setup.sh"
    echo "     or: curl -sSL <url> | sudo bash"
    exit 1
fi

# Detect the real user (for Rust install) when running under sudo
REAL_USER="${SUDO_USER:-root}"
REAL_HOME=$(eval echo "~$REAL_USER")

# ─── Detect package manager ─────────────────────────────────────────────────
echo "Checking system requirements..."

if command -v apt >/dev/null 2>&1; then
    PKG_MANAGER="apt"
    echo "✓ Detected Debian/Ubuntu (apt)"
elif command -v dnf >/dev/null 2>&1; then
    PKG_MANAGER="dnf"
    echo "✓ Detected Fedora/RHEL (dnf)"
elif command -v yum >/dev/null 2>&1; then
    PKG_MANAGER="yum"
    echo "✓ Detected RHEL/CentOS (yum)"
elif command -v zypper >/dev/null 2>&1; then
    PKG_MANAGER="zypper"
    echo "✓ Detected SLES/openSUSE (zypper)"
elif command -v pacman >/dev/null 2>&1; then
    PKG_MANAGER="pacman"
    echo "✓ Detected Arch Linux (pacman)"
else
    echo "✗ Could not detect package manager (apt/dnf/yum/zypper/pacman)"
    echo "  Please install dependencies manually."
    exit 1
fi

# ─── Update system packages first ─────────────────────────────────────────
# Ensures package index is in sync and avoids dependency mismatches
# Refresh package index (needed to install dependencies) but do NOT upgrade existing packages.
# A full system upgrade can break things, takes ages, and the user didn't ask for it.
echo ""
echo "Refreshing package index..."
if [ "$PKG_MANAGER" = "apt" ]; then
    if ! apt update -qq 2>/dev/null; then
        echo "  ⚠ Some repositories failed to update."
        echo "    This is usually caused by a third-party repo (e.g. Docker) that doesn't"
        echo "    support your distro version. WolfStack will still install, but you may"
        echo "    need to fix the broken repo afterwards:"
        echo "      sudo apt update    (to see which repo is failing)"
        echo "      Check /etc/apt/sources.list.d/ for the problematic .list file"
        echo "    Continuing installation..."
    fi
elif [ "$PKG_MANAGER" = "dnf" ]; then
    dnf makecache -q 2>/dev/null || true
elif [ "$PKG_MANAGER" = "zypper" ]; then
    zypper refresh -q 2>/dev/null || true
elif [ "$PKG_MANAGER" = "pacman" ]; then
    pacman -Sy --noconfirm 2>/dev/null || true
fi
echo "✓ Package index refreshed"

# ─── Detect Proxmox VE host ─────────────────────────────────────────────────
IS_PROXMOX=false
if command -v pveversion >/dev/null 2>&1 || [ -f /etc/pve/.version ] || dpkg -l proxmox-ve >/dev/null 2>&1 2>&1; then
    IS_PROXMOX=true
    PVE_VER=$(pveversion 2>/dev/null || echo "unknown")
    echo "✓ Detected Proxmox VE host ($PVE_VER)"
    echo "  Skipping packages already provided by Proxmox (QEMU, LXC)"
fi

# ─── Install system dependencies ────────────────────────────────────────────
echo ""
echo "Installing system dependencies..."

if [ "$PKG_MANAGER" = "apt" ]; then
    apt update -qq 2>/dev/null || true
    # On Proxmox hosts, QEMU and LXC are already provided by pve-qemu-kvm and lxc-pve.
    # Many Debian packages conflict with PVE equivalents, causing APT to try removing
    # the proxmox-ve metapackage. We must be very conservative on PVE hosts.
    if [ "$IS_PROXMOX" = true ]; then
        # Only install build dependencies needed for compiling Rust/WolfStack.
        # Proxmox already provides QEMU, LXC, socat, bridge-utils, etc.
        apt install -y --no-install-recommends git curl build-essential pkg-config libssl-dev libcrypt-dev || {
            echo "⚠ Some build dependencies failed to install. Trying individually..."
            for pkg in git curl build-essential pkg-config libssl-dev libcrypt-dev; do
                dpkg -s "$pkg" >/dev/null 2>&1 || apt install -y --no-install-recommends "$pkg" 2>/dev/null || true
            done
        }
        # Install optional runtime deps one-by-one — skip if already provided by PVE
        for pkg in dnsmasq-base bridge-utils socat nfs-common fuse3; do
            if dpkg -s "$pkg" >/dev/null 2>&1; then
                echo "  ✓ $pkg already installed"
            else
                echo "  Installing $pkg..."
                apt install -y --no-install-recommends "$pkg" 2>/dev/null || \
                    echo "  ⚠ Could not install $pkg (may conflict with PVE) — skipping"
            fi
        done
        # s3fs — try both package names (s3fs-fuse on Kali/some Debian, s3fs on Ubuntu/Proxmox)
        if ! dpkg -s s3fs-fuse >/dev/null 2>&1 && ! dpkg -s s3fs >/dev/null 2>&1; then
            apt install -y --no-install-recommends s3fs-fuse 2>/dev/null || \
                apt install -y --no-install-recommends s3fs 2>/dev/null || \
                echo "  ⚠ s3fs not available — S3 mounts will use built-in sync"
        fi
    else
        # Select architecture-appropriate QEMU package
        ARCH=$(uname -m)
        if [ "$ARCH" = "ppc64le" ] || [ "$ARCH" = "ppc64" ]; then
            QEMU_PKG="qemu-system-ppc qemu-utils"
        elif [ "$ARCH" = "aarch64" ]; then
            QEMU_PKG="qemu-system-arm qemu-utils qemu-efi-aarch64"
        else
            QEMU_PKG="qemu-system-x86 qemu-utils"
        fi
        apt install -y git curl build-essential pkg-config libssl-dev libcrypt-dev lxc lxc-templates dnsmasq-base bridge-utils $QEMU_PKG socat nfs-common fuse3
        apt install -y s3fs-fuse 2>/dev/null || apt install -y s3fs 2>/dev/null || echo "  ⚠ s3fs not available — S3 mounts will use built-in sync"
    fi
elif [ "$PKG_MANAGER" = "dnf" ]; then
    ARCH=$(uname -m)
    if [ "$ARCH" = "aarch64" ]; then
        QEMU_DNF="qemu-system-aarch64 qemu-img edk2-aarch64"
    else
        QEMU_DNF="qemu-kvm qemu-img"
    fi
    dnf install -y git curl gcc gcc-c++ make openssl-devel pkg-config libxcrypt-devel lxc lxc-templates lxc-extra dnsmasq bridge-utils $QEMU_DNF socat s3fs-fuse nfs-utils fuse3
elif [ "$PKG_MANAGER" = "yum" ]; then
    ARCH=$(uname -m)
    if [ "$ARCH" = "aarch64" ]; then
        QEMU_YUM="qemu-system-aarch64 qemu-img"
    else
        QEMU_YUM="qemu-kvm qemu-img"
    fi
    yum install -y git curl gcc gcc-c++ make openssl-devel pkgconfig lxc lxc-templates lxc-extra dnsmasq bridge-utils $QEMU_YUM socat s3fs-fuse nfs-utils fuse
elif [ "$PKG_MANAGER" = "zypper" ]; then
    ARCH=$(uname -m)
    if [ "$ARCH" = "aarch64" ]; then
        QEMU_ZYPP="qemu-arm qemu-tools qemu-uefi-aarch64"
    else
        QEMU_ZYPP="qemu-kvm qemu-tools"
    fi
    zypper install -y git curl gcc gcc-c++ make libopenssl-devel pkg-config lxc dnsmasq bridge-utils $QEMU_ZYPP socat s3fs nfs-client fuse3
elif [ "$PKG_MANAGER" = "pacman" ]; then
    ARCH=$(uname -m)
    if [ "$ARCH" = "aarch64" ]; then
        QEMU_PAC="qemu-system-aarch64 qemu-img edk2-aarch64"
    else
        QEMU_PAC="qemu-full"
    fi
    pacman -Sy --noconfirm --needed git curl base-devel openssl pkg-config lxc dnsmasq $QEMU_PAC socat s3fs-fuse nfs-utils fuse3 rustup
fi

echo "✓ System dependencies installed"

# ─── Install Proxmox Backup Client (optional, for PBS integration) ──────────
echo ""
echo "Installing Proxmox Backup Client..."

# Helper: detect the best PBS codename to use based on current OS
pbs_detect_codename() {
    local codename=""
    if [ -r /etc/os-release ]; then
        . /etc/os-release
        codename="${VERSION_CODENAME:-}"
    fi
    [ -z "$codename" ] && codename=$(lsb_release -sc 2>/dev/null || echo "")
    case "$codename" in
        # Debian — use as-is (proxmox publishes for these)
        trixie|bookworm|bullseye) echo "$codename" ;;
        # Ubuntu codenames mapped to closest Debian
        noble|oracular|plucky) echo "trixie" ;;       # 24.04+/25.04 → Debian 13
        jammy|lunar|mantic)    echo "bookworm" ;;     # 22.04–23.10 → Debian 12
        focal|impish)          echo "bullseye" ;;     # 20.04–21.10 → Debian 11
        # Unknown/everything else — trixie is newest; extraction will fallback
        *) echo "trixie" ;;
    esac
}

# Helper: extract proxmox-backup-client binary from Debian .deb
# Used on non-Debian systems (Fedora, Arch fallback, openSUSE)
pbs_extract_from_deb() {
    local codename="$1"
    local arch="${2:-amd64}"
    local tmp
    tmp=$(mktemp -d)
    local base_url="http://download.proxmox.com/debian/pbs/dists/${codename}/pbs-no-subscription/binary-${arch}/"
    local deb_name
    deb_name=$(curl -fsSL "$base_url" 2>/dev/null | grep -oP 'proxmox-backup-client_[^"]+\.deb' | grep -v dbgsym | sort -V | tail -1)
    if [ -z "$deb_name" ]; then
        rm -rf "$tmp"
        return 1
    fi
    echo "  Downloading $deb_name (${codename}/${arch})..." >&2
    if ! curl -fsSL "${base_url}${deb_name}" -o "${tmp}/${deb_name}" 2>/dev/null; then
        rm -rf "$tmp"; return 1
    fi
    ( cd "$tmp" && ar x "$deb_name" 2>/dev/null ) || { rm -rf "$tmp"; return 1; }
    local data_tar
    data_tar=$(ls "$tmp"/data.tar.* 2>/dev/null | head -1)
    [ -z "$data_tar" ] && { rm -rf "$tmp"; return 1; }
    case "$data_tar" in
        *.zst) zstd -d -q "$data_tar" -o "$tmp/data.tar" 2>/dev/null || { rm -rf "$tmp"; return 1; } ;;
        *.xz)  xz -d -k "$data_tar" 2>/dev/null || { rm -rf "$tmp"; return 1; } ;;
        *.gz)  gzip -dk "$data_tar" 2>/dev/null || { rm -rf "$tmp"; return 1; } ;;
    esac
    if ! tar -C "$tmp" -xf "$tmp/data.tar" ./usr/bin/proxmox-backup-client 2>/dev/null; then
        rm -rf "$tmp"; return 1
    fi
    install -m 0755 "$tmp/usr/bin/proxmox-backup-client" /usr/local/bin/proxmox-backup-client
    rm -rf "$tmp"
    return 0
}

pbs_install_success=false

if command -v proxmox-backup-client >/dev/null 2>&1; then
    echo "✓ proxmox-backup-client already installed ($(proxmox-backup-client --version 2>&1 | head -1))"
    pbs_install_success=true

elif command -v apt-get >/dev/null 2>&1; then
    # ─── Debian / Ubuntu / Proxmox VE ──────────────────────────────────────
    CODENAME=$(pbs_detect_codename)
    echo "  Using Proxmox PBS repo for: $CODENAME"
    mkdir -p /etc/apt/sources.list.d /etc/apt/trusted.gpg.d
    echo "deb http://download.proxmox.com/debian/pbs $CODENAME pbs-no-subscription" > /etc/apt/sources.list.d/pbs-client.list
    curl -fsSL "https://enterprise.proxmox.com/debian/proxmox-release-${CODENAME}.gpg" \
        -o "/etc/apt/trusted.gpg.d/proxmox-release-${CODENAME}.gpg" 2>/dev/null || true

    if apt-get update -qq 2>/dev/null && apt-get install -y proxmox-backup-client 2>/dev/null; then
        echo "✓ proxmox-backup-client installed from $CODENAME repo"
        pbs_install_success=true
    elif [ "$CODENAME" != "bookworm" ]; then
        echo "  ⚠ $CODENAME install failed — trying bookworm repo"
        echo "deb http://download.proxmox.com/debian/pbs bookworm pbs-no-subscription" > /etc/apt/sources.list.d/pbs-client.list
        curl -fsSL "https://enterprise.proxmox.com/debian/proxmox-release-bookworm.gpg" \
            -o "/etc/apt/trusted.gpg.d/proxmox-release-bookworm.gpg" 2>/dev/null || true
        if apt-get update -qq 2>/dev/null && apt-get install -y proxmox-backup-client 2>/dev/null; then
            echo "✓ proxmox-backup-client installed from bookworm repo"
            pbs_install_success=true
        fi
    fi

elif command -v pacman >/dev/null 2>&1; then
    # ─── Arch / CachyOS / Manjaro ──────────────────────────────────────────
    # Required libraries for the binary: libfuse3, openssl 3, acl, zstd
    pacman -S --needed --noconfirm fuse3 openssl acl zstd 2>/dev/null || true

    # Try AUR helper first (paru or yay) — builds clean Arch package
    AUR_HELPER=""
    for h in paru yay; do
        if command -v "$h" >/dev/null 2>&1; then AUR_HELPER="$h"; break; fi
    done
    if [ -n "$AUR_HELPER" ] && [ -n "$SUDO_USER" ] && [ "$SUDO_USER" != "root" ]; then
        echo "  Using $AUR_HELPER to build proxmox-backup-client-bin from AUR..."
        if su - "$SUDO_USER" -c "$AUR_HELPER -S --needed --noconfirm proxmox-backup-client-bin" 2>/dev/null; then
            pbs_install_success=true
        fi
    fi
    if [ "$pbs_install_success" != "true" ]; then
        echo "  Falling back to Debian .deb extraction..."
        if pbs_extract_from_deb trixie amd64 || pbs_extract_from_deb bookworm amd64; then
            echo "✓ proxmox-backup-client installed to /usr/local/bin/"
            pbs_install_success=true
        fi
    fi

elif command -v dnf >/dev/null 2>&1; then
    # ─── Fedora / RHEL / Rocky / AlmaLinux ─────────────────────────────────
    dnf install -y fuse3 openssl libacl zstd 2>/dev/null || true
    ARCH=$(uname -m | sed 's/x86_64/amd64/;s/aarch64/arm64/')
    if pbs_extract_from_deb trixie "$ARCH" || pbs_extract_from_deb bookworm "$ARCH"; then
        echo "✓ proxmox-backup-client installed to /usr/local/bin/"
        pbs_install_success=true
    fi

elif command -v zypper >/dev/null 2>&1; then
    # ─── openSUSE ──────────────────────────────────────────────────────────
    zypper install -y fuse3 openssl libacl1 libzstd1 2>/dev/null || true
    ARCH=$(uname -m | sed 's/x86_64/amd64/;s/aarch64/arm64/')
    if pbs_extract_from_deb trixie "$ARCH" || pbs_extract_from_deb bookworm "$ARCH"; then
        echo "✓ proxmox-backup-client installed to /usr/local/bin/"
        pbs_install_success=true
    fi

else
    # ─── Unknown distro — try generic .deb extract ─────────────────────────
    ARCH=$(uname -m | sed 's/x86_64/amd64/;s/aarch64/arm64/')
    if pbs_extract_from_deb trixie "$ARCH" || pbs_extract_from_deb bookworm "$ARCH"; then
        echo "✓ proxmox-backup-client installed to /usr/local/bin/"
        pbs_install_success=true
    fi
fi

if [ "$pbs_install_success" != "true" ]; then
    echo "⚠ Could not install proxmox-backup-client. PBS integration will be unavailable."
    echo "  You can install manually later. See: https://pbs.proxmox.com/docs/backup-client.html"
fi

# Fix libfuse3 soname: proxmox-backup-client links against libfuse3.so.3
# but some distros (CachyOS, rolling Arch) have soname 4 (libfuse3.so.4)
if command -v proxmox-backup-client >/dev/null 2>&1; then
    for libdir in /usr/lib /usr/lib64 /usr/lib/x86_64-linux-gnu /usr/lib/aarch64-linux-gnu; do
        if [ ! -e "$libdir/libfuse3.so.3" ] && [ -e "$libdir/libfuse3.so.4" ]; then
            FUSE3_REAL=$(readlink -f "$libdir/libfuse3.so.4")
            ln -sf "$FUSE3_REAL" "$libdir/libfuse3.so.3"
            echo "  ✓ Created $libdir/libfuse3.so.3 symlink (soname compat)"
        fi
    done
fi

# ─── Configure FUSE for storage mounts ──────────────────────────────────────
# Enable allow_other in FUSE (needed for s3fs mounts accessible by containers)
if [ -f /etc/fuse.conf ]; then
    if ! grep -q "^user_allow_other" /etc/fuse.conf; then
        echo "user_allow_other" >> /etc/fuse.conf
    fi
fi

# Create storage directories
# rust-s3 syncs bucket contents to /var/cache/wolfstack/s3/<mount-id>/
mkdir -p /etc/wolfstack/s3 /etc/wolfstack/pbs /mnt/wolfstack /var/cache/wolfstack/s3
echo "✓ Storage directories configured"

# Lock down /etc/wolfstack — it holds the cluster secret, PVE API
# tokens inside nodes.json, the join-token, and license.key. Before
# v18.7.27 these files were world-readable (inherited process umask),
# which let any unprivileged local user impersonate a cluster member
# or siphon PVE credentials. The running binary also enforces this
# on startup (paths::harden_existing); tightening here too closes
# the window on very first install before wolfstack has started.
chmod 700 /etc/wolfstack 2>/dev/null || true
for f in /etc/wolfstack/custom-cluster-secret \
         /etc/wolfstack/cluster-secret \
         /etc/wolfstack/nodes.json \
         /etc/wolfstack/join-token \
         /etc/wolfstack/license.key \
         /etc/wolfstack/key.pem; do
    if [ -e "$f" ]; then
        chmod 600 "$f" 2>/dev/null || true
    fi
done

# ─── Install Docker if missing ──────────────────────────────────────────────
if ! command -v docker >/dev/null 2>&1; then
    echo ""
    echo "Installing Docker..."
    DOCKER_INSTALLED=false

    # Try the official convenience script first (works for Ubuntu, Debian, Fedora, CentOS, RHEL)
    if curl -fsSL https://get.docker.com | sh 2>/dev/null; then
        DOCKER_INSTALLED=true
    else
        # Convenience script failed — likely a derivative distro (Nobara, Rocky, Alma, etc.)
        # Detect the distro family from /etc/os-release and set up Docker repo manually
        echo "  Convenience script failed — trying manual Docker repo setup..."
        DISTRO_ID=$(. /etc/os-release 2>/dev/null && echo "$ID")
        DISTRO_LIKE=$(. /etc/os-release 2>/dev/null && echo "$ID_LIKE")
        DISTRO_VERSION=$(. /etc/os-release 2>/dev/null && echo "$VERSION_ID")

        if echo "$DISTRO_ID $DISTRO_LIKE" | grep -qiE "fedora"; then
            # Fedora family (Nobara, Ultramarine, etc.) — use Fedora Docker repo
            # Use the major Fedora version the derivative is based on
            FEDORA_VER="$DISTRO_VERSION"
            # For derivatives, try to detect the Fedora base version
            if [ "$DISTRO_ID" != "fedora" ]; then
                PLATFORM_ID=$(. /etc/os-release 2>/dev/null && echo "$PLATFORM_ID")
                if echo "$PLATFORM_ID" | grep -q "fedora"; then
                    FEDORA_VER=$(echo "$PLATFORM_ID" | grep -oP 'f\K[0-9]+' || echo "$DISTRO_VERSION")
                fi
            fi
            echo "  Detected Fedora family (${DISTRO_ID} based on Fedora ${FEDORA_VER})"
            dnf -y install dnf-plugins-core 2>/dev/null || true
            dnf config-manager addrepo --from-repofile=https://download.docker.com/linux/fedora/docker-ce.repo 2>/dev/null || \
                dnf config-manager --add-repo https://download.docker.com/linux/fedora/docker-ce.repo 2>/dev/null || true
            # Override $releasever with actual Fedora version for derivatives
            if [ "$DISTRO_ID" != "fedora" ] && [ -n "$FEDORA_VER" ]; then
                sed -i "s/\$releasever/${FEDORA_VER}/g" /etc/yum.repos.d/docker-ce.repo 2>/dev/null || true
            fi
            if dnf install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin; then
                DOCKER_INSTALLED=true
            fi
        elif echo "$DISTRO_ID $DISTRO_LIKE" | grep -qiE "rhel|centos"; then
            # RHEL family (Rocky, Alma, Oracle, etc.) — use CentOS Docker repo
            echo "  Detected RHEL/CentOS family (${DISTRO_ID})"
            yum install -y yum-utils 2>/dev/null || true
            yum-config-manager --add-repo https://download.docker.com/linux/centos/docker-ce.repo 2>/dev/null || \
                dnf config-manager addrepo --from-repofile=https://download.docker.com/linux/centos/docker-ce.repo 2>/dev/null || true
            if yum install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin 2>/dev/null || \
               dnf install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin 2>/dev/null; then
                DOCKER_INSTALLED=true
            fi
        elif echo "$DISTRO_ID $DISTRO_LIKE" | grep -qiE "suse|sles"; then
            # openSUSE/SLES — use distro docker packages
            echo "  Detected SUSE family (${DISTRO_ID})"
            if zypper install -y docker docker-compose 2>/dev/null; then
                DOCKER_INSTALLED=true
            fi
        elif echo "$DISTRO_ID $DISTRO_LIKE" | grep -qiE "debian|ubuntu"; then
            # Debian derivatives (Mint, Pop!_OS, etc.) — use Debian/Ubuntu Docker repo
            UPSTREAM="debian"
            CODENAME=$(. /etc/os-release 2>/dev/null && echo "$UBUNTU_CODENAME")
            if [ -n "$CODENAME" ]; then
                UPSTREAM="ubuntu"
            else
                CODENAME=$(. /etc/os-release 2>/dev/null && echo "$VERSION_CODENAME")
                # For rolling/derivative distros (Kali, Parrot, etc.), Docker has no matching repo
                # Fall back to a known Debian stable codename
                if [ -z "$CODENAME" ] || echo "$CODENAME" | grep -qiE "rolling|sid|unstable"; then
                    CODENAME="bookworm"
                fi
            fi
            echo "  Detected Debian family (${DISTRO_ID}, using ${UPSTREAM}/${CODENAME})"
            # Try distro's own docker package first (works on Kali, Parrot, ARM, etc.)
            if apt install -y docker.io 2>/dev/null; then
                echo "  ✓ Installed docker.io from distro repos"
                DOCKER_INSTALLED=true
            else
                # Fall back to Docker's official repo
                apt install -y ca-certificates curl gnupg 2>/dev/null
                install -m 0755 -d /etc/apt/keyrings
                curl -fsSL "https://download.docker.com/linux/${UPSTREAM}/gpg" | gpg --dearmor -o /etc/apt/keyrings/docker.gpg 2>/dev/null
                chmod a+r /etc/apt/keyrings/docker.gpg
                echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/${UPSTREAM} ${CODENAME} stable" > /etc/apt/sources.list.d/docker.list
                apt update -qq 2>/dev/null
                if apt install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin 2>/dev/null; then
                    DOCKER_INSTALLED=true
                fi
            fi
        elif echo "$DISTRO_ID $DISTRO_LIKE" | grep -qiE "arch|manjaro"; then
            # Arch family
            echo "  Detected Arch family (${DISTRO_ID})"
            if pacman -Sy --noconfirm docker docker-compose 2>/dev/null; then
                DOCKER_INSTALLED=true
            fi
        fi
    fi

    if [ "$DOCKER_INSTALLED" = true ]; then
        systemctl enable docker 2>/dev/null || true
        systemctl start docker 2>/dev/null || true
        echo "✓ Docker installed"
    else
        echo ""
        echo "  ⚠ Docker could not be installed automatically."
        echo "    Your distro may not be directly supported by Docker's official repos."
        echo "    WolfStack will still work for LXC containers, VMs, and server management."
        echo "    To add Docker support, install it manually:"
        echo "      https://docs.docker.com/engine/install/"
        if [ "$PKG_MANAGER" = "apt" ]; then
            echo "    Or try: sudo apt install docker.io  (community package, may be older)"
        elif [ "$PKG_MANAGER" = "dnf" ]; then
            echo "    Or try: sudo dnf install docker  (community package)"
        elif [ "$PKG_MANAGER" = "pacman" ]; then
            echo "    Or try: sudo pacman -S docker  (community package)"
        elif [ "$PKG_MANAGER" = "zypper" ]; then
            echo "    Or try: sudo zypper install docker  (community package)"
        fi
        echo ""
    fi
else
    echo "✓ Docker already installed"
fi

# ─── Install WolfNet (cluster network layer) ────────────────────────────────

# Helper: download prebuilt WolfNet binaries or build from source.
# Sets WOLFNET_BUILT=true on success.
# Requires WOLFNET_SRC_DIR to be set if building from source.
build_or_download_wolfnet() {
    # Try prebuilt first
    if download_prebuilt "wolfsoftwaresystemsltd/WolfNet" "wolfnet" "/usr/local/bin/wolfnet"; then
        download_prebuilt "wolfsoftwaresystemsltd/WolfNet" "wolfnetctl" "/usr/local/bin/wolfnetctl" || true
        WOLFNET_BUILT=true
        return 0
    fi

    # Fall back to source build
    if [ -z "$WOLFNET_SRC_DIR" ] || [ ! -d "$WOLFNET_SRC_DIR" ]; then
        echo "  ✗ WolfNet source not available and no prebuilt binary — skipping"
        WOLFNET_BUILT=false
        return 1
    fi

    export PATH="${CARGO_HOME:-$REAL_HOME/.cargo}/bin:/usr/local/bin:/usr/bin:$PATH"
    if ! command -v cargo >/dev/null 2>&1; then
        echo "  ⚠ Cargo not found — skipping WolfNet rebuild"
        WOLFNET_BUILT=false
        return 1
    fi

    echo "  Building WolfNet from source..."
    cd "$WOLFNET_SRC_DIR"
    if [ -n "$CUSTOM_INSTALL_DIR" ]; then
        chown -R "$REAL_USER:$REAL_USER" "$WOLFNET_SRC_DIR" "$CARGO_HOME" "$RUSTUP_HOME" "$TMPDIR" 2>/dev/null || true
        if [ "$REAL_USER" != "root" ]; then
            su - "$REAL_USER" -c "export CARGO_HOME='$CARGO_HOME' RUSTUP_HOME='$RUSTUP_HOME' TMPDIR='$TMPDIR' PATH='$CARGO_HOME/bin:/usr/local/bin:/usr/bin:\$PATH' && cd $WOLFNET_SRC_DIR && cargo build --release"
        else
            cargo build --release
        fi
    elif [ "$REAL_USER" != "root" ] && [ -f "$REAL_HOME/.cargo/bin/cargo" ]; then
        chown -R "$REAL_USER:$REAL_USER" "$WOLFNET_SRC_DIR"
        su - "$REAL_USER" -c "cd $WOLFNET_SRC_DIR && $REAL_HOME/.cargo/bin/cargo build --release"
    else
        cargo build --release
    fi

    cp "$WOLFNET_SRC_DIR/target/release/wolfnet" /usr/local/bin/wolfnet
    chmod +x /usr/local/bin/wolfnet
    if [ -f "$WOLFNET_SRC_DIR/target/release/wolfnetctl" ]; then
        cp "$WOLFNET_SRC_DIR/target/release/wolfnetctl" /usr/local/bin/wolfnetctl
        chmod +x /usr/local/bin/wolfnetctl
    fi
    WOLFNET_BUILT=true
    return 0
}

echo ""
echo "Checking WolfNet (cluster networking)..."

if command -v wolfnet >/dev/null 2>&1 && systemctl is-active --quiet wolfnet 2>/dev/null; then
    # Already installed and running — check for upgrades
    echo "✓ WolfNet already installed and running"
    WOLFNET_IP=$(ip -4 addr show wolfnet0 2>/dev/null | awk '/inet / {split($2,a,"/"); print a[1]}' || echo "")
    if [ -n "$WOLFNET_IP" ]; then
        echo "  WolfNet IP: $WOLFNET_IP"
    fi

    # Always update WolfNet when WolfStack updates
    WOLFNET_SRC_DIR="${CUSTOM_INSTALL_DIR:-/opt}/wolfnet-src"
    if [ ! -d "$WOLFNET_SRC_DIR" ]; then
        echo "  WolfNet source not found — cloning..."
        git clone https://github.com/wolfsoftwaresystemsltd/WolfNet.git "$WOLFNET_SRC_DIR"
        git config --global --add safe.directory "$WOLFNET_SRC_DIR" 2>/dev/null || true
    fi

    echo "  Updating WolfNet..."
    cd "$WOLFNET_SRC_DIR"
    git config --global --add safe.directory "$WOLFNET_SRC_DIR" 2>/dev/null || true
    git fetch origin 2>&1 || true
    git reset --hard origin/main 2>&1 || true

    # If the existing source dir is a WolfScale clone (old layout), replace it
    if [ -f "$WOLFNET_SRC_DIR/Cargo.toml" ] && ! grep -q 'name = "wolfnet"' "$WOLFNET_SRC_DIR/Cargo.toml"; then
        echo "  Replacing old WolfScale clone with standalone WolfNet repo..."
        rm -rf "$WOLFNET_SRC_DIR"
        git clone https://github.com/wolfsoftwaresystemsltd/WolfNet.git "$WOLFNET_SRC_DIR"
        git config --global --add safe.directory "$WOLFNET_SRC_DIR" 2>/dev/null || true
    fi

    # Update binaries (prebuilt or source)
    systemctl stop wolfnet 2>/dev/null || true
    if build_or_download_wolfnet; then
        systemctl start wolfnet 2>/dev/null || true
        echo "  ✓ WolfNet updated and restarted"
    else
        systemctl start wolfnet 2>/dev/null || true
    fi

elif command -v wolfnet >/dev/null 2>&1 && [ -f "/etc/systemd/system/wolfnet.service" ]; then
    # Installed but not running — check for upgrades, then start
    echo "✓ WolfNet installed (not running)"

    # Always update WolfNet when WolfStack updates
    WOLFNET_SRC_DIR="${CUSTOM_INSTALL_DIR:-/opt}/wolfnet-src"
    if [ ! -d "$WOLFNET_SRC_DIR" ]; then
        echo "  WolfNet source not found — cloning..."
        git clone https://github.com/wolfsoftwaresystemsltd/WolfNet.git "$WOLFNET_SRC_DIR"
        git config --global --add safe.directory "$WOLFNET_SRC_DIR" 2>/dev/null || true
    fi

    echo "  Updating WolfNet..."
    cd "$WOLFNET_SRC_DIR"
    git config --global --add safe.directory "$WOLFNET_SRC_DIR" 2>/dev/null || true
    git fetch origin 2>&1 || true
    git reset --hard origin/main 2>&1 || true

    # If the existing source dir is a WolfScale clone (old layout), replace it
    if [ -f "$WOLFNET_SRC_DIR/Cargo.toml" ] && ! grep -q 'name = "wolfnet"' "$WOLFNET_SRC_DIR/Cargo.toml"; then
        echo "  Replacing old WolfScale clone with standalone WolfNet repo..."
        rm -rf "$WOLFNET_SRC_DIR"
        git clone https://github.com/wolfsoftwaresystemsltd/WolfNet.git "$WOLFNET_SRC_DIR"
        git config --global --add safe.directory "$WOLFNET_SRC_DIR" 2>/dev/null || true
    fi

    if build_or_download_wolfnet; then
        echo "  ✓ WolfNet updated"
    fi

    echo "  Starting WolfNet..."
    systemctl start wolfnet 2>/dev/null || true
    sleep 2
    if systemctl is-active --quiet wolfnet; then
        WOLFNET_IP=$(ip -4 addr show wolfnet0 2>/dev/null | awk '/inet / {split($2,a,"/"); print a[1]}' || echo "")
        echo "  ✓ WolfNet started. IP: ${WOLFNET_IP:-unknown}"
    else
        echo "  ⚠ WolfNet failed to start. Check: journalctl -u wolfnet -n 20"
    fi

else
    # WolfNet NOT installed — must install it
    echo "  WolfNet not found — installing for cluster networking..."
    echo ""

    # WolfNet needs /dev/net/tun
    SKIP_WOLFNET=false
    if [ ! -e /dev/net/tun ]; then
        echo ""
        echo "  ⚠  /dev/net/tun is NOT available!"
        echo "  ─────────────────────────────────────"
        echo ""
        echo "  This is almost certainly a Proxmox LXC container."
        echo "  WolfNet needs TUN/TAP to create its network overlay."
        echo ""
        echo "  To fix this, run the following on the Proxmox HOST (not inside the container):"
        echo ""
        echo "  1. Edit the container config:"
        echo "     nano /etc/pve/lxc/<CTID>.conf"
        echo ""
        echo "  2. Add these lines:"
        echo "     lxc.cgroup2.devices.allow: c 10:200 rwm"
        echo "     lxc.mount.entry: /dev/net dev/net none bind,create=dir"
        echo ""
        echo "  3. Restart the container:"
        echo "     pct restart <CTID>"
        echo ""
        echo "  4. Inside the container, create the device if needed:"
        echo "     mkdir -p /dev/net"
        echo "     mknod /dev/net/tun c 10 200"
        echo "     chmod 666 /dev/net/tun"
        echo ""
        echo "  Then re-run this installer."
        echo ""
        echo "  ✗ Cannot continue without WolfNet. Fix /dev/net/tun and re-run."
        exit 1
    fi

    # Try prebuilt binary first, fall back to source build
    WOLFNET_SRC_DIR="${CUSTOM_INSTALL_DIR:-/opt}/wolfnet-src"
    if ! build_or_download_wolfnet; then
        # Prebuilt failed and source wasn't available yet — clone and retry
        echo "  Downloading WolfNet source..."
        if [ -d "$WOLFNET_SRC_DIR" ]; then
            git config --global --add safe.directory "$WOLFNET_SRC_DIR" 2>/dev/null || true
            cd "$WOLFNET_SRC_DIR" && git fetch origin && git reset --hard origin/main
        else
            git clone https://github.com/wolfsoftwaresystemsltd/WolfNet.git "$WOLFNET_SRC_DIR"
            git config --global --add safe.directory "$WOLFNET_SRC_DIR" 2>/dev/null || true
            cd "$WOLFNET_SRC_DIR"
        fi

        # Ensure Rust is available for building WolfNet
        export PATH="${CARGO_HOME:-$REAL_HOME/.cargo}/bin:/usr/local/bin:/usr/bin:$PATH"

        if ! command -v cargo >/dev/null 2>&1; then
            echo "  Installing Rust first..."
            if [ -n "$CUSTOM_INSTALL_DIR" ] || [ "$REAL_USER" = "root" ]; then
                curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
            else
                su - "$REAL_USER" -c "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y"
            fi
            export PATH="${CARGO_HOME:-$REAL_HOME/.cargo}/bin:$PATH"
        fi

        build_or_download_wolfnet
    fi
    echo "  ✓ WolfNet binary installed"

    # Configure WolfNet for cluster use
    mkdir -p /etc/wolfnet /var/run/wolfnet

    if [ ! -f "/etc/wolfnet/config.toml" ]; then
        # Auto-assign a cluster IP based on the last octet of the host IP
        HOST_IP=$(ip -4 route show default 2>/dev/null | awk '/src/ {for(i=1;i<=NF;i++) if($i=="src") print $(i+1)}' | head -1)
        [ -z "$HOST_IP" ] && HOST_IP=$(ip -4 route get 1.1.1.1 2>/dev/null | awk '{for(i=1;i<=NF;i++) if($i=="src") {print $(i+1); exit}}')
        [ -z "$HOST_IP" ] && HOST_IP=$(ip -4 addr show scope global 2>/dev/null | awk '/inet / {split($2,a,"/"); print a[1]; exit}')
        [ -z "$HOST_IP" ] && HOST_IP=$(hostname -I 2>/dev/null | awk '{print $1}')
        [ -z "$HOST_IP" ] && HOST_IP=$(hostname -i 2>/dev/null | awk '{print $1}')
        LAST_OCTET=$(echo "$HOST_IP" | awk -F. '{print $4}')
        # Ensure last octet is valid (1-254); default to 1 if detection fails
        if [ -z "$LAST_OCTET" ] || [ "$LAST_OCTET" -lt 1 ] 2>/dev/null || [ "$LAST_OCTET" -gt 254 ] 2>/dev/null; then
            LAST_OCTET=1
        fi

        # Find a /24 subnet that doesn't conflict with existing networks
        # Preferred: 10.10.10.0/24, fallback: 10.10.20.0/24, 10.10.30.0/24, etc.
        WOLFNET_SUBNET=""
        for THIRD_OCTET in 10 20 30 40 50 60 70 80 90; do
            CANDIDATE="10.10.${THIRD_OCTET}.0/24"
            # Check if this subnet is already routed or has addresses assigned
            if ! ip route show 2>/dev/null | grep -q "10.10.${THIRD_OCTET}\." && \
               ! ip addr show 2>/dev/null | grep -q "10.10.${THIRD_OCTET}\."; then
                WOLFNET_SUBNET="10.10.${THIRD_OCTET}"
                break
            fi
            echo "  ⚠ Subnet $CANDIDATE already in use, trying next..."
        done

        if [ -z "$WOLFNET_SUBNET" ]; then
            echo "  ✗ Could not find a free 10.10.x.0/24 subnet!"
            echo "  Please configure WolfNet manually: /etc/wolfnet/config.toml"
            WOLFNET_SUBNET="10.10.10"  # fallback anyway
        fi

        # Check the candidate IP isn't already taken by another node
        # (e.g. two servers with the same last octet on different subnets)
        WOLFNET_IP="${WOLFNET_SUBNET}.${LAST_OCTET}"
        TRIES=0
        while [ $TRIES -lt 253 ]; do
            # Quick ping check — if nobody responds, it's free
            if ! ping -c 1 -W 1 "$WOLFNET_IP" >/dev/null 2>&1; then
                break
            fi
            echo "  ⚠ ${WOLFNET_IP} already in use, trying next..."
            LAST_OCTET=$(( (LAST_OCTET % 254) + 1 ))
            WOLFNET_IP="${WOLFNET_SUBNET}.${LAST_OCTET}"
            TRIES=$((TRIES + 1))
        done

        # Ask about LAN auto-discovery
        echo ""
        echo "  ──────────────────────────────────────────────────"
        echo "  LAN Auto-Discovery"
        echo "  ──────────────────────────────────────────────────"
        echo ""
        echo "  WolfNet can broadcast discovery packets on your local"
        echo "  network to automatically find other WolfNet nodes."
        echo ""
        echo "  ⚠  Do NOT enable on public/datacenter networks!"
        echo "     (Proxmox VLANs, Hetzner, OVH, etc.)"
        echo "     Only enable on private LANs (home, office)."
        echo ""
        echo -n "Enable LAN auto-discovery? [y/N]: "
        prompt_read ENABLE_DISCOVERY
        if [ "$ENABLE_DISCOVERY" = "y" ] || [ "$ENABLE_DISCOVERY" = "Y" ]; then
            WOLFNET_DISCOVERY="true"
        else
            WOLFNET_DISCOVERY="false"
        fi

        # Generate keys
        KEY_FILE="/etc/wolfnet/private.key"
        /usr/local/bin/wolfnet genkey --output "$KEY_FILE" 2>/dev/null || true

        cat <<EOF > /etc/wolfnet/config.toml
# WolfNet Configuration
# Auto-generated by WolfStack installer
# Provides cluster overlay network

[network]
interface = "wolfnet0"
address = "$WOLFNET_IP"
subnet = 24
listen_port = 9600
gateway = false
discovery = $WOLFNET_DISCOVERY
mtu = 1400

[security]
private_key_file = "$KEY_FILE"

# Peers will be added automatically when you add servers to WolfStack
EOF
        echo "  ✓ WolfNet configured: $WOLFNET_IP/24 (subnet: ${WOLFNET_SUBNET}.0/24)"
        if [ "$WOLFNET_DISCOVERY" = "false" ]; then
            echo "  ℹ  Discovery disabled. You can enable it later in WolfStack → WolfNet → Network Settings."
        fi
    fi

    # Create systemd service
    if [ ! -f "/etc/systemd/system/wolfnet.service" ]; then
        cat > /etc/systemd/system/wolfnet.service <<EOF
[Unit]
Description=WolfNet - Secure Private Mesh Networking
Before=wolfstack.service
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
ExecStart=/usr/local/bin/wolfnet --config /etc/wolfnet/config.toml
Restart=on-failure
RestartSec=5
LimitNOFILE=65535
DeviceAllow=/dev/net/tun rw
RuntimeDirectory=wolfnet
RuntimeDirectoryMode=0755

[Install]
WantedBy=multi-user.target
EOF
        systemctl daemon-reload
    fi

    systemctl enable wolfnet 2>/dev/null || true
    systemctl start wolfnet 2>/dev/null || true
    sleep 2

    if systemctl is-active --quiet wolfnet; then
        WOLFNET_IP=$(ip -4 addr show wolfnet0 2>/dev/null | awk '/inet / {split($2,a,"/"); print a[1]}' || echo "${WOLFNET_IP:-unknown}")
        echo "  ✓ WolfNet running! Cluster IP: $WOLFNET_IP"
    else
        echo "  ⚠ WolfNet may not have started. Check: journalctl -u wolfnet -n 20"
    fi
fi


# ─── Install Rust if not present ────────────────────────────────────────────
CARGO_BIN="${CARGO_HOME:-$REAL_HOME/.cargo}/bin/cargo"

if [ -f "$CARGO_BIN" ]; then
    echo "✓ Rust already installed"
elif command -v cargo >/dev/null 2>&1; then
    CARGO_BIN="$(command -v cargo)"
    echo "✓ Rust already installed (system-wide)"
elif command -v rustup >/dev/null 2>&1; then
    # rustup installed (e.g. via pacman on Arch) but no toolchain set yet
    echo "  Setting default Rust toolchain via rustup..."
    rustup default stable
    echo "✓ Rust installed via rustup"
else
    echo ""
    if [ -n "$CUSTOM_INSTALL_DIR" ]; then
        echo "Installing Rust to $CUSTOM_INSTALL_DIR..."
    else
        echo "Installing Rust for user '$REAL_USER'..."
    fi
    if [ -n "$CUSTOM_INSTALL_DIR" ] || [ "$REAL_USER" = "root" ]; then
        curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
    else
        su - "$REAL_USER" -c "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y"
    fi
    echo "✓ Rust installed"
fi

# Ensure cargo is found
export PATH="${CARGO_HOME:-$REAL_HOME/.cargo}/bin:/usr/local/bin:/usr/bin:$PATH"

if ! command -v cargo >/dev/null 2>&1; then
    echo "✗ cargo not found after installation. Check Rust install."
    exit 1
fi

echo "✓ Using cargo: $(command -v cargo)"

# ─── Clone or update repository ─────────────────────────────────────────────
INSTALL_DIR="${CUSTOM_INSTALL_DIR:-/opt}/wolfstack-src"
if [ -n "$CUSTOM_INSTALL_DIR" ]; then
    export CARGO_TARGET_DIR="$CUSTOM_INSTALL_DIR/wolfstack-target"
    mkdir -p "$CARGO_TARGET_DIR"
    chown -R "$REAL_USER:$REAL_USER" "$CARGO_TARGET_DIR" 2>/dev/null || true
    echo ""
    echo "  External drive build paths:"
    echo "    Source:    $INSTALL_DIR"
    echo "    Target:    $CARGO_TARGET_DIR"
    echo "    Cargo:     $CARGO_HOME"
    echo "    Rustup:    $RUSTUP_HOME"
    echo "    Tmpdir:    $TMPDIR"
fi
echo ""
echo "Cloning WolfStack repository..."

if [ -d "$INSTALL_DIR" ]; then
    echo "  Updating existing installation..."
    cd "$INSTALL_DIR"
    if ! git fetch origin 2>/dev/null; then
        echo "  ⚠ Git repo corrupted — re-cloning..."
        cd /
        rm -rf "$INSTALL_DIR"
        git clone -b $BRANCH https://github.com/wolfsoftwaresystemsltd/WolfStack.git "$INSTALL_DIR"
        cd "$INSTALL_DIR"
    else
        git checkout -B $BRANCH origin/$BRANCH
        git reset --hard origin/$BRANCH
    fi
else
    git clone -b $BRANCH https://github.com/wolfsoftwaresystemsltd/WolfStack.git "$INSTALL_DIR"
    cd "$INSTALL_DIR"
fi

# Show what we're building
BUILT_VERSION=$(grep '^version' Cargo.toml | head -1 | sed 's/.*"\(.*\)"/\1/')
echo "✓ Repository ready ($INSTALL_DIR)"
echo "  Branch: $BRANCH | Version: $BUILT_VERSION"

# ─── Build or download WolfStack ───────────────────────────────────────────
echo ""

# Flag restart if service is running (for upgrades)
if systemctl is-active --quiet wolfstack 2>/dev/null; then
    echo "WolfStack service is running — will restart after upgrade."
    RESTART_SERVICE=true
else
    RESTART_SERVICE=false
fi

echo ""
if [ -f "/usr/local/bin/wolfstack" ]; then
    echo "Upgrading WolfStack..."
else
    echo "Installing WolfStack..."
fi

# Try prebuilt binary first — saves disk space, memory, and build time
WOLFSTACK_PREBUILT=false
if download_prebuilt "wolfsoftwaresystemsltd/WolfStack" "wolfstack" "/usr/local/bin/wolfstack"; then
    WOLFSTACK_PREBUILT=true
else
    # Fall back to building from source
    echo "Building WolfStack from source (this may take a few minutes)..."

    # Force full rebuild to ensure the new version takes effect
    echo "  Cleaning previous build..."
    CLEAN_TARGET="${CARGO_TARGET_DIR:-$INSTALL_DIR/target}"
    rm -rf "$CLEAN_TARGET/release/wolfstack" "$CLEAN_TARGET/release/.fingerprint/wolfstack-"*

    # Low-memory systems (< 4GB): create swap and limit parallelism to avoid OOM
    TOTAL_MEM_KB=$(grep MemTotal /proc/meminfo | awk '{print $2}')
    TOTAL_SWAP_KB=$(grep SwapTotal /proc/meminfo | awk '{print $2}')
    TOTAL_AVAILABLE_KB=$((TOTAL_MEM_KB + TOTAL_SWAP_KB))
    CARGO_JOBS=""
    CREATED_SWAP=""

    if [ "$TOTAL_AVAILABLE_KB" -lt 4000000 ]; then
        echo "  Low memory detected ($(( TOTAL_MEM_KB / 1024 ))MB RAM + $(( TOTAL_SWAP_KB / 1024 ))MB swap)"
        CARGO_JOBS="-j 1"

        # Create a temporary swap file if total memory + swap < 4GB
        SWAP_DIR="${CUSTOM_INSTALL_DIR:-/var}"
        SWAP_FILE="$SWAP_DIR/.wolfstack-build-swap"
        NEEDED_SWAP_MB=$(( (4000000 - TOTAL_AVAILABLE_KB) / 1024 + 512 ))
        if [ "$NEEDED_SWAP_MB" -gt 4096 ]; then
            NEEDED_SWAP_MB=4096
        fi

        echo "  Creating ${NEEDED_SWAP_MB}MB temporary swap file for build..."
        dd if=/dev/zero of="$SWAP_FILE" bs=1M count="$NEEDED_SWAP_MB" status=none 2>/dev/null && \
        chmod 600 "$SWAP_FILE" && \
        mkswap "$SWAP_FILE" >/dev/null 2>&1 && \
        swapon "$SWAP_FILE" 2>/dev/null && \
        CREATED_SWAP="$SWAP_FILE" && \
        echo "  ✓ Temporary swap enabled" || \
        echo "  ⚠ Could not create swap file (build may be slow or fail)"
    fi

    if [ -n "$CUSTOM_INSTALL_DIR" ]; then
        # Custom install dir — all build I/O goes to external drive
        chown -R "$REAL_USER:$REAL_USER" "$INSTALL_DIR" "$CARGO_HOME" "$RUSTUP_HOME" "$TMPDIR" "$CARGO_TARGET_DIR" 2>/dev/null || true
        if [ "$REAL_USER" != "root" ]; then
            su - "$REAL_USER" -c "export CARGO_HOME='$CARGO_HOME' RUSTUP_HOME='$RUSTUP_HOME' TMPDIR='$TMPDIR' CARGO_TARGET_DIR='$CARGO_TARGET_DIR' PATH='$CARGO_HOME/bin:/usr/local/bin:/usr/bin:\$PATH' && cd $INSTALL_DIR && cargo build --release $CARGO_JOBS"
        else
            cargo build --release $CARGO_JOBS
        fi
    elif [ "$REAL_USER" != "root" ] && [ -f "$REAL_HOME/.cargo/bin/cargo" ]; then
        chown -R "$REAL_USER:$REAL_USER" "$INSTALL_DIR"
        su - "$REAL_USER" -c "cd $INSTALL_DIR && $REAL_HOME/.cargo/bin/cargo build --release $CARGO_JOBS"
    else
        cargo build --release $CARGO_JOBS
    fi

    # Clean up temporary swap file
    if [ -n "$CREATED_SWAP" ]; then
        swapoff "$CREATED_SWAP" 2>/dev/null
        rm -f "$CREATED_SWAP"
        echo "  ✓ Temporary swap removed"
    fi

    BUILD_TARGET_DIR="${CARGO_TARGET_DIR:-$INSTALL_DIR/target}"
    cp "$BUILD_TARGET_DIR/release/wolfstack" /usr/local/bin/wolfstack
    chmod +x /usr/local/bin/wolfstack
fi

echo "✓ wolfstack installed to /usr/local/bin/wolfstack"

# AI knowledge base is now compiled into the binary — no separate install needed
echo "✓ AI knowledge base embedded in binary"

# ─── Install WolfUSB ────────────────────────────────────────────────────────
echo ""
echo "Installing WolfUSB..."

# Install libusb (required by wolfusb)
if command -v pacman >/dev/null 2>&1; then
    pacman -S --noconfirm libusb 2>/dev/null || true
elif command -v apt-get >/dev/null 2>&1; then
    apt-get install -y libusb-1.0-0 2>/dev/null || true
elif command -v dnf >/dev/null 2>&1; then
    dnf install -y libusbx 2>/dev/null || dnf install -y libusb1 2>/dev/null || true
elif command -v zypper >/dev/null 2>&1; then
    zypper install -y libusb-1_0-0 2>/dev/null || true
fi

# Stop service for upgrade
systemctl stop wolfusb 2>/dev/null || true

# Install/update wolfusb binary via its official setup.sh (handles platform detection)
if curl -fsSL https://raw.githubusercontent.com/wolfsoftwaresystemsltd/wolfusb/main/setup.sh | bash; then
    echo "  ✓ WolfUSB binary installed"
else
    echo "  ⚠ WolfUSB install failed (non-critical)"
fi

# Configure WolfUSB with the cluster secret as its auth key
mkdir -p /etc/wolfusb
CLUSTER_SECRET_FILE="/etc/wolfstack/custom-cluster-secret"
if [ -f "$CLUSTER_SECRET_FILE" ] && [ -s "$CLUSTER_SECRET_FILE" ]; then
    WOLFUSB_KEY_VALUE=$(cat "$CLUSTER_SECRET_FILE" | tr -d '\n\r')
else
    # Use the compiled-in default (wolfstack will use this as well)
    WOLFUSB_KEY_VALUE="wsk_a7f3b9e2c1d4f6a8b0e3d5c7f9a1b3d5e7f9a1c3b5d7e9f0a2b4c6d8e0f1a3"
fi
cat > /etc/wolfusb/wolfusb.env << ENV
WOLFUSB_BIND=0.0.0.0
WOLFUSB_PORT=3240
WOLFUSB_KEY=${WOLFUSB_KEY_VALUE}
ENV
chmod 600 /etc/wolfusb/wolfusb.env

# Install systemd unit
cat > /etc/systemd/system/wolfusb.service << 'UNIT'
[Unit]
Description=WolfUSB Server
After=network.target

[Service]
Type=simple
Environment=WOLFUSB_BIND=0.0.0.0
Environment=WOLFUSB_PORT=3240
EnvironmentFile=-/etc/wolfusb/wolfusb.env
ExecStart=/usr/local/bin/wolfusb server --bind ${WOLFUSB_BIND} --port ${WOLFUSB_PORT}
Restart=on-failure
RestartSec=5

[Install]
WantedBy=multi-user.target
UNIT

# Udev rules for USB device access
mkdir -p /etc/udev/rules.d
echo 'SUBSYSTEM=="usb", MODE="0666", GROUP="plugdev"' > /etc/udev/rules.d/99-wolfusb.rules
udevadm control --reload-rules 2>/dev/null || true

# USB/IP kernel module setup.
# Each wolfstack node is both a potential client (needs vhci-hcd) AND server
# (needs usbip-host). On most distros these live in a "kernel-modules-extra"
# style package that ISN'T installed by default. First try loading; if that
# fails, install the right package for the distro, then try again.
wolfusb_try_modprobe() {
    modprobe vhci-hcd 2>/dev/null || modprobe vhci_hcd 2>/dev/null || true
    modprobe usbip-host 2>/dev/null || modprobe usbip_host 2>/dev/null || true
    [ -d /sys/devices/platform/vhci_hcd.0 ] && \
        [ -d /sys/bus/usb/drivers/usbip-host ]
}

wolfusb_install_modules_pkg() {
    # Read distro once
    local ID="" LIKE=""
    if [ -r /etc/os-release ]; then
        ID=$(. /etc/os-release && echo "${ID:-}")
        LIKE=$(. /etc/os-release && echo "${ID_LIKE:-}")
    fi
    case "$ID $LIKE" in
        *arch*|*manjaro*|*cachyos*|*endeavouros*)
            # Arch: modules come with the kernel package; nothing extra needed.
            return 1
            ;;
        *fedora*|*rhel*|*centos*|*rocky*|*alma*)
            # Fedora/RHEL family: kernel-modules-extra
            dnf install -y kernel-modules-extra 2>/dev/null || \
                yum install -y kernel-modules-extra 2>/dev/null || return 1
            ;;
        *debian*|*ubuntu*|*pop*|*linuxmint*|*elementary*|*raspbian*)
            # Debian/Ubuntu family: linux-modules-extra-$(uname -r)
            # Fall back to the meta-package when the exact kernel build isn't
            # available (common with third-party/self-built kernels).
            apt-get update 2>/dev/null || true
            apt-get install -y "linux-modules-extra-$(uname -r)" 2>/dev/null || \
                apt-get install -y linux-modules-extra-generic 2>/dev/null || \
                apt-get install -y linux-image-extra-"$(uname -r)" 2>/dev/null || \
                return 1
            ;;
        *suse*|*sles*|*opensuse*)
            zypper install -y kernel-default-extra 2>/dev/null || return 1
            ;;
        *alpine*)
            # Alpine's default kernel (linux-lts/linux-virt) has usbip built in
            # for some image variants, missing on others. Try the modules pkg.
            apk add --no-cache linux-lts 2>/dev/null || return 1
            ;;
        *)
            # Unknown distro — can't guess the package name
            return 1
            ;;
    esac
    return 0
}

mkdir -p /etc/modules-load.d
printf 'vhci-hcd\nusbip-core\nusbip-host\n' > /etc/modules-load.d/wolfusb.conf

if ! wolfusb_try_modprobe; then
    echo "  USB/IP kernel modules not available — installing modules package..."
    if wolfusb_install_modules_pkg && wolfusb_try_modprobe; then
        :  # success
    fi
fi

if [ -d /sys/devices/platform/vhci_hcd.0 ] && \
   [ -d /sys/bus/usb/drivers/usbip-host ]; then
    echo "  ✓ USB/IP kernel modules loaded (vhci-hcd + usbip-host)"
    echo "    Node can both share local USB devices and mount remote ones."
else
    echo "  ⚠ USB/IP kernel modules unavailable on this kernel."
    echo "    Remote USB device passthrough will not work until these are"
    echo "    installed. Try:"
    echo "      Fedora/RHEL:   dnf install kernel-modules-extra && reboot"
    echo "      Debian/Ubuntu: apt install linux-modules-extra-\$(uname -r)"
    echo "      openSUSE:      zypper install kernel-default-extra && reboot"
    echo "      Arch:          usually already present; ensure stock linux kernel"
    echo "    Container/cloud-optimised kernels (GCP COS, Bottlerocket, etc.)"
    echo "    generally cannot run usbip-host and are not supported."
fi

systemctl daemon-reload
systemctl enable wolfusb 2>/dev/null || true
systemctl restart wolfusb 2>/dev/null || systemctl start wolfusb 2>/dev/null || true

if systemctl is-active --quiet wolfusb 2>/dev/null; then
    echo "  ✓ WolfUSB service running on port 3240"
else
    echo "  ⚠ WolfUSB service not running — check: journalctl -u wolfusb -n 20"
fi

# ─── Install web UI ─────────────────────────────────────────────────────────
echo ""
echo "Installing web UI..."
mkdir -p /opt/wolfstack/web
cp -r "$INSTALL_DIR/web/"* /opt/wolfstack/web/
echo "✓ Web UI installed to /opt/wolfstack/web"

# ─── Configuration ──────────────────────────────────────────────────────────
if [ ! -f "/etc/wolfstack/config.toml" ]; then
    echo ""
    echo "  ──────────────────────────────────────────────────"
    echo "  WolfStack Configuration"
    echo "  ──────────────────────────────────────────────────"
    echo ""

    # Prompt for port
    echo -n "Dashboard port [8553]: "
    prompt_read WS_PORT
    WS_PORT=$(echo "$WS_PORT" | tr -d '[:space:][:cntrl:]')
    WS_PORT=${WS_PORT:-8553}
    # Validate port is a number between 1-65535, fallback to default
    if ! echo "$WS_PORT" | grep -qE '^[0-9]+$' || [ "$WS_PORT" -lt 1 ] 2>/dev/null || [ "$WS_PORT" -gt 65535 ] 2>/dev/null; then
        echo "  ⚠ Invalid port '$WS_PORT' — using default 8553"
        WS_PORT=8553
    fi

    # Prompt for bind address
    echo -n "Bind address [0.0.0.0]: "
    prompt_read WS_BIND
    WS_BIND=$(echo "$WS_BIND" | tr -d '[:cntrl:]' | xargs)
    WS_BIND=${WS_BIND:-0.0.0.0}
    # Validate bind is a valid IP pattern, fallback to default
    if ! echo "$WS_BIND" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$'; then
        echo "  ⚠ Invalid bind address '$WS_BIND' — using default 0.0.0.0"
        WS_BIND="0.0.0.0"
    fi

    # Write config
    mkdir -p /etc/wolfstack
    cat <<EOF > /etc/wolfstack/config.toml
# WolfStack Configuration
# Generated by setup.sh

[server]
port = $WS_PORT
bind = "$WS_BIND"
web_dir = "/opt/wolfstack/web"
EOF
    echo "✓ Config created at /etc/wolfstack/config.toml"
    echo ""
    echo "  Dashboard: http://$WS_BIND:$WS_PORT"
else
    echo ""
    echo "✓ Config already exists at /etc/wolfstack/config.toml"
    echo "  (Upgrade mode - skipping configuration prompts)"
    # Read port and bind from existing config
    WS_PORT=$(grep "^port" /etc/wolfstack/config.toml 2>/dev/null | head -1 | awk '{print $3}' | tr -d '[:space:][:cntrl:]' || echo "8553")
    WS_PORT=${WS_PORT:-8553}
    WS_BIND=$(grep "^bind" /etc/wolfstack/config.toml 2>/dev/null | head -1 | sed 's/.*"\(.*\)"/\1/' | tr -d '[:space:][:cntrl:]' || echo "0.0.0.0")
    WS_BIND=${WS_BIND:-0.0.0.0}
fi

# ─── Create systemd service ─────────────────────────────────────────────────
if [ ! -f "/etc/systemd/system/wolfstack.service" ]; then
    echo ""
    echo "  ──────────────────────────────────────────────────"
    echo "  Creating systemd service..."
    echo "  ──────────────────────────────────────────────────"
    echo ""

    cat > /etc/systemd/system/wolfstack.service <<EOF
[Unit]
Description=WolfStack - Server Management Platform
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
ExecStart=/usr/local/bin/wolfstack --port $WS_PORT --bind $WS_BIND
WorkingDirectory=/opt/wolfstack
Restart=on-failure
RestartSec=5
LimitNOFILE=65535

# Must run as root for Linux auth and service management
User=root
Group=root

# Logging
StandardOutput=journal
StandardError=journal
SyslogIdentifier=wolfstack

[Install]
WantedBy=multi-user.target
EOF

    systemctl daemon-reload
    echo "✓ Systemd service created"

    # Enable and optionally start
    echo ""
    echo -n "Start WolfStack now? [Y/n]: "
    prompt_read start_now
    if [ "$start_now" != "n" ] && [ "$start_now" != "N" ]; then
        systemctl enable wolfstack
        systemctl start wolfstack
        sleep 2
        if systemctl is-active --quiet wolfstack; then
            echo "✓ WolfStack is running!"
        else
            echo "⚠ WolfStack may have failed to start. Check: journalctl -u wolfstack -n 20"
        fi
    else
        systemctl enable wolfstack
        echo "✓ WolfStack enabled (will start on boot)"
    fi
else
    echo ""
    echo "✓ Service already installed - reloading systemd"
    systemctl daemon-reload
fi

# ─── Firewall ───────────────────────────────────────────────────────────────
echo ""
if command -v ufw >/dev/null 2>&1; then
    ufw allow "$WS_PORT/tcp" 2>/dev/null && echo "✓ Firewall: Opened port $WS_PORT/tcp (ufw)" || true
    ufw allow 9600/udp 2>/dev/null && echo "✓ Firewall: Opened port 9600/udp for WolfNet (ufw)" || true
elif command -v firewall-cmd >/dev/null 2>&1; then
    firewall-cmd --permanent --add-port="$WS_PORT/tcp" 2>/dev/null && \
    firewall-cmd --permanent --add-port="9600/udp" 2>/dev/null && \
    firewall-cmd --reload 2>/dev/null && \
    echo "✓ Firewall: Opened port $WS_PORT/tcp and 9600/udp (firewalld)" || true
fi

# ─── Set up lxcbr0 bridge for LXC containers ────────────────────────────────
if command -v lxc-ls >/dev/null 2>&1; then
    # Only configure lxc-net on fresh installs — restarting lxc-net on upgrades
    # destroys lxcbr0 and all container kernel routes, breaking WolfNet routing.
    # WolfStack's reapply_wolfnet_routes() handles route restoration on startup.
    if ip link show lxcbr0 >/dev/null 2>&1 && ip -4 addr show lxcbr0 2>/dev/null | grep -q "inet "; then
        echo "✓ LXC networking already active (lxcbr0 up)"
    else
        echo ""
        echo "Configuring LXC networking (lxc-net)..."
        
        # Ensure USE_LXC_BRIDGE="true" in /etc/default/lxc-net
        if [ -f "/etc/default/lxc-net" ]; then
            if grep -q "USE_LXC_BRIDGE" /etc/default/lxc-net; then
                sed -i 's/^#\?USE_LXC_BRIDGE=.*/USE_LXC_BRIDGE="true"/' /etc/default/lxc-net
            else
                echo 'USE_LXC_BRIDGE="true"' >> /etc/default/lxc-net
            fi
        else
            echo 'USE_LXC_BRIDGE="true"' > /etc/default/lxc-net
        fi

        # Enable and start lxc-net service
        systemctl enable lxc-net 2>/dev/null || true
        systemctl restart lxc-net 2>/dev/null || true
        
        # Check if dnsmasq is running on lxcbr0
        sleep 2
        if pgrep -f "dnsmasq.*lxcbr0" > /dev/null; then
            echo "✓ LXC networking active (lxcbr0 + dnsmasq)"
        else
            echo "⚠ LXC networking service started but dnsmasq not detected on lxcbr0."
            echo "  Attempting manual fallback..."
            systemctl stop lxc-net 2>/dev/null || true
            
            ip link add lxcbr0 type bridge 2>/dev/null || true
            ip addr add 10.0.3.1/24 dev lxcbr0 2>/dev/null || true
            ip link set lxcbr0 up 2>/dev/null || true
            
            # NAT
            echo 1 > /proc/sys/net/ipv4/ip_forward 2>/dev/null || true
            iptables -t nat -A POSTROUTING -s 10.0.3.0/24 ! -d 10.0.3.0/24 -j MASQUERADE 2>/dev/null || true
            iptables -A FORWARD -i lxcbr0 -j ACCEPT 2>/dev/null || true
            iptables -A FORWARD -o lxcbr0 -m state --state RELATED,ESTABLISHED -j ACCEPT 2>/dev/null || true
            
            # DNSMasq
            mkdir -p /run/lxc
            dnsmasq --strict-order --bind-interfaces --pid-file=/run/lxc/dnsmasq.pid \
                --listen-address 10.0.3.1 --dhcp-range 10.0.3.2,10.0.3.254 \
                --dhcp-lease-max=253 --dhcp-no-override --except-interface=lo \
                --interface=lxcbr0 --conf-file= 2>/dev/null || true
                
            echo "✓ Manually configured lxcbr0 and dnsmasq"
        fi
    fi
fi

# ─── Done ────────────────────────────────────────────────────────────────────
echo ""
# Portable IP detection — `hostname -I` is GNU-only; Arch/BSD use `hostname -i`
# `ip` is the most reliable fallback on modern Linux.
get_primary_ip() {
    local ip=""
    # Default route src (works without internet connectivity)
    ip=$(ip -4 route show default 2>/dev/null | awk '/src/ {for(i=1;i<=NF;i++) if($i=="src") print $(i+1)}' | head -1)
    if [ -n "$ip" ]; then echo "$ip"; return; fi
    # Route to public IP (needs connectivity, but works everywhere)
    ip=$(ip -4 route get 1.1.1.1 2>/dev/null | awk '{for(i=1;i<=NF;i++) if($i=="src") {print $(i+1); exit}}')
    if [ -n "$ip" ]; then echo "$ip"; return; fi
    # First global IPv4 (may be VPN/tailscale, but beats nothing)
    ip=$(ip -4 addr show scope global 2>/dev/null | awk '/inet / {split($2,a,"/"); print a[1]; exit}')
    if [ -n "$ip" ]; then echo "$ip"; return; fi
    # GNU hostname -I
    ip=$(hostname -I 2>/dev/null | awk '{print $1}')
    if [ -n "$ip" ]; then echo "$ip"; return; fi
    # Non-GNU hostname -i
    ip=$(hostname -i 2>/dev/null | awk '{print $1}')
    if [ -n "$ip" ] && [ "$ip" != "127.0.0.1" ]; then echo "$ip"; return; fi
    echo "localhost"
}

echo "  🐺 Installation Complete!"
echo "  ─────────────────────────────────────"
echo "  Dashboard:  http://$(get_primary_ip):${WS_PORT}"
echo "  Login:      Use your Linux system username and password"
echo ""
echo "  Manage:"
echo "  Status:     sudo systemctl status wolfstack"
echo "  Logs:       sudo journalctl -u wolfstack -f"
echo "  Restart:    sudo systemctl restart wolfstack"
echo "  Config:     /etc/wolfstack/config.toml"
echo ""
echo "**** UPGRADE COMPLETE ****"
echo ""
echo "Please Refresh your browser if upgrading..."

# ─── Restart service if upgrading (must be last!) ────────────────────────────
if [ "$RESTART_SERVICE" = "true" ]; then
    nohup bash -c "sleep 3 && systemctl restart wolfstack" >/dev/null 2>&1 &
fi