#!/usr/bin/env bash
#
# battlenet-umu-setup.sh
#
# Installs the Battle.net launcher on Arch Linux via umu-launcher.
# Idempotent: safe to re-run. Skips steps that are already done.
#
# Usage:
#   ./battlenet-umu-setup.sh                                # interactive
#   ./battlenet-umu-setup.sh --yes                          # non-interactive, assume yes
#   ./battlenet-umu-setup.sh --reinstall                    # wipe prefix and reinstall
#   ./battlenet-umu-setup.sh --proton-version=GE-Proton9-20 # pin a specific Proton version
#   ./battlenet-umu-setup.sh --gpu=nvidia|amd|intel         # override GPU detection
#   ./battlenet-umu-setup.sh --prefix-dir=PATH              # override prefix location
#   ./battlenet-umu-setup.sh --skip-packages                # skip prerequisite package check
# END_HELP
#
# Does NOT run itself as root. It will call sudo only for pacman.

set -euo pipefail

# ---------- config ----------
PREFIX_DIR="${BATTLENET_PREFIX:-$HOME/Games/battlenet-umu}"
GAMEID="umu-battlenet"
PROTON_VERSION="GE-Proton"   # default; overridden by --proton-version or PROTONPATH env var
INSTALLER_URL="https://downloader.battle.net/download/getInstaller?os=win&installer=Battle.net-Setup.exe"
INSTALLER_PATH="$HOME/Downloads/Battle.net-Setup.exe"
LAUNCHER_EXE_REL="drive_c/Program Files (x86)/Battle.net/Battle.net Launcher.exe"
BIN_DIR="$HOME/.local/bin"
DESKTOP_DIR="$HOME/.local/share/applications"
ICON_DIR="$HOME/.local/share/icons/hicolor/256x256/apps"
LAUNCH_SCRIPT="$BIN_DIR/battlenet"
KILL_SCRIPT="$BIN_DIR/battlenetkill"
DESKTOP_FILE="$DESKTOP_DIR/battlenet.desktop"

ASSUME_YES=0
REINSTALL=0
SKIP_PACKAGES=0
GPU_TYPE=""   # auto-detected; override with --gpu=nvidia|amd|intel

# Honour PROTONPATH env var as a default, but let --proton-version override it
[[ -n "${PROTONPATH:-}" ]] && PROTON_VERSION="$PROTONPATH"

# ---------- helpers ----------
c_reset=$'\033[0m'; c_blue=$'\033[1;34m'; c_green=$'\033[1;32m'
c_yellow=$'\033[1;33m'; c_red=$'\033[1;31m'; c_dim=$'\033[2m'

log()  { printf '%s==>%s %s\n' "$c_blue"   "$c_reset" "$*"; }
ok()   { printf '%s✓%s %s\n'   "$c_green"  "$c_reset" "$*"; }
warn() { printf '%s!%s %s\n'   "$c_yellow" "$c_reset" "$*"; }
err()  { printf '%s✗%s %s\n'   "$c_red"    "$c_reset" "$*" >&2; }
skip() { printf '  %s(skip) %s%s\n' "$c_dim" "$*" "$c_reset"; }

confirm() {
    # confirm "Question?" -> returns 0 for yes, 1 for no
    if [[ $ASSUME_YES -eq 1 ]]; then return 0; fi
    local prompt="$1 [Y/n] "
    local reply
    read -r -p "$prompt" reply
    [[ -z "$reply" || "$reply" =~ ^[Yy] ]]
}

die() { err "$*"; exit 1; }

# ---------- arg parsing ----------
for arg in "$@"; do
    case "$arg" in
        -y|--yes)           ASSUME_YES=1 ;;
        --reinstall)        REINSTALL=1 ;;
        --skip-packages)    SKIP_PACKAGES=1 ;;
        --proton-version=*) PROTON_VERSION="${arg#*=}" ;;
        --prefix-dir=*)     PREFIX_DIR="${arg#*=}" ;;
        --gpu=nvidia|--gpu=amd|--gpu=intel)
                            GPU_TYPE="${arg#*=}" ;;
        -h|--help)
            sed -n '/^# Usage:/,/^# END_HELP/p' "$0" | sed 's/^# \{0,1\}//'
            exit 0
            ;;
        *) die "Unknown argument: $arg" ;;
    esac
done

# Expose PROTONPATH for umu-run
PROTONPATH="$PROTON_VERSION"
export PROTONPATH

# ---------- preflight ----------
log "Preflight checks"

[[ $EUID -ne 0 ]] || die "Do not run this script as root. It will sudo when it needs to."

command -v pacman >/dev/null || die "pacman not found — this script is for Arch Linux."
command -v curl   >/dev/null || die "curl not installed. Run: sudo pacman -S curl"

# multilib check
if ! pacman -Sl multilib >/dev/null 2>&1; then
    err "[multilib] repository is not enabled in /etc/pacman.conf."
    err "Edit /etc/pacman.conf, uncomment the [multilib] section, then run:"
    err "  sudo pacman -Syu"
    exit 1
fi
ok "[multilib] is enabled"

# ---------- detect GPU ----------
log "Detecting GPU"

if [[ -z "$GPU_TYPE" ]]; then
    if command -v lspci >/dev/null 2>&1; then
        if lspci | grep -qi "NVIDIA"; then
            GPU_TYPE="nvidia"
        elif lspci | grep -qi "AMD\|Radeon"; then
            GPU_TYPE="amd"
        else
            GPU_TYPE="intel"
        fi
    else
        warn "lspci not found (install pciutils for auto-detection). Defaulting to intel."
        warn "Override with --gpu=nvidia|amd|intel if that is wrong."
        GPU_TYPE="intel"
    fi
fi
ok "GPU type: $GPU_TYPE (override with --gpu=nvidia|amd|intel)"

# ---------- install prerequisites ----------
if [[ $SKIP_PACKAGES -eq 1 ]]; then
    skip "Package check skipped (--skip-packages)"
else
    log "Checking 32-bit prerequisite libraries"

    COMMON_PKGS=(
        giflib lib32-giflib
        libpng lib32-libpng
        libldap lib32-libldap
        gnutls lib32-gnutls
        mpg123 lib32-mpg123
        openal lib32-openal
        v4l-utils lib32-v4l-utils
        libpulse lib32-libpulse
        alsa-plugins lib32-alsa-plugins
        alsa-lib lib32-alsa-lib
        libjpeg-turbo lib32-libjpeg-turbo
        sqlite lib32-sqlite
        libxcomposite lib32-libxcomposite
        libxinerama lib32-libxinerama
        ncurses lib32-ncurses
        opencl-icd-loader lib32-opencl-icd-loader
        libxslt lib32-libxslt
        libva lib32-libva
        gtk3 lib32-gtk3
        gst-plugins-base-libs lib32-gst-plugins-base-libs
        vulkan-icd-loader lib32-vulkan-icd-loader
        cups samba
    )

    GPU_PKGS=()
    case "$GPU_TYPE" in
        nvidia) GPU_PKGS=(lib32-nvidia-utils nvidia-utils) ;;
        amd)    GPU_PKGS=(lib32-mesa lib32-vulkan-radeon vulkan-radeon mesa) ;;
        intel)  GPU_PKGS=(lib32-mesa lib32-vulkan-intel vulkan-intel mesa) ;;
    esac

    ALL_PKGS=("${COMMON_PKGS[@]}" "${GPU_PKGS[@]}")

    MISSING_PKGS=()
    for pkg in "${ALL_PKGS[@]}"; do
        pacman -Qi "$pkg" >/dev/null 2>&1 || MISSING_PKGS+=("$pkg")
    done

    if [[ ${#MISSING_PKGS[@]} -eq 0 ]]; then
        skip "All prerequisite packages already installed"
    else
        warn "Missing packages: ${MISSING_PKGS[*]}"
        confirm "Install missing prerequisite packages?" || die "Aborted."
        sudo pacman -S --needed --noconfirm "${MISSING_PKGS[@]}"
        ok "Prerequisite packages installed"
    fi
fi

# ---------- install umu-launcher ----------
log "Installing umu-launcher"

if pacman -Qi umu-launcher >/dev/null 2>&1; then
    skip "umu-launcher already installed"
else
    confirm "Install umu-launcher via pacman?" || die "Aborted."
    sudo pacman -S --needed --noconfirm umu-launcher
    ok "umu-launcher installed"
fi

command -v umu-run >/dev/null || die "umu-run not on PATH after install — something went wrong."

# ---------- download installer ----------
log "Fetching Battle.net installer"

mkdir -p "$(dirname "$INSTALLER_PATH")"

_download_installer() {
    curl -L --fail --progress-bar -o "$INSTALLER_PATH" "$INSTALLER_URL"
    # Sanity check: Battle.net-Setup.exe is always several MB; <100 KB means a redirect page or error.
    local size
    size=$(stat -c%s "$INSTALLER_PATH" 2>/dev/null || echo 0)
    if (( size < 102400 )); then
        rm -f "$INSTALLER_PATH"
        die "Downloaded file is suspiciously small (${size} bytes) — download may have failed."
    fi
}

if [[ -f "$INSTALLER_PATH" ]]; then
    local_size=$(stat -c%s "$INSTALLER_PATH" 2>/dev/null || echo 0)
    if (( local_size < 102400 )); then
        warn "Existing installer looks corrupt (${local_size} bytes). Re-downloading."
        _download_installer
        ok "Installer downloaded to $INSTALLER_PATH"
    else
        skip "Installer already at $INSTALLER_PATH ($(( local_size / 1024 / 1024 )) MB)"
        # When running non-interactively (--yes), keep the existing file; only re-download on explicit prompt.
        if [[ $ASSUME_YES -eq 0 ]] && confirm "Re-download to get the latest?"; then
            _download_installer
            ok "Installer re-downloaded"
        fi
    fi
else
    _download_installer
    ok "Installer saved to $INSTALLER_PATH"
fi

# ---------- prefix ----------
log "Preparing prefix at $PREFIX_DIR"

if [[ -d "$PREFIX_DIR" ]]; then
    if [[ $REINSTALL -eq 1 ]]; then
        warn "Reinstall requested — removing existing prefix"
        if confirm "Really delete $PREFIX_DIR ?"; then
            rm -rf "$PREFIX_DIR"
            ok "Prefix removed"
        else
            die "Aborted."
        fi
    else
        skip "Prefix already exists (use --reinstall to wipe)"
    fi
fi

mkdir -p "$PREFIX_DIR"

# ---------- run installer ----------
LAUNCHER_PATH="$PREFIX_DIR/$LAUNCHER_EXE_REL"

if [[ -f "$LAUNCHER_PATH" ]]; then
    skip "Battle.net Launcher.exe already present in prefix"
else
    log "Running Battle.net installer through umu (this takes a few minutes)"
    warn "The Battle.net installer window will appear. Accept the defaults"
    warn "and let it finish. The installer may auto-close — that's fine."
    confirm "Continue?" || die "Aborted."

    WINEPREFIX="$PREFIX_DIR" \
    GAMEID="$GAMEID" \
    PROTONPATH="$PROTONPATH" \
        umu-run "$INSTALLER_PATH" || {
            warn "umu-run exited non-zero. That's often normal for the Battle.net installer."
        }

    # The installer sometimes exits before the launcher EXE is flushed to disk.
    # Poll for up to 30 s rather than assuming a fixed delay is enough.
    i=0
    while [[ ! -f "$LAUNCHER_PATH" && $i -lt 30 ]]; do
        sleep 1
        (( i++ )) || true
    done

    if [[ ! -f "$LAUNCHER_PATH" ]]; then
        err "Battle.net Launcher.exe not found at expected path:"
        err "  $LAUNCHER_PATH"
        err "The installer may not have completed. Re-run with --reinstall."
        exit 1
    fi
    ok "Battle.net installed into prefix"
fi

# ---------- icon ----------
log "Installing icon"

ICON_SRC="$PREFIX_DIR/drive_c/Program Files (x86)/Battle.net/Battle.net Launcher.exe"
ICON_PNG="$ICON_DIR/battlenet.png"

_install_icon() {
    local tmpdir
    tmpdir=$(mktemp -d)
    trap 'rm -rf "$tmpdir"' RETURN

    # Extract the largest icon group from the EXE, then convert to PNG.
    # wrestool and icotool come from the 'icoutils' package.
    if command -v wrestool >/dev/null && command -v icotool >/dev/null; then
        if wrestool -x -t 14 "$ICON_SRC" -o "$tmpdir" 2>/dev/null; then
            local ico
            ico=$(find "$tmpdir" -name "*.ico" | head -1)
            if [[ -n "$ico" ]]; then
                mkdir -p "$ICON_DIR"
                icotool -x --index=1 -o "$ICON_DIR" "$ico" 2>/dev/null \
                    && mv "$ICON_DIR/"*_1_*.png "$ICON_PNG" 2>/dev/null \
                    && ok "Icon installed to $ICON_PNG" \
                    && return 0
            fi
        fi
        warn "Icon extraction failed — Battle.net may not be installed yet, or the EXE layout changed."
    else
        warn "icoutils (wrestool/icotool) not installed. Skipping icon extraction."
        warn "Install with: sudo pacman -S icoutils"
    fi
    return 1
}

if [[ -f "$ICON_PNG" ]]; then
    skip "Icon already at $ICON_PNG"
elif [[ -f "$ICON_SRC" ]]; then
    _install_icon || true
else
    skip "Launcher EXE not found; skipping icon (will retry on next run)"
fi

# Refresh icon cache best-effort
if command -v gtk-update-icon-cache >/dev/null; then
    gtk-update-icon-cache -f -t "$HOME/.local/share/icons/hicolor" 2>/dev/null || true
fi

# ---------- launch scripts ----------
log "Installing launch scripts to $BIN_DIR"

mkdir -p "$BIN_DIR"

cat > "$LAUNCH_SCRIPT" <<EOF
#!/bin/sh
# Auto-generated by battlenet-umu-setup.sh
# Proton version: $PROTON_VERSION
export WINEPREFIX="$PREFIX_DIR"
export GAMEID="$GAMEID"
export PROTONPATH="$PROTON_VERSION"
exec umu-run "\$WINEPREFIX/$LAUNCHER_EXE_REL" "\$@"
EOF
chmod +x "$LAUNCH_SCRIPT"
ok "Wrote $LAUNCH_SCRIPT"

cat > "$KILL_SCRIPT" <<'KILL_EOF'
#!/bin/sh
# Auto-generated by battlenet-umu-setup.sh
# Gracefully stops Battle.net, then force-kills after a timeout.

_graceful_kill() {
    local pattern="$1"
    # Try SIGTERM first
    pkill -15 -f "$pattern" 2>/dev/null || return 0
    local i=0
    while pkill -0 -f "$pattern" 2>/dev/null && [ $i -lt 5 ]; do
        sleep 1
        i=$(( i + 1 ))
    done
    # Force-kill anything still alive
    pkill -9 -f "$pattern" 2>/dev/null || true
}

_graceful_kill 'Battle\.net'
_graceful_kill 'Agent\.exe'
_graceful_kill 'Blizzard'

# wineserver may not be on PATH in umu environments; try both.
if command -v wineserver >/dev/null; then
    wineserver -k 2>/dev/null || true
fi
KILL_EOF
chmod +x "$KILL_SCRIPT"
ok "Wrote $KILL_SCRIPT"

# PATH check — show shell-appropriate instructions
case ":$PATH:" in
    *":$BIN_DIR:"*) ok "$BIN_DIR is on PATH" ;;
    *)
        warn "$BIN_DIR is not on your PATH."
        _shell_name="$(basename "${SHELL:-bash}")"
        case "$_shell_name" in
            fish)
                warn "Add it permanently with:"
                warn "  fish_add_path \$HOME/.local/bin"
                ;;
            zsh)
                warn "Add this to ~/.zshrc:"
                warn "  export PATH=\"\$HOME/.local/bin:\$PATH\""
                ;;
            *)
                warn "Add this to ~/.bashrc (or ~/.profile for login shells):"
                warn "  export PATH=\"\$HOME/.local/bin:\$PATH\""
                ;;
        esac
        ;;
esac

# ---------- desktop entry ----------
log "Installing desktop entry"

# Use extracted icon if available, fall back to generic game icon
ICON_NAME="battlenet"
[[ -f "$ICON_PNG" ]] || ICON_NAME="applications-games"

mkdir -p "$DESKTOP_DIR"
cat > "$DESKTOP_FILE" <<EOF
[Desktop Entry]
Name=Battle.net
Comment=Blizzard launcher (umu/Proton)
Exec=$LAUNCH_SCRIPT
Icon=$ICON_NAME
Type=Application
Categories=Game;
StartupWMClass=battle.net.exe
EOF
ok "Wrote $DESKTOP_FILE"

if command -v update-desktop-database >/dev/null; then
    update-desktop-database "$DESKTOP_DIR" 2>/dev/null || true
fi

# ---------- done ----------
echo
ok "Setup complete."
echo
echo "  Prefix:       $PREFIX_DIR"
echo "  Proton:       $PROTON_VERSION"
echo "  Launcher:     $LAUNCH_SCRIPT"
echo "  Kill:         $KILL_SCRIPT"
echo "  Desktop:      $DESKTOP_FILE"
echo
echo "Run 'battlenet' to start it, or launch from your app menu."
echo "If it hangs on 'Update Agent went to sleep', run: battlenetkill && battlenet"