#!/bin/bash
#
# Seamless Hyprland <-> Gamescope Session Switcher for Arch Linux
# Updated: 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"
# UPDATED: Correct config path for systemd user environment
GS_ENV_DIR="$USER_HOME/.config/environment.d"
GS_ENV_FILE="$GS_ENV_DIR/gamescope-session-plus.conf"

# UPDATED: Removed lib32-networkmanager
OFFICIAL_PACKAGES=( "hyprland" "wofi" "sddm" "uwsm" "networkmanager" )
AUR_PACKAGES=( "gamescope-git" "gamescope-session-git" "steam" "gamescope-session-steam-git" "mangohud" )

#=======================================================
# 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
# We don't force start here to avoid killing the user's current SSH/Session immediately, 
# but it will be active on reboot.

#=======================================================
# 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..."
        
        # --- LOAD USER CONFIG FROM ENVIRONMENT.D ---
        # Since we are running manually, we must explicitly load the systemd env config
        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
            
            # TRIPLE-FORCE MONITOR CONFIGURATION
            # 1. Environment variable (from sourced file or manual export)
            export GAMESCOPECMD="/usr/bin/gamescope ${GAMESCOPE_ARGS:-"-O DP-1 --steam -e"}"
            export OUTPUT_CONNECTOR="${OUTPUT_CONNECTOR:-"DP-1"}"
            
            # 2. Run blocking, do not exec
            # 3. Pass CLI args just in case
            gamescope-session-plus steam -- -O DP-1
            
            # Cleanup
            unset GAMESCOPECMD
            unset OUTPUT_CONNECTOR
        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 (COLOR RESET FIX)
#=======================================================
echo -e "${C_BLUE}==> Creating Session Switching Helper...${C_NC}"

mkdir -p "$(dirname "$SWITCH_SCRIPT_PATH")"
cat > "$SWITCH_SCRIPT_PATH" <<'EOSWITCH'
#!/bin/bash

# Set environment variables for Wofi to work properly
export WAYLAND_DISPLAY="${WAYLAND_DISPLAY:-wayland-1}"
export XDG_RUNTIME_DIR="${XDG_RUNTIME_DIR:-/run/user/$(id -u)}"

# Create Wofi config directory if it doesn't exist
mkdir -p "$HOME/.config/wofi"

# Create custom Wofi stylesheet for SteamOS appearance
cat > "$HOME/.config/wofi/style.css" <<'EOSTYLE'
window {
    background-color: rgba(30, 30, 46, 0.85);
    border: 1px solid rgba(60, 125, 210, 0.5);
    border-radius: 16px;
    font-family: "Noto Sans", sans-serif;
    font-size: 18px;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.4), inset 0 1px 0 rgba(255, 255, 255, 0.1);
    min-height: 250px;
}
#input {
    background-color: rgba(42, 42, 62, 0.6);
    border: 1px solid rgba(60, 125, 210, 0.3);
    border-radius: 8px;
    color: #cdd6f4;
    font-size: 16px;
    margin: 12px;
    padding: 12px 16px;
}
#inner-box { background-color: transparent; margin: 8px; min-height: 150px; }
#outer-box { background-color: rgba(255, 255, 255, 0.02); border-radius: 12px; margin: 4px; padding: 8px; }
#scroll { background-color: transparent; min-height: 120px; }
#text { color: rgba(205, 214, 244, 0.95); padding: 12px 16px; font-weight: 500; font-size: 18px; }
#entry {
    background-color: rgba(255, 255, 255, 0.05);
    border: 1px solid rgba(255, 255, 255, 0.1);
    border-radius: 10px;
    margin: 8px 12px;
    padding: 12px;
    min-height: 48px;
    transition: all 0.2s ease;
}
#entry:hover { background-color: rgba(255, 255, 255, 0.08); border-color: rgba(60, 125, 210, 0.4); transform: translateX(4px); }
#entry:selected {
    background: linear-gradient(135deg, rgba(60, 125, 210, 0.8) 0%, rgba(37, 99, 235, 0.8) 100%);
    border: 1px solid rgba(60, 125, 210, 0.6);
}
#entry:selected #text { color: #ffffff; font-weight: 600; }
EOSTYLE

# Show session switcher with Wofi
choice=$(printf "🎮 SteamOS" | wofi --dmenu \
    --prompt "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")
        # --- FIX: RESET SCREEN COLOR ---
        notify-send "Session Switcher" "Disabling Night Light..." -t 1000
        pkill hyprsunset
        pkill wlsunset
        pkill gammastep
        sleep 1  # Give GPU time to reset gamma
        # -------------------------------

        # 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"
        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..."
        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
}

# 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. REMOVED 'lib32-networkmanager' (Using 64-bit standard package)."
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