#!/bin/bash

# WiFi Power Save Disable Script
# Supports Arch Linux and Fedora
# Author: Asher Payn :)

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Colour

# Function to print coloured output
print_status() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Check if running as root
check_root() {
    if [[ $EUID -eq 0 ]]; then
        print_error "This script should not be run as root. Please run as a regular user."
        print_status "The script will use sudo when needed."
        exit 1
    fi
}

# Detect the Linux distribution
detect_distro() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        DISTRO=$ID
    else
        print_error "Cannot detect Linux distribution"
        exit 1
    fi
    
    case $DISTRO in
        arch|manjaro|cachyos|artix|endeavouros)
            DISTRO_TYPE="arch"
            print_status "Detected Arch-based distribution: $NAME"
            ;;
        fedora|centos|rhel|rocky|almalinux)
            DISTRO_TYPE="fedora"
            print_status "Detected Red Hat-based distribution: $NAME"
            ;;
        *)
            print_error "Unsupported distribution: $DISTRO"
            print_status "This script supports Arch Linux and Fedora-based distributions only."
            exit 1
            ;;
    esac
}

# Check if NetworkManager is running
check_networkmanager() {
    if systemctl is-active --quiet NetworkManager; then
        print_status "NetworkManager is running"
        return 0
    else
        print_warning "NetworkManager is not running"
        return 1
    fi
}

# Get WiFi interface name
get_wifi_interface() {
    WIFI_INTERFACE=$(iw dev | grep Interface | head -n1 | awk '{print $2}')
    if [[ -z "$WIFI_INTERFACE" ]]; then
        print_error "No WiFi interface found"
        exit 1
    fi
    print_status "Found WiFi interface: $WIFI_INTERFACE"
}

# Install required packages
install_packages() {
    case $DISTRO_TYPE in
        arch)
            if ! command -v iw &> /dev/null; then
                print_status "Installing iw..."
                sudo pacman -S --needed --noconfirm iw
            fi
            ;;
        fedora)
            if ! command -v iw &> /dev/null; then
                print_status "Installing iw..."
                sudo dnf install -y iw
            fi
            ;;
    esac
}

# Method 1: NetworkManager configuration
configure_networkmanager() {
    print_status "Configuring NetworkManager to disable WiFi power save..."
    
    # Create NetworkManager conf.d directory if it doesn't exist
    sudo mkdir -p /etc/NetworkManager/conf.d/
    
    # Create the configuration file
    sudo tee /etc/NetworkManager/conf.d/wifi-powersave-off.conf > /dev/null <<EOF
[connection]
wifi.powersave = 2
EOF
    
    print_success "NetworkManager configuration created"
    
    # Restart NetworkManager
    print_status "Restarting NetworkManager..."
    sudo systemctl restart NetworkManager
    sleep 3
    
    print_success "NetworkManager restarted"
}

# Method 2: systemd service (fallback)
create_systemd_service() {
    print_status "Creating systemd service as fallback method..."
    
    sudo tee /etc/systemd/system/wifi-powersave-off.service > /dev/null <<EOF
[Unit]
Description=Disable WiFi Power Save
After=network.target
Wants=network.target

[Service]
Type=oneshot
ExecStart=/usr/bin/iw dev ${WIFI_INTERFACE} set power_save off
RemainAfterExit=yes
StandardOutput=journal
StandardError=journal

[Install]
WantedBy=multi-user.target
EOF
    
    # Reload systemd and enable the service
    sudo systemctl daemon-reload
    sudo systemctl enable wifi-powersave-off.service
    
    print_success "Systemd service created and enabled"
}

# Method 3: udev rule (additional method)
create_udev_rule() {
    print_status "Creating udev rule for automatic power save disable..."
    
    sudo tee /etc/udev/rules.d/70-wifi-powersave-off.rules > /dev/null <<EOF
# Disable WiFi power save on interface up
ACTION=="add", SUBSYSTEM=="net", KERNEL=="wlan*", RUN+="/usr/bin/iw dev %k set power_save off"
ACTION=="add", SUBSYSTEM=="net", KERNEL=="wlp*", RUN+="/usr/bin/iw dev %k set power_save off"
EOF
    
    # Reload udev rules
    sudo udevadm control --reload-rules
    sudo udevadm trigger
    
    print_success "Udev rule created"
}

# Verify the configuration
verify_configuration() {
    print_status "Verifying WiFi power save configuration..."
    
    sleep 5  # Wait for changes to take effect
    
    # Check current power management status
    if command -v iw &> /dev/null; then
        POWER_STATUS=$(iw dev "$WIFI_INTERFACE" get power_save 2>/dev/null || echo "")

        if [[ $POWER_STATUS == *"off"* ]]; then
            print_success "WiFi power save is DISABLED"
            echo "Current status: $POWER_STATUS"
        elif [[ $POWER_STATUS == *"on"* ]]; then
            print_warning "WiFi power save is still ENABLED"
            echo "Current status: $POWER_STATUS"
            print_status "You may need to reboot or reconnect to WiFi for changes to take full effect"
        else
            print_warning "Could not determine power save status"
        fi
    else
        print_warning "iw not available, cannot verify status"
    fi
}

# Create a script to manually disable power save
create_manual_script() {
    print_status "Creating manual disable script..."
    
    sudo tee /usr/local/bin/disable-wifi-powersave > /dev/null <<EOF
#!/bin/bash
# Manual WiFi power save disable script

WIFI_INTERFACE=\$(iw dev | grep Interface | head -n1 | awk '{print \$2}')

if [[ -n "\$WIFI_INTERFACE" ]]; then
    echo "Disabling power save on \$WIFI_INTERFACE..."
    iw dev "\$WIFI_INTERFACE" set power_save off
    echo "Power save disabled."
    iw dev "\$WIFI_INTERFACE" get power_save
else
    echo "No WiFi interface found"
    exit 1
fi
EOF
    
    sudo chmod +x /usr/local/bin/disable-wifi-powersave
    print_success "Manual script created at /usr/local/bin/disable-wifi-powersave"
}

# Show final instructions
show_final_instructions() {
    echo
    print_success "WiFi power save disable setup completed!"
    echo
    echo "What was configured:"
    if check_networkmanager; then
        echo "  ✓ NetworkManager configuration (primary method)"
    fi
    echo "  ✓ Systemd service (backup/redundancy method)"
    echo "  ✓ Udev rules (automatic method)"
    echo "  ✓ Manual script for immediate disable"
    echo
    echo "Next steps:"
    echo "  1. Reboot your system for all changes to take effect"
    echo "  2. After reboot, verify with: iw dev $WIFI_INTERFACE get power_save"
    echo "  3. For immediate disable (without reboot): sudo disable-wifi-powersave"
    echo
    echo "If issues persist after reboot:"
    echo "  • Check NetworkManager status: systemctl status NetworkManager"
    echo "  • Check systemd service: systemctl status wifi-powersave-off"
    echo "  • Run manual script: sudo disable-wifi-powersave"
    echo
}

# Main execution
main() {
    echo "================================================="
    echo "    WiFi Power Save Disable Script"
    echo "    Supports Arch Linux and Fedora"
    echo "================================================="
    echo
    
    check_root
    detect_distro
    get_wifi_interface
    install_packages
    
    if check_networkmanager; then
        configure_networkmanager
        print_status "Also creating systemd service for redundancy..."
    else
        print_warning "NetworkManager not available, using alternative methods"
    fi
    
    create_systemd_service
    create_udev_rule
    create_manual_script
    verify_configuration
    show_final_instructions
}

# Run main function
main "$@"