#!/bin/bash
set -e

# =============================================================================
# iClient installer
#
# The script is organised into discrete steps. By default each step pauses and
# asks the user before running, so the operator can skip parts that are not
# needed on a given machine.
#
# Command-line options:
#   -y, --yes, --unattended    Do not prompt. Run every step that is not
#                              explicitly skipped.
#   --skip=step1,step2,...     Comma-separated list of step IDs to skip.
#   --only=step1,step2,...     Run only the listed step IDs (skip everything
#                              else). Mutually exclusive with --skip.
#   --list                     List the available step IDs and exit.
#   -h, --help                 Show this help text and exit.
#
# Available step IDs (in execution order):
#   user             Create the 'iteslive' system user
#   packages         Install base packages (window manager, audio, etc.)
#   gpu              Detect GPU(s) and install matching drivers
#   download         Download and install the iClient binaries
#   displaymanager   Configure the display manager for auto-login
#   windowmanager    Configure the window manager autostart for iwatcher
#   apparmor         Install AppArmor profile (Ubuntu 23.10+/24.04+)
# =============================================================================

# -------- CLI parsing --------------------------------------------------------

export PATH=$PATH:/usr/sbin

UNATTENDED=0
SKIP_LIST=""
ONLY_LIST=""

ALL_STEPS=(user packages gpu download displaymanager windowmanager apparmor)

print_help() {
    sed -n '3,28p' "$0"
}

print_steps() {
    echo "Available steps:"
    for s in "${ALL_STEPS[@]}"; do
        echo "  - $s"
    done
}

for arg in "$@"; do
    case "$arg" in
        -y|--yes|--unattended)
            UNATTENDED=1
            ;;
        --skip=*)
            SKIP_LIST="${arg#--skip=}"
            ;;
        --only=*)
            ONLY_LIST="${arg#--only=}"
            ;;
        --list)
            print_steps
            exit 0
            ;;
        -h|--help)
            print_help
            exit 0
            ;;
        *)
            echo "Unknown argument: $arg"
            print_help
            exit 1
            ;;
    esac
done

if [ -n "$SKIP_LIST" ] && [ -n "$ONLY_LIST" ]; then
    echo "Error: --skip and --only cannot be used together."
    exit 1
fi

# Returns 0 if the given step ID should run, 1 otherwise.
should_run_step() {
    local step="$1"
    local item

    if [ -n "$ONLY_LIST" ]; then
        IFS=',' read -ra _ONLY <<<"$ONLY_LIST"
        for item in "${_ONLY[@]}"; do
            [ "$item" = "$step" ] && return 0
        done
        return 1
    fi

    if [ -n "$SKIP_LIST" ]; then
        IFS=',' read -ra _SKIP <<<"$SKIP_LIST"
        for item in "${_SKIP[@]}"; do
            [ "$item" = "$step" ] && return 1
        done
    fi
    return 0
}

# Prompts the user about a step. Honours --unattended.
# Usage: confirm_step <step-id> "<human description>"
# Returns 0 to run, 1 to skip.
confirm_step() {
    local step="$1"
    local desc="$2"

    echo
    echo "============================================================"
    echo "STEP: $step"
    echo "$desc"
    echo "============================================================"

    if ! should_run_step "$step"; then
        echo ">>> Skipping step '$step' (excluded via --skip/--only)."
        return 1
    fi

    if [ "$UNATTENDED" -eq 1 ]; then
        echo ">>> Unattended mode: running step '$step'."
        return 0
    fi

    local reply
    read -r -p "Run this step? [Y/n/q to quit] " reply
    case "$reply" in
        ""|y|Y|yes|YES)
            return 0
            ;;
        q|Q|quit|QUIT)
            echo "Aborted by user."
            exit 0
            ;;
        *)
            echo ">>> Skipping step '$step' by user request."
            return 1
            ;;
    esac
}

# -------- Root / sudo detection ---------------------------------------------

# Detect if running as root; if so, no need to prefix commands with sudo
if [ "$(id -u)" -eq 0 ]; then
    SUDO=""
else
    if ! command -v sudo &>/dev/null; then
        echo "This script must be run as root or with sudo available. Exiting."
        exit 1
    fi
    SUDO="sudo"
fi

# Sanity check: this installer is Debian-based only
if ! command -v apt-get &>/dev/null; then
    echo "This script only supports Debian-based distributions. Exiting."
    exit 1
fi

# Shared variables populated by early steps and consumed by later ones.
WINDOW_MANAGER=""
TERMINAL_EMULATOR=""
TARGET_USER="iteslive"

# =============================================================================
# STEP: user
# =============================================================================
step_user() {
    if id "$TARGET_USER" &>/dev/null; then
        echo "User '$TARGET_USER' already exists. Will use that user."
    else
        $SUDO useradd "$TARGET_USER" -m -s /bin/bash
    fi
}

# =============================================================================
# STEP: packages
# =============================================================================
step_packages() {
    # Returns the name of an available apt package, picking the first one that exists.
    pick_pkg() {
        for p in "$@"; do
            if apt-cache show "$p" 2>/dev/null | grep -q '^Package:'; then
                echo "$p"
                return 0
            fi
        done
        echo "$1"
        return 1
    }

    echo "Updating APT cache..."
    $SUDO apt-get update

    # Debian 13+ / Ubuntu 24.04+ ship the time_t-64 transition package "libasound2t64".
    # Older releases still provide "libasound2". Pick whichever the repos expose.
    ALSA_PKG="$(pick_pkg libasound2t64 libasound2)"
    echo "Using ALSA package: $ALSA_PKG"

    # Check if labwc is available in the repos
    if apt-cache show "labwc" &>/dev/null; then
        echo "Package 'labwc' is available in APT repositories."
        WINDOW_MANAGER="labwc"
        TERMINAL_EMULATOR="alacritty"
    else
        echo "Package 'labwc' is NOT available. Falling back to sway."
        WINDOW_MANAGER="sway"
        TERMINAL_EMULATOR="foot"
    fi

    echo "Installing base packages..."
    $SUDO apt-get -y install libicu-dev unzip $WINDOW_MANAGER grim libnss3 wlr-randr \
        $ALSA_PKG alsa-utils pulseaudio $TERMINAL_EMULATOR polkitd kanshi
}

# Make sure WINDOW_MANAGER is set even if the 'packages' step was skipped,
# because later steps need it. Best-effort: use labwc if installed, else sway.
ensure_window_manager_var() {
    if [ -z "$WINDOW_MANAGER" ]; then
        if command -v labwc &>/dev/null; then
            WINDOW_MANAGER="labwc"
            TERMINAL_EMULATOR="alacritty"
        elif command -v sway &>/dev/null; then
            WINDOW_MANAGER="sway"
            TERMINAL_EMULATOR="foot"
        else
            WINDOW_MANAGER="labwc"
            TERMINAL_EMULATOR="alacritty"
        fi
        echo "(packages step skipped) Assuming WINDOW_MANAGER=$WINDOW_MANAGER"
    fi
}

# =============================================================================
# STEP: gpu
# =============================================================================
step_gpu() {
    # Detect distro (debian | ubuntu) and codename
    DISTRO_ID="$(. /etc/os-release && echo "$ID")"
    DISTRO_LIKE="$(. /etc/os-release && echo "$ID_LIKE")"
    DISTRO_CODENAME="$(lsb_release -cs 2>/dev/null || . /etc/os-release && echo "${VERSION_CODENAME:-}")"

    case "$DISTRO_ID" in
        ubuntu|debian) ;;
        *)
            case "$DISTRO_LIKE" in
                *ubuntu*|*debian*) ;;
                *)
                    echo "Unsupported distribution '$DISTRO_ID'. GPU driver install supports Debian/Ubuntu only."
                    return 0
                    ;;
            esac
            ;;
    esac

    # Grab the VGA / 3D / Display controller line(s) from lspci
    GPU_INFO="$(lspci -nn | grep -Ei 'vga|3d|display' || true)"
    echo "Detected GPU(s):"
    echo "$GPU_INFO"

    # Classify vendor(s) present
    HAS_NVIDIA=0; HAS_AMD=0; HAS_INTEL=0; HAS_VMWARE=0; HAS_VIRTIO=0
    echo "$GPU_INFO" | grep -qi 'nvidia'                       && HAS_NVIDIA=1
    echo "$GPU_INFO" | grep -Eqi 'amd|advanced micro|ati|radeon' && HAS_AMD=1
    echo "$GPU_INFO" | grep -qi 'intel'                        && HAS_INTEL=1
    echo "$GPU_INFO" | grep -Eqi 'vmware|virtualbox|qxl|cirrus' && HAS_VMWARE=1
    echo "$GPU_INFO" | grep -Eqi 'virtio|red hat.*virtio'      && HAS_VIRTIO=1

    install_intel() {
        echo "Installing Intel GPU drivers..."
        # i965 (older) + iris (Gen8+) userspace drivers, plus VAAPI for HW video decode
        $SUDO apt-get -y install \
            intel-media-va-driver-non-free \
            i965-va-driver \
            vainfo \
            libva-drm2 libva-x11-2 || \
        $SUDO apt-get -y install intel-media-va-driver i965-va-driver vainfo libva-drm2 libva-x11-2
    }

    install_amd() {
        echo "Installing AMD GPU drivers..."
        # Open-source AMDGPU/radeon stack + VAAPI/VDPAU
        $SUDO apt-get -y install \
            firmware-amd-graphics \
            libgl1-mesa-dri \
            mesa-va-drivers \
            mesa-vdpau-drivers \
            vainfo || true
        # firmware-amd-graphics lives in non-free-firmware on Debian; ignore failure on Ubuntu
    }

    install_nvidia() {
        echo "Installing NVIDIA GPU drivers..."
        if [ "$DISTRO_ID" = "ubuntu" ] || echo "$DISTRO_LIKE" | grep -qi ubuntu; then
            # Ubuntu: ubuntu-drivers picks the best supported proprietary driver
            $SUDO apt-get -y install ubuntu-drivers-common
            $SUDO ubuntu-drivers autoinstall || {
                echo "ubuntu-drivers autoinstall failed; falling back to nvidia-driver-535"
                $SUDO apt-get -y install nvidia-driver-535
            }
        else
            # Debian: requires contrib + non-free-firmware components
            if ! grep -REqs 'non-free-firmware' /etc/apt/sources.list /etc/apt/sources.list.d; then
                echo "Enabling contrib + non-free-firmware for Debian..."
                $SUDO sed -i -E 's/^(deb .* main)( contrib)?( non-free)?( non-free-firmware)?/\1 contrib non-free non-free-firmware/' /etc/apt/sources.list
                $SUDO apt-get update
            fi
            $SUDO apt-get -y install nvidia-driver firmware-misc-nonfree
        fi

        # Enable DRM kernel-mode setting (required for Wayland/labwc on NVIDIA)
        echo "options nvidia-drm modeset=1" | $SUDO tee /etc/modprobe.d/nvidia-drm.conf >/dev/null
        if command -v update-initramfs &>/dev/null; then
            $SUDO update-initramfs -u
        fi
    }

    install_virt() {
        echo "Installing virtualized GPU drivers (VMware/VirtualBox/QXL/virtio)..."
        $SUDO apt-get -y install \
            xserver-xorg-video-qxl \
            xserver-xorg-video-vmware \
            spice-vdagent || true
    }

    # Install in priority order
    [ "$HAS_INTEL"  = "1" ] && install_intel
    [ "$HAS_AMD"    = "1" ] && install_amd
    [ "$HAS_NVIDIA" = "1" ] && install_nvidia
    { [ "$HAS_VMWARE" = "1" ] || [ "$HAS_VIRTIO" = "1" ]; } && install_virt

    if [ "$HAS_NVIDIA$HAS_AMD$HAS_INTEL$HAS_VMWARE$HAS_VIRTIO" = "00000" ]; then
        echo "WARNING: No known GPU vendor detected."
    fi
}

# =============================================================================
# STEP: download
# =============================================================================
step_download() {
    # Detect arch as either x64 or arm64
    case "$(uname -m)" in
        x86_64|amd64) arch=x64 ;;
        aarch64|arm64) arch=arm64 ;;
        *)
            echo "Unsupported architecture: $(uname -m). Exiting."
            exit 1
            ;;
    esac
    echo "Detected architecture: $arch"

    echo "Downloading iteslive..."

    DOWNLOAD_URL="https://files-cac.itesmedia.tv/iclient/main/linux-$arch"
    if command -v curl &>/dev/null; then
        curl -fL "$DOWNLOAD_URL" -o /tmp/iteslive.zip
    elif command -v wget &>/dev/null; then
        wget -nv "$DOWNLOAD_URL" -O /tmp/iteslive.zip
    else
        echo "Neither curl nor wget is available. Exiting."
        exit 1
    fi

    $SUDO mkdir -p /home/iteslive/iteslive
    $SUDO unzip -q -o /tmp/iteslive.zip -d /home/iteslive/iteslive/
    rm /tmp/iteslive.zip
    $SUDO chmod +x /home/iteslive/iteslive/iwatcher || true
    $SUDO chown -R iteslive:iteslive /home/iteslive/iteslive/
}

# =============================================================================
# STEP: displaymanager
# =============================================================================
step_displaymanager() {
    ensure_window_manager_var

    configure_lightdm() {
        echo "Configuring LightDM for auto login..."
        $SUDO systemctl enable lightdm

        $SUDO mkdir -p /etc/lightdm/lightdm.conf.d
        $SUDO tee /etc/lightdm/lightdm.conf.d/50-iteslive-autologin.conf >/dev/null <<EOF
[Seat:*]
autologin-user=${TARGET_USER}
autologin-user-timeout=0
autologin-session=${WINDOW_MANAGER}
user-session=${WINDOW_MANAGER}
greeter-session=lightdm-gtk-greeter
EOF
    }

    configure_gdm3() {
        echo "Configuring GDM3 for auto login..."
        GDM3_CONFIG="/etc/gdm3/daemon.conf"
        $SUDO cp "$GDM3_CONFIG" "${GDM3_CONFIG}.bak"

        $SUDO sed -i '/^\[daemon\]/a AutomaticLoginEnable=true\nAutomaticLogin='"$TARGET_USER" "$GDM3_CONFIG"
        $SUDO sed -i '/^AutomaticLoginEnable=/d; /^AutomaticLogin=/d' "$GDM3_CONFIG"
        $SUDO sed -i 's/^#\s*AutomaticLoginEnable.*/AutomaticLoginEnable=true/' "$GDM3_CONFIG"
        $SUDO sed -i 's/^#\s*AutomaticLogin.*/AutomaticLogin='"$TARGET_USER"'/' "$GDM3_CONFIG"

        $SUDO mkdir -p /var/lib/AccountsService/users
        $SUDO tee /var/lib/AccountsService/users/"$TARGET_USER" >/dev/null <<EOF
[User]
Session=${WINDOW_MANAGER}
EOF
    }

    configure_sddm() {
        echo "Configuring SDDM for auto login..."
        $SUDO mkdir -p /etc/sddm.conf.d
        $SUDO tee /etc/sddm.conf.d/autologin.conf >/dev/null <<EOF
[Autologin]
User=${TARGET_USER}
Session=${WINDOW_MANAGER}
EOF
    }

    echo "Checking running display manager..."
    displayManager=$(basename $(test -L /etc/systemd/system/display-manager.service && readlink /etc/systemd/system/display-manager.service || echo "none") .service)
    echo "Running display manager: $displayManager"
    case "$displayManager" in
        lightdm) configure_lightdm ;;
        gdm3)    configure_gdm3 ;;
        sddm)    configure_sddm ;;
        lxdm|xdm)
            echo "Unsupported display manager '$displayManager'. Please switch to LightDM, GDM3 or SDDM manually and re-run this step."
            exit 1
            ;;
        *)
            echo "No supported display manager detected. Installing LightDM..."
            $SUDO apt-get -y install lightdm
            configure_lightdm
            ;;
    esac
}

# =============================================================================
# STEP: windowmanager
# =============================================================================
step_windowmanager() {
    ensure_window_manager_var

    case "$WINDOW_MANAGER" in
        labwc)
            echo "Configuring labwc as the window manager..."
            $SUDO mkdir -p /home/$TARGET_USER/.config/labwc
            echo -e "cd ~/iteslive\n./iwatcher &" | $SUDO tee /home/$TARGET_USER/.config/labwc/autostart >/dev/null
            $SUDO chmod +x /home/$TARGET_USER/.config/labwc/autostart
            ;;
        sway)
            echo "Configuring sway as the window manager..."
            $SUDO mkdir -p /home/$TARGET_USER/.config/sway
            $SUDO cp /etc/sway/config /home/$TARGET_USER/.config/sway/config
            echo -e "exec cd ~/iteslive && exec ./iwatcher" | $SUDO tee -a /home/$TARGET_USER/.config/sway/config >/dev/null
            ;;
        *)
            echo "Unsupported window manager '$WINDOW_MANAGER'. Exiting."
            exit 1
            ;;
    esac

    # Fix permissions
    $SUDO chown -R $TARGET_USER:$TARGET_USER /home/$TARGET_USER/.config
}

# =============================================================================
# STEP: apparmor
# =============================================================================
step_apparmor() {
    if [ -d /etc/apparmor.d ] && command -v apparmor_parser &>/dev/null; then
        echo "Configuring AppArmor for $TARGET_USER user namespaces"

        $SUDO tee /etc/apparmor.d/home.$TARGET_USER.$TARGET_USER >/dev/null <<EOF
abi <abi/4.0>,
include <tunables/global>

profile $TARGET_USER /home/$TARGET_USER/iteslive/** flags=(unconfined) {
userns,
}
EOF

        $SUDO apparmor_parser -r /etc/apparmor.d/home.$TARGET_USER.$TARGET_USER || true
    else
        echo "AppArmor not present on this system; skipping userns profile."
    fi
}

# =============================================================================
# Main dispatcher
# =============================================================================

run_step() {
    local id="$1"
    local desc="$2"
    local fn="$3"
    if confirm_step "$id" "$desc"; then
        "$fn"
    fi
}

echo "iClient installer starting."
if [ "$UNATTENDED" -eq 1 ]; then echo "Mode: UNATTENDED"; else echo "Mode: interactive"; fi
[ -n "$SKIP_LIST" ] && echo "Skipping steps: $SKIP_LIST"
[ -n "$ONLY_LIST" ] && echo "Running ONLY:   $ONLY_LIST"

run_step "user" \
    "Create the '$TARGET_USER' system user (login shell /bin/bash, with home directory)." \
    step_user

run_step "packages" \
    "Install base packages: unzip, window manager (labwc or sway), grim, libnss3, wlr-randr, ALSA, PulseAudio, terminal emulator, polkitd." \
    step_packages

run_step "gpu" \
    "Detect installed GPU(s) via lspci and install matching drivers (Intel / AMD / NVIDIA / virtualized)." \
    step_gpu

run_step "download" \
    "Download the iClient binaries from files-cac.itesmedia.tv and install them into /home/$TARGET_USER/iteslive/." \
    step_download

run_step "displaymanager" \
    "Configure the active display manager (LightDM / GDM3 / SDDM) to auto-login the '$TARGET_USER' user. Installs LightDM if none is present." \
    step_displaymanager

run_step "windowmanager" \
    "Configure the chosen window manager (labwc or sway) to autostart ~/iteslive/iwatcher on login." \
    step_windowmanager

run_step "apparmor" \
    "Install an AppArmor profile allowing iteslive's binaries to create user namespaces (required on Ubuntu 23.10+/24.04+ for Chromium/Electron sandboxing)." \
    step_apparmor

echo

# Run only on Raspberry Pi hardware
if grep -qi "Raspberry Pi" /proc/device-tree/model 2>/dev/null; then
    echo "Raspberry Pi detected"

    # Disable autologin (requires Raspberry Pi OS and raspi-config)
    if command -v raspi-config >/dev/null 2>&1; then
        sudo raspi-config nonint do_boot_behaviour B2
        echo "Autologin disabled"
    else
        echo "raspi-config not found"
    fi
else
    echo "Not a Raspberry Pi, skipping"
fi

echo "iClient installation complete. Reboot to start the iteslive session."