#!/bin/bash
#
# Seamless Hyprland <-> Gamescope Session Switcher for Arch Linux
# Updated: Walker, Infinite Loop, env.d Config, Color Reset, IWD Backend
#

set -e

# --- Pre-flight Checks and Setup ---
C_GREEN='\033[0;32m'
C_YELLOW='\033[1;33m'
C_RED='\033[0;31m'
C_BLUE='\033[0;34m'
C_NC='\033[0m' # No Color

if [ "$EUID" -eq 0 ]; then
  echo -e "${C_RED}Error: Do not run this script as root! It will use 'sudo' as needed.${C_NC}"
  exit 1
fi

USER_NAME=$(whoami)
USER_HOME=$(eval echo "~$USER_NAME")

# --- Banner ---
echo -e "${C_BLUE}===================================================================${C_NC}"
echo -e "${C_BLUE} Hyprland <-> Gamescope Session Switcher Setup for: ${C_YELLOW}$USER_NAME${C_NC}"
echo -e "${C_BLUE}===================================================================${C_NC}\n"

# --- USER CHOICE FOR AUTOLOGIN ---
echo -e "${C_YELLOW}Configuration Choice:${C_NC}"
echo "Do you want to enable automatic login? [y/N]: "
read AUTOLOGIN_CHOICE

# --- Script Variables ---
HYPR_CONF="$USER_HOME/.config/hypr/hyprland.conf"
SWITCH_SCRIPT_PATH="$USER_HOME/.local/bin/switch-session.sh"
XSESSION_PATH="$USER_HOME/.xsession"
SERVICE_OVERRIDE_DIR="/etc/systemd/user/gamescope-session-plus@.service.d"
SERVICE_OVERRIDE_FILE="$SERVICE_OVERRIDE_DIR/override.conf"
GS_ENV_DIR="$USER_HOME/.config/environment.d"
GS_ENV_FILE="$GS_ENV_DIR/gamescope-session-plus.conf"

# UPDATED: Removed wofi, added walker to AUR packages
OFFICIAL_PACKAGES=( "hyprland" "sddm" "uwsm" "networkmanager" )
AUR_PACKAGES=( "gamescope-git" "gamescope-session-git" "steam" "gamescope-session-steam-git" "mangohud" "walker" )

#=======================================================
# STEP 1: DEPENDENCY INSTALLATION
#=======================================================
echo -e "${C_BLUE}==> Installing Dependencies...${C_NC}"
if command -v yay &> /dev/null; then AUR_HELPER="yay"; elif command -v paru &> /dev/null; then AUR_HELPER="paru"; else
    echo -e "${C_RED}Error: No AUR helper found (yay or paru). Please install one.${C_NC}"; exit 1; fi
echo -e "${C_GREEN}Found AUR helper: ${AUR_HELPER}${C_NC}"
sudo pacman -Syudd --needed "${OFFICIAL_PACKAGES[@]}" --noconfirm
$AUR_HELPER -Sdd --needed "${AUR_PACKAGES[@]}" --noconfirm

#=======================================================
# STEP 2: CLEAN UP OLD CONFIGURATIONS
#=======================================================
echo -e "${C_BLUE}==> Cleaning Up Old Configurations...${C_NC}"
sudo rm -rf "$SERVICE_OVERRIDE_DIR"
rm -f "$XSESSION_PATH"
rm -f "$SWITCH_SCRIPT_PATH"

#=======================================================
# STEP 3: APPLY SYSTEMD FIXES
#=======================================================
echo -e "${C_BLUE}==> Applying Core Fix for Wayland Socket Errors...${C_NC}"
sudo mkdir -p "$SERVICE_OVERRIDE_DIR"
sudo tee "$SERVICE_OVERRIDE_FILE" > /dev/null <<'EOF'
[Service]
ExecStart=
ExecStart=/usr/bin/env -u WAYLAND_DISPLAY /usr/share/gamescope-session-plus/gamescope-session-plus %i
EOF
systemctl --user daemon-reload

#=======================================================
# STEP 4: CONFIGURE GAMESCOPE ENVIRONMENT (MONITOR FIX)
#=======================================================
echo -e "${C_BLUE}==> Configuring Gamescope Monitor Settings (DP-1)...${C_NC}"
echo "Writing to: $GS_ENV_FILE"
mkdir -p "$GS_ENV_DIR"
cat > "$GS_ENV_FILE" <<EOF
# Force Output to DP-1 (Fixes display issue)
OUTPUT_CONNECTOR=DP-1

# Force specific arguments for gamescope-session-plus
GAMESCOPE_ARGS="-O DP-1 --steam -e"
EOF

#=======================================================
# STEP 5: CONFIGURE NETWORK MANAGER (IWD BACKEND)
#=======================================================
echo -e "${C_BLUE}==> Configuring NetworkManager for SteamOS compatibility...${C_NC}"
# Create config to use iwd backend (prevents conflicts)
sudo mkdir -p /etc/NetworkManager/conf.d
sudo tee /etc/NetworkManager/conf.d/wifi_backend.conf > /dev/null <<EOF
[device]
wifi.backend=iwd
EOF

echo "Enabling NetworkManager..."
sudo systemctl enable NetworkManager

#=======================================================
# STEP 6: INSTALL SESSION SWITCHING LOGIC
#=======================================================
echo -e "${C_BLUE}==> Installing Session Switching Workflow...${C_NC}"

# --- SDDM AUTOLOGIN (CONDITIONAL) ---
if [[ "$AUTOLOGIN_CHOICE" =~ ^[Yy]$ ]]; then
    echo "--> Configuring SDDM for autologin..."
    sudo tee /etc/sddm.conf > /dev/null <<EOF
[Autologin]
User=${USER_NAME}
Session=switcher.desktop
Relogin=true
EOF
    AUTOLOGIN_ENABLED=true
else
    AUTOLOGIN_ENABLED=false
fi

# --- WAYLAND SESSION DESKTOP ENTRY ---
sudo tee /usr/share/wayland-sessions/switcher.desktop > /dev/null <<EOF
[Desktop Entry]
Name=Auto Session Switcher
Exec=${XSESSION_PATH}
Type=Application
EOF

#=======================================================
# STEP 7: CREATE XSESSION LAUNCH SCRIPT (THE LOOP)
#=======================================================
echo -e "${C_BLUE}==> Creating Session Launch Script (Infinite Loop)...${C_NC}"

cat > "$XSESSION_PATH" <<'EOS'
#!/bin/bash

# Redirect logs for debugging
exec > >(tee -a "$HOME/.xsession.log") 2>&1

# Function to handle Wayland socket waiting
wait_for_wayland() {
    local socket_path="$XDG_RUNTIME_DIR/wayland-1"
    local max_attempts=60
    local attempt=1
    
    echo "Waiting for Wayland socket at $socket_path"
    while [ $attempt -le $max_attempts ]; do
        if [ -e "$socket_path" ] && [ -S "$socket_path" ]; then
            return 0
        fi
        sleep 0.1
        ((attempt++))
    done
    return 1
}

# --- THE ENDLESS LOOP ---
while true; do
    echo "--- New Session Cycle Starting ---"
    
    # 1. READ THE NEXT SESSION
    if [ -f "$HOME/.next-session" ]; then
        SESSION=$(cat "$HOME/.next-session")
        rm -f "$HOME/.next-session"
    else
        SESSION="hyprland"
    fi

    # 2. UPDATE ENVIRONMENT
    dbus-update-activation-environment --systemd DISPLAY WAYLAND_DISPLAY

    # 3. LAUNCH THE SESSION (WITHOUT EXEC)
    if [[ "$SESSION" == *"gamescope-session-steam"* ]]; then
        echo "Starting Gamescope session..."
        
        # Check for the executable
        if command -v gamescope-session-plus &> /dev/null; then
            # Run blocking, do not exec
            export GAMESCOPECMD="/usr/bin/gamescope -O DP-1 --steam -e --"
            gamescope-session-plus steam
            unset GAMESCOPECMD
        else
            echo "Error: gamescope-session-plus not found, falling back to Hyprland"
            SESSION="hyprland"
        fi
    fi

    # Note: We use a separate 'if' here to catch the fallback above
    if [[ "$SESSION" != *"gamescope-session-steam"* ]]; then
        echo "Starting Hyprland session with UWSM..."
        
        # Run blocking, do not exec
        uwsm start hyprland-uwsm.desktop
    fi

    echo "Session exited. Looping..."
    
    # 4. CRASH PROTECTION
    sleep 2
done

EOS
chmod +x "$XSESSION_PATH"

#=======================================================
# STEP 8: CREATE SWITCHING SCRIPT (WALKER EDITION)
#=======================================================
echo -e "${C_BLUE}==> Creating Session Switching Helper...${C_NC}"

mkdir -p "$(dirname "$SWITCH_SCRIPT_PATH")"
cat > "$SWITCH_SCRIPT_PATH" <<'EOSWITCH'

#!/bin/bash

export WAYLAND_DISPLAY="${WAYLAND_DISPLAY:-wayland-1}"
export XDG_RUNTIME_DIR="${XDG_RUNTIME_DIR:-/run/user/$(id -u)}"

choice=$(printf "🎮 SteamOS" | walker --dmenu \
    --placeholder "Switch Mode:" \
    # --width 400 \
    # --height 250 \
    # --cache-file /dev/null \
    # --style "$HOME/.config/wofi/style.css" \
    # --hide-scroll \
    # --no-actions \
    # --lines 2)
  )

case "$choice" in
    "🎮 SteamOS")
        # Validate we can write the session file
        if echo "gamescope-session-steam.desktop" > "$HOME/.next-session" 2>/dev/null; then
            notify-send "Session Switcher" "Switching to SteamOS..." -t 2000
            echo "Session file written successfully"
            pkill hyprsunset
            sleep 1
        else
            notify-send "Session Switcher" "Error: Failed to write session file" -t 3000
            echo "Error: Cannot write to $HOME/.next-session"
            exit 1
        fi
        ;;
    *)
        echo "No valid choice made"
        exit 1
        ;;
esac


# Improved session cleanup with better reliability
cleanup_session() {
    local max_wait=5
    local wait_count=0
    
    # Detect current session more reliably
    if pgrep -x "gamescope" > /dev/null || pgrep -f "gamescope-session" > /dev/null; then
        echo "Detected Gamescope session, shutting down..."
        
        # Try graceful shutdown first
        systemctl --user stop gamescope-session-plus@steam 2>/dev/null || true
        sleep 1
        
        # Force kill if still running
        if pgrep -x "gamescope" > /dev/null; then
            pkill -TERM gamescope
            sleep 0.5
            pkill -KILL gamescope 2>/dev/null || true
        fi
        
        # Wait for processes to fully exit
        while pgrep -x "gamescope" > /dev/null && [ $wait_count -lt $max_wait ]; do
            sleep 1
            ((wait_count++))
        done
        
    elif [ -n "$HYPRLAND_INSTANCE_SIGNATURE" ] || pgrep -x "Hyprland" > /dev/null; then
        echo "Detected Hyprland session, shutting down..."
        
        # Try hyprctl first if available
        if command -v hyprctl &> /dev/null && [ -n "$HYPRLAND_INSTANCE_SIGNATURE" ]; then
            hyprctl dispatch exit 2>/dev/null || true
        else
            # Fallback to pkill
            pkill -TERM Hyprland 2>/dev/null || true
        fi
        
        sleep 1
        
        # Force kill if still running
        if pgrep -x "Hyprland" > /dev/null; then
            pkill -KILL Hyprland 2>/dev/null || true
        fi
        
        # Wait for processes to fully exit
        while pgrep -x "Hyprland" > /dev/null && [ $wait_count -lt $max_wait ]; do
            sleep 1
            ((wait_count++))
        done
        
    else
        echo "No specific session detected, attempting general cleanup..."
        # Fallback cleanup
        pkill -u $USER Hyprland 2>/dev/null || true
        pkill -u $USER gamescope 2>/dev/null || true
        systemctl --user stop gamescope-session-plus@steam 2>/dev/null || true
        sleep 1
    fi
    
    # Clean up any remaining compositor processes
    pkill -u $USER -f "gamescope-session" 2>/dev/null || true
    pkill -u $USER -f "steam" 2>/dev/null || true
    
    # Give time for cleanup
    sleep 0.5
}

# Execute cleanup
cleanup_session
EOSWITCH
chmod +x "$SWITCH_SCRIPT_PATH"

#=======================================================
# STEP 9: CONFIGURE HYPRLAND KEYBINDING
#=======================================================
echo -e "${C_BLUE}==> Configuring Hyprland Keybinding...${C_NC}"

# Create hyprland config directory if it doesn't exist
mkdir -p "$(dirname "$HYPR_CONF")"

if [ ! -f "$HYPR_CONF" ]; then
    echo "# Hyprland Configuration" > "$HYPR_CONF"
    echo "" >> "$HYPR_CONF"
fi

sed -i '/# Session Switcher Keybinding/d' "$HYPR_CONF"
sed -i '/bind = SUPER, F12, exec,.*switch-session.sh/d' "$HYPR_CONF"

cat >> "$HYPR_CONF" <<EOF

# Session Switcher Keybinding
bind = SUPER, F12, exec, $SWITCH_SCRIPT_PATH
EOF

echo -e "${C_GREEN}Added SUPER+F12 keybinding to $HYPR_CONF${C_NC}"


#=======================================================
# FINALIZATION
#=======================================================
echo -e "${C_GREEN}✅ Setup Complete! ✅${C_NC}\n"

echo -e "${C_YELLOW}Important Fixes Applied:${C_NC}"
echo "1. REPLACED 'wofi' with 'walker'."
echo "2. CONFIGURED NetworkManager to use 'iwd' backend."
echo "3. UPDATED '.xsession' to use an infinite loop (prevents logout)."
echo "4. CREATED '~/.config/environment.d/gamescope-session-plus.conf' to force DP-1."
echo "5. MODIFIED Switch script to kill 'hyprsunset' before switching."

echo -e "\n${C_YELLOW}A REBOOT IS REQUIRED to activate NetworkManager and the new Session structure.${C_NC}\n"

if [ "$AUTOLOGIN_ENABLED" = true ]; then
    echo -e "-> Autologin is enabled. System will boot directly into Hyprland."
else
    echo -e "-> Autologin is disabled. Choose 'Auto Session Switcher' at login."
fi