#!/bin/bash

# 00-system-prep.sh
# Updates system, bootstraps sudo, enables Debian extra components,
# and sets up Flathub.
# Relies on DISTRO and PACKAGE_MANAGER being set by the caller.

echo "--- Starting System Preparation ---"

if [ -z "$DISTRO" ] || [ -z "$PACKAGE_MANAGER" ]; then
    echo "ERROR: DISTRO and PACKAGE_MANAGER must be set in the environment."
    exit 1
fi

# Ensure 'sudo' exists so all later scripts can use it.
# - Running as root on a fresh Debian: installing sudo lets root invoke
#   `sudo ...` directly (root is not subject to sudoers).
# - Non-root users without sudo are handled by main-setup.sh's preflight.
if ! command -v sudo &>/dev/null; then
    echo "sudo not found. Bootstrapping it..."
    if [ "$(id -u)" -eq 0 ]; then
        if [ "$PACKAGE_MANAGER" == "apt" ]; then
            apt-get update && apt-get install -y sudo
        elif [ "$PACKAGE_MANAGER" == "dnf" ]; then
            dnf install -y sudo
        fi
        echo "sudo installed. Add your normal user to the sudo group if desired:"
        echo "  usermod -aG sudo <your-user>   (then log out/in)"
    else
        echo "ERROR: 'sudo' is required but missing, and you are not root."
        echo "As root, run:  apt install -y sudo && usermod -aG sudo \$USER"
        exit 1
    fi
fi

# Update system before installing packages
echo "Updating system packages..."
if [ "$PACKAGE_MANAGER" == "dnf" ]; then
    sudo dnf update -y && sudo dnf upgrade -y
elif [ "$PACKAGE_MANAGER" == "apt" ]; then
    sudo apt update && sudo apt upgrade -y
else
    echo "WARNING: Unknown package manager '$PACKAGE_MANAGER'. Skipping system update."
fi

# ---------------------------------------------------------------------------
# Debian extras (only on Debian; idempotent)
# ---------------------------------------------------------------------------
if [ "$DISTRO" == "debian" ]; then

    # Enable contrib + non-free-firmware components on the main archive.
    # Needed for some WiFi/GPU firmware and non-free tooling. Skip with SKIP_NONFREE=1.
    if [ -z "$SKIP_NONFREE" ]; then
        echo "Enabling 'contrib' and 'non-free-firmware' apt components..."
        _enable_debian_components() {
            local f="/etc/apt/sources.list.d/debian.sources"
            local tmp
            # deb822 format used by Debian 12+ (trixie ships debian.sources)
            if [ -f "$f" ]; then
                tmp="$(mktemp)"
                if awk '
                    /^URIs:[[:space:]]/ { in_archive = ($0 ~ /deb\.debian\.org\/debian/) }
                    in_archive && /^Components:/ && $0 !~ /contrib/ {
                        print $0 " contrib non-free-firmware"
                        in_archive = 0
                        next
                    }
                    { print }
                ' "$f" > "$tmp"; then
                    if ! diff -q "$f" "$tmp" >/dev/null; then
                        sudo cp "$f" "${f}.bak"
                        sudo mv "$tmp" "$f"
                        echo "  -> updated ${f}"
                    else
                        rm -f "$tmp"
                        echo "  -> components already present in ${f}"
                    fi
                else
                    rm -f "$tmp"
                    echo "WARNING: could not parse ${f}; skipping component changes."
                fi
            fi
            # Legacy single-line /etc/apt/sources.list fallback
            if [ -f /etc/apt/sources.list ] \
                && grep -qE '^deb[[:space:]]+https?://deb\.debian\.org/debian' /etc/apt/sources.list \
                && ! grep -qE '^deb[[:space:]]+https?://deb\.debian\.org/debian[^#]*non-free-firmware' /etc/apt/sources.list; then
                sudo sed -i -E 's/^(deb[[:space:]]+https?:\/\/deb\.debian\.org\/debian[^#]*main)([[:space:]]|$)/\1 contrib non-free-firmware/' /etc/apt/sources.list
                echo "  -> updated /etc/apt/sources.list"
            fi
            sudo apt update
        }
        _enable_debian_components
    else
        echo "SKIP_NONFREE is set - leaving apt components unchanged."
    fi

    # needrestart prompts (or auto-restarts) services after library upgrades.
    if ! dpkg-query -W -f='${Status}' needrestart 2>/dev/null | grep -q "ok installed"; then
        echo "Installing needrestart..."
        sudo apt install -y needrestart
    else
        echo "needrestart already installed."
    fi
fi

# Setup Flatpak
echo "Setting up Flathub repository..."
if command -v flatpak &> /dev/null; then
    flatpak remote-add --if-not-exists flathub https://dl.flathub.org/repo/flathub.flatpakrepo
else
    echo "WARNING: flatpak command not found. Skipping Flathub setup."
    echo "Ensure Flatpak is installed via 01-package-install.sh or manually."
fi

echo "--- System Preparation Finished ---"