#!/bin/bash
#
# Seamless Hyprland <-> Gamescope Session Switcher for Arch Linux
# Updated: Clean Config, User Inputs for HDR/VRR
#

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=$HOME

# --- 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 INPUTS ---
echo -e "${C_YELLOW}Configuration Choice:${C_NC}"
echo "Do you want to enable automatic login? [y/N]: "
read AUTOLOGIN_CHOICE

echo -e "\n${C_YELLOW}Monitor Configuration:${C_NC}"

# 1. Output Port
echo "Enter your Output Connector (default: DP-1):"
read USER_OUTPUT
USER_OUTPUT=${USER_OUTPUT:-DP-1}

# 2. Width
echo "Enter Screen Width (default: 3440):"
read USER_WIDTH
USER_WIDTH=${USER_WIDTH:-3440}

# 3. Height
echo "Enter Screen Height (default: 1440):"
read USER_HEIGHT
USER_HEIGHT=${USER_HEIGHT:-1440}

# 4. Refresh Rate
echo "Enter Refresh Rate (default: 180):"
read USER_REFRESH
USER_REFRESH=${USER_REFRESH:-180}

# 5. VRR (Adaptive Sync)
echo "Enable VRR / Adaptive Sync? [Y/n] (Default: Yes)"
read USER_VRR
if [[ "$USER_VRR" =~ ^[Nn]$ ]]; then
    VRR_FLAG=""
    echo -e "-> VRR Disabled"
else
    VRR_FLAG="--adaptive-sync"
    echo -e "-> VRR Enabled"
fi

# 6. HDR
echo "Enable HDR? [Y/n] (Default: Yes)"
read USER_HDR
if [[ "$USER_HDR" =~ ^[Nn]$ ]]; then
    HDR_FLAG=""
    echo -e "-> HDR Disabled"
else
    HDR_FLAG="--hdr-enabled"
    echo -e "-> HDR Enabled"
fi

echo -e "${C_GREEN}Configured: ${USER_OUTPUT} @ ${USER_WIDTH}x${USER_HEIGHT} (${USER_REFRESH}Hz)${C_NC}\n"


# --- Script Variables ---
HYPR_CONF="$USER_HOME/.config/hypr/bindings.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"

OFFICIAL_PACKAGES=( "hyprland" "sddm" "uwsm" "networkmanager" )
AUR_PACKAGES=( "gamescope-git" "gamescope-session-git" "steam" "gamescope-session-steam-git" "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 3.5: GRANT REAL-TIME PRIORITY CAPABILITIES
#=======================================================
echo -e "${C_BLUE}==> Granting CAP_SYS_NICE to Gamescope...${C_NC}"
# This allows gamescope to use --rt (realtime) scheduling without root
GAME_BIN=$(which gamescope)
if [ -f "$GAME_BIN" ]; then
    sudo setcap 'CAP_SYS_NICE=eip' "$GAME_BIN"
    echo -e "${C_GREEN}Capability set on $GAME_BIN${C_NC}"
else
    echo -e "${C_RED}Error: Gamescope binary not found!${C_NC}"
fi

#=======================================================
# STEP 4: CONFIGURE GAMESCOPE ENVIRONMENT (CLEAN)
#=======================================================
echo -e "${C_BLUE}==> Configuring Gamescope Command...${C_NC}"

if [ -f "$GS_ENV_DIR" ]; then rm "$GS_ENV_DIR"; fi
mkdir -p "$GS_ENV_DIR"

echo "Writing to: $GS_ENV_FILE"

# NOTE: We do NOT use quotes around EOF here.
# This allows the variables ($USER_OUTPUT, $VRR_FLAG, etc.) to expanded
# RIGHT NOW, so the resulting file contains only hardcoded values.

cat > "$GS_ENV_FILE" <<EOF
# --- GAMESCOPE CONFIGURATION ---
# Generated by setup script on $(date)
#
# Flags Explanation:
# -O : Output Connector
# -r : Refresh Rate
# -W/-H : Resolution the Game sees
# -w/-h : Resolution sent to Monitor
# -e : Steam Integration
# -f : Force Fullscreen
# --xwayland-count 2 : Vital for Steam Deck UI + Game
# --immediate-flips : Low latency mode
# --rt : Real-time priority

GAMESCOPECMD="/usr/bin/gamescope \\
  -O $USER_OUTPUT \\
  -r $USER_REFRESH \\
  -W $USER_WIDTH -H $USER_HEIGHT \\
  -w $USER_WIDTH -h $USER_HEIGHT \\
  $VRR_FLAG \\
  $HDR_FLAG \\
  --immediate-flips \\
  --rt \\
  -e \\
  --mangoapp \\
  --xwayland-count 2 \\
  --default-touch-mode 4 \\
  --hide-cursor-delay 3000 \\
  --fade-out-duration 200"
EOF

#=======================================================
# STEP 5: CONFIGURE NETWORK MANAGER (IWD BACKEND)
#=======================================================
echo -e "${C_BLUE}==> Configuring NetworkManager for SteamOS compatibility...${C_NC}"
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}"

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

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 + Safe Priority)...${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..."
        
        # --- LOAD USER CONFIG FROM ENVIRONMENT.D ---
        if [ -f "$HOME/.config/environment.d/gamescope-session-plus.conf" ]; then
            echo "Loading gamescope config from environment.d..."
            set -a
            source "$HOME/.config/environment.d/gamescope-session-plus.conf"
            set +a
        fi

        # Check for the executable
        if command -v gamescope-session-plus &> /dev/null; then
            
            # Use the constructed variable from the conf file
            echo "Executing Command: $GAMESCOPECMD"
            
            # Strip the leading binary path to get just the args
            # We strip "/usr/bin/gamescope" OR just "gamescope" to be safe
            ARGS=$(echo "$GAMESCOPECMD" | sed 's|^/usr/bin/gamescope ||' | sed 's|^gamescope ||')
            
            gamescope-session-plus steam -- $ARGS
            
        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)}"

# Walker in dmenu mode
choice=$(printf "🎮 SteamOS" | walker --dmenu --placeholder "Switch Mode")

case "$choice" in
    "🎮 SteamOS")
        notify-send "Session Switcher" "Disabling Night Light..." -t 1000
        pkill hyprsunset
        pkill wlsunset
        pkill gammastep
        sleep 1
        if echo "gamescope-session-steam.desktop" > "$HOME/.next-session" 2>/dev/null; then
            notify-send "Session Switcher" "Switching to SteamOS..." -t 2000
        else
            notify-send "Session Switcher" "Error: Failed to write session file" -t 3000
            exit 1
        fi
        ;;
    *)
        echo "No valid choice made"
        exit 1
        ;;
esac

cleanup_session() {
    local max_wait=5
    local wait_count=0
    
    if pgrep -x "gamescope" > /dev/null || pgrep -f "gamescope-session" > /dev/null; then
        echo "Detected Gamescope session, shutting down..."
        systemctl --user stop gamescope-session-plus@steam 2>/dev/null || true
        sleep 1
        if pgrep -x "gamescope" > /dev/null; then
            pkill -TERM gamescope
            sleep 0.5
            pkill -KILL gamescope 2>/dev/null || true
        fi
    elif [ -n "$HYPRLAND_INSTANCE_SIGNATURE" ] || pgrep -x "Hyprland" > /dev/null; then
        echo "Detected Hyprland session, shutting down..."
        if command -v hyprctl &> /dev/null && [ -n "$HYPRLAND_INSTANCE_SIGNATURE" ]; then
            hyprctl dispatch exit 2>/dev/null || true
        else
            pkill -TERM Hyprland 2>/dev/null || true
        fi
        sleep 1
        if pgrep -x "Hyprland" > /dev/null; then
            pkill -KILL Hyprland 2>/dev/null || true
        fi
    else
        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
    pkill -u $USER -f "gamescope-session" 2>/dev/null || true
    pkill -u $USER -f "steam" 2>/dev/null || true
    sleep 0.5
}
cleanup_session
EOSWITCH
chmod +x "$SWITCH_SCRIPT_PATH"

#=======================================================
# STEP 9: CONFIGURE HYPRLAND KEYBINDING
#=======================================================
echo -e "${C_BLUE}==> Configuring Hyprland Keybinding...${C_NC}"
mkdir -p "$(dirname "$HYPR_CONF")"
if [ ! -f "$HYPR_CONF" ]; then echo "" > "$HYPR_CONF"; fi

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

cat >> "$HYPR_CONF" <<EOF
# Session Switcher Keybinding
bindd = SUPER, F12, Switch to SteamOS, 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}Restored Functionality:${C_NC}"
echo "1. REMOVED systemd-run wrapper (it was causing permissions crashes)."
echo "2. ADDED safe priority boost (renice) for smoother performance."
echo "3. FIXED environment.d directory creation error."
echo -e "\n${C_YELLOW}Run this script, then REBOOT.${C_NC}\n"