#!/usr/bin/env bash
# RustDesk Remote Desktop Setup Script

set -euo pipefail

# Load environment variables
source /tmp/env-desktop

log() {
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] RUSTDESK: $1"
}

error() {
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] RUSTDESK ERROR: $1" >&2
    exit 1
}

# Detect OS and architecture
detect_system() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS=$ID
        VERSION=$VERSION_ID
    else
        error "Cannot detect OS"
    fi
    
    ARCH=$(uname -m)
    case $ARCH in
        x86_64)
            RUSTDESK_ARCH="x86_64"
            ;;
        aarch64|arm64)
            RUSTDESK_ARCH="aarch64"
            ;;
        *)
            error "Unsupported architecture: $ARCH"
            ;;
    esac
    
    log "Detected system: $OS $VERSION ($RUSTDESK_ARCH)"
}

# Download and install RustDesk
install_rustdesk() {
    log "Installing RustDesk for $OS..."
    
    local temp_dir="/tmp/rustdesk-install"
    mkdir -p "$temp_dir"
    cd "$temp_dir"
    
    case $OS in
        ubuntu|debian)
            # Download RustDesk .deb package
            local rustdesk_url="https://github.com/rustdesk/rustdesk/releases/latest/download/rustdesk-${RUSTDESK_ARCH}.deb"
            log "Downloading RustDesk from $rustdesk_url"
            
            curl -fsSL -o rustdesk.deb "$rustdesk_url" || error "Failed to download RustDesk"
            
            # Install dependencies
            apt-get update
            apt-get install -y libgtk-3-0 libxcb-randr0 libxdo3 libxfixes3 libasound2-dev libsystemd0
            
            # Install RustDesk
            dpkg -i rustdesk.deb || apt-get install -f -y
            ;;
            
        centos|rhel|fedora)
            # Download RustDesk .rpm package  
            local rustdesk_url="https://github.com/rustdesk/rustdesk/releases/latest/download/rustdesk-${RUSTDESK_ARCH}.rpm"
            log "Downloading RustDesk from $rustdesk_url"
            
            curl -fsSL -o rustdesk.rpm "$rustdesk_url" || error "Failed to download RustDesk"
            
            # Install dependencies
            if command -v dnf >/dev/null 2>&1; then
                dnf install -y gtk3 libxcb libXfixes alsa-lib systemd
                dnf install -y rustdesk.rpm
            else
                yum install -y gtk3 libxcb libXfixes alsa-lib systemd
                yum localinstall -y rustdesk.rpm
            fi
            ;;
            
        *)
            error "Unsupported OS for RustDesk installation: $OS"
            ;;
    esac
    
    # Clean up
    cd /
    rm -rf "$temp_dir"
    
    log "RustDesk installation completed"
}

# Configure RustDesk
configure_rustdesk() {
    local desktop_user="${DESKTOP_USER:-desktop}"
    local desktop_home="${DESKTOP_HOME:-/home/$desktop_user}"
    
    log "Configuring RustDesk for user $desktop_user"
    
    # Create RustDesk config directory
    sudo -u "$desktop_user" mkdir -p "$desktop_home/.config/rustdesk"
    
    # Create RustDesk configuration
    local config_file="$desktop_home/.config/rustdesk/RustDesk2.toml"
    
    cat > "$config_file" << EOF
[options]
custom-rendezvous-server = "${RUSTDESK_CUSTOM_SERVER:-}"
relay-server = "${RUSTDESK_CUSTOM_SERVER:-}"
api-server = "${RUSTDESK_CUSTOM_SERVER:-}"
key = ""
auto-disconnect-timeout = "10"
keep-screen-on = "Y"
wake-on-lan = "Y"
allow-guest-access = "${RUSTDESK_ALLOW_GUEST:-N}"

[ui]
hide-cm = ""
hide-connection-management = ""
hide-network-setting = ""
hide-password-setting = ""
hide-about-link = ""
hide-software-update = ""

[network]
rendezvous-server = "${RUSTDESK_CUSTOM_SERVER:-}"
nat-type = ""
serial = ""

[security]
access-mode = "custom"
EOF

    # Set custom server if provided
    if [[ -n "${RUSTDESK_CUSTOM_SERVER:-}" ]]; then
        log "Using custom RustDesk server: $RUSTDESK_CUSTOM_SERVER"
    fi
    
    # Set permanent password if provided
    if [[ -n "${RUSTDESK_PERMANENT_PASSWORD:-}" ]]; then
        log "Setting permanent password for RustDesk"
        # Note: RustDesk permanent password is set via GUI or command line
        # This is a placeholder for the configuration
        echo "permanent_password = \"$RUSTDESK_PERMANENT_PASSWORD\"" >> "$config_file"
    fi
    
    chown -R "$desktop_user:$desktop_user" "$desktop_home/.config/rustdesk"
    
    log "RustDesk configuration created"
}

# Create RustDesk systemd service
create_rustdesk_service() {
    local desktop_user="${DESKTOP_USER:-desktop}"
    
    log "Creating RustDesk systemd service for user $desktop_user"
    
    # Create systemd user service
    local service_dir="/home/$desktop_user/.config/systemd/user"
    mkdir -p "$service_dir"
    
    cat > "$service_dir/rustdesk.service" << EOF
[Unit]
Description=RustDesk Remote Desktop
After=graphical-session.target

[Service]
Type=simple
ExecStart=/usr/bin/rustdesk --service
Restart=always
RestartSec=5
Environment=DISPLAY=:0

[Install]
WantedBy=default.target
EOF
    
    chown -R "$desktop_user:$desktop_user" "/home/$desktop_user/.config/systemd"
    
    # Enable user service
    sudo -u "$desktop_user" systemctl --user daemon-reload
    
    if [[ "${RUSTDESK_AUTO_START:-true}" == "true" ]]; then
        sudo -u "$desktop_user" systemctl --user enable rustdesk.service
        log "RustDesk service enabled for auto-start"
    fi
    
    log "RustDesk systemd service created"
}

# Setup RustDesk desktop shortcut
create_desktop_shortcut() {
    local desktop_user="${DESKTOP_USER:-desktop}"
    local desktop_home="${DESKTOP_HOME:-/home/$desktop_user}"
    
    log "Creating RustDesk desktop shortcut"
    
    cat > "$desktop_home/Desktop/rustdesk.desktop" << 'EOF'
[Desktop Entry]
Version=1.0
Type=Application
Name=RustDesk
Comment=Remote Desktop Software
Exec=rustdesk
Icon=rustdesk
Terminal=false
StartupNotify=true
Categories=Network;RemoteAccess;
Keywords=remote;desktop;vnc;connection;
EOF
    
    chmod +x "$desktop_home/Desktop/rustdesk.desktop"
    chown "$desktop_user:$desktop_user" "$desktop_home/Desktop/rustdesk.desktop"
    
    log "RustDesk desktop shortcut created"
}

# Setup firewall rules for RustDesk
setup_firewall() {
    log "Setting up firewall rules for RustDesk"
    
    local rustdesk_port="${RUSTDESK_PORT:-21116}"
    local rustdesk_hbbr_port="${RUSTDESK_HBBR_PORT:-21117}"
    
    # Try different firewall tools
    if command -v ufw >/dev/null 2>&1; then
        ufw allow "$rustdesk_port/tcp" comment "RustDesk"
        ufw allow "$rustdesk_port/udp" comment "RustDesk"  
        ufw allow "$rustdesk_hbbr_port/tcp" comment "RustDesk hbbr"
        log "UFW rules added for RustDesk ports $rustdesk_port and $rustdesk_hbbr_port"
    elif command -v firewall-cmd >/dev/null 2>&1; then
        firewall-cmd --permanent --add-port="$rustdesk_port/tcp"
        firewall-cmd --permanent --add-port="$rustdesk_port/udp"
        firewall-cmd --permanent --add-port="$rustdesk_hbbr_port/tcp"
        firewall-cmd --reload
        log "FirewallD rules added for RustDesk ports $rustdesk_port and $rustdesk_hbbr_port"
    else
        log "WARNING: No supported firewall tool found. Manual firewall configuration may be needed."
    fi
}

# Get RustDesk ID and password
get_rustdesk_info() {
    log "RustDesk installation completed!"
    log "To get your RustDesk ID and password, run:"
    log "  sudo -u $DESKTOP_USER rustdesk --get-id"
    log "  sudo -u $DESKTOP_USER rustdesk --password"
    log ""
    log "RustDesk will be available on ports:"
    log "  Main port: ${RUSTDESK_PORT:-21116}"
    log "  hbbr port: ${RUSTDESK_HBBR_PORT:-21117}"
}

# Main installation function
main() {
    if [[ "${RUSTDESK_ENABLED:-true}" != "true" ]]; then
        log "RustDesk is disabled, skipping installation"
        return 0
    fi
    
    log "Starting RustDesk installation and configuration..."
    
    detect_system
    install_rustdesk
    configure_rustdesk
    create_rustdesk_service
    create_desktop_shortcut
    setup_firewall
    get_rustdesk_info
    
    log "RustDesk setup completed successfully!"
}

# Run main function if script is executed directly
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi