#!/usr/bin/env bash
# Desktop Environment Installation Script
# Installs minimal desktop environment with essential applications

set -euo pipefail

# Load environment variables
source /tmp/env-desktop

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

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

# Detect OS
detect_os() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS=$ID
        VERSION=$VERSION_ID
    else
        error "Cannot detect OS"
    fi
    log "Detected OS: $OS $VERSION"
}

# Update system packages
update_system() {
    log "Updating system packages..."
    case $OS in
        ubuntu|debian)
            apt-get update -y
            apt-get upgrade -y
            ;;
        centos|rhel|fedora)
            if command -v dnf >/dev/null 2>&1; then
                dnf update -y
            else
                yum update -y
            fi
            ;;
        *)
            error "Unsupported OS: $OS"
            ;;
    esac
}

# Install desktop environment
install_desktop_environment() {
    log "Installing $DESKTOP_TYPE desktop environment..."
    
    case $OS in
        ubuntu|debian)
            case $DESKTOP_TYPE in
                xfce)
                    apt-get install -y xfce4 xfce4-goodies
                    if [[ "$DISPLAY_MANAGER" == "lightdm" ]]; then
                        apt-get install -y lightdm lightdm-gtk-greeter
                    fi
                    ;;
                gnome)
                    apt-get install -y ubuntu-desktop-minimal
                    ;;
                kde)
                    apt-get install -y kde-plasma-desktop
                    ;;
                lxde)
                    apt-get install -y lxde
                    ;;
                mate)
                    apt-get install -y ubuntu-mate-desktop
                    ;;
            esac
            ;;
        centos|rhel|fedora)
            case $DESKTOP_TYPE in
                xfce)
                    if command -v dnf >/dev/null 2>&1; then
                        dnf groupinstall -y "Xfce Desktop"
                    else
                        yum groupinstall -y "Xfce Desktop"
                    fi
                    ;;
                gnome)
                    if command -v dnf >/dev/null 2>&1; then
                        dnf groupinstall -y "GNOME Desktop Environment"
                    else
                        yum groupinstall -y "GNOME Desktop Environment"
                    fi
                    ;;
            esac
            ;;
    esac
}

# Install VNC server
install_vnc_server() {
    if [[ "$VNC_ENABLED" == "true" ]]; then
        log "Installing VNC server..."
        
        case $OS in
            ubuntu|debian)
                apt-get install -y tightvncserver
                ;;
            centos|rhel|fedora)
                if command -v dnf >/dev/null 2>&1; then
                    dnf install -y tigervnc-server
                else
                    yum install -y tigervnc-server
                fi
                ;;
        esac
        
        # Configure VNC for desktop user
        setup_vnc_user
    fi
}

# Setup VNC for desktop user
setup_vnc_user() {
    log "Setting up VNC for user $DESKTOP_USER..."
    
    # Create user if not exists
    if ! id "$DESKTOP_USER" &>/dev/null; then
        useradd -m -s /bin/bash "$DESKTOP_USER"
        log "Created user $DESKTOP_USER"
    fi
    
    # Setup VNC directory
    sudo -u "$DESKTOP_USER" mkdir -p "$DESKTOP_HOME/.vnc"
    
    # Create VNC startup script
    cat > "$DESKTOP_HOME/.vnc/xstartup" << EOF
#!/bin/bash
xrdb \$HOME/.Xresources
startxfce4 &
EOF
    
    chmod +x "$DESKTOP_HOME/.vnc/xstartup"
    chown "$DESKTOP_USER:$DESKTOP_USER" "$DESKTOP_HOME/.vnc/xstartup"
    
    # Set VNC password if provided
    if [[ -n "${VNC_PASSWORD:-}" ]]; then
        echo "$VNC_PASSWORD" | sudo -u "$DESKTOP_USER" vncpasswd -f > "$DESKTOP_HOME/.vnc/passwd"
        chmod 600 "$DESKTOP_HOME/.vnc/passwd"
        chown "$DESKTOP_USER:$DESKTOP_USER" "$DESKTOP_HOME/.vnc/passwd"
    fi
    
    # Create VNC service
    create_vnc_service
}

# Create VNC systemd service
create_vnc_service() {
    log "Creating VNC systemd service..."
    
    cat > "/etc/systemd/system/vncserver@.service" << EOF
[Unit]
Description=Start TightVNC server at startup
After=syslog.target network.target

[Service]
Type=forking
User=$DESKTOP_USER
Group=$DESKTOP_USER
WorkingDirectory=$DESKTOP_HOME

PIDFile=$DESKTOP_HOME/.vnc/%H:%i.pid
ExecStartPre=-/usr/bin/vncserver -kill :%i > /dev/null 2>&1
ExecStart=/usr/bin/vncserver -depth $VNC_DEPTH -geometry $VNC_GEOMETRY :%i
ExecStop=/usr/bin/vncserver -kill :%i

[Install]
WantedBy=multi-user.target
EOF
    
    systemctl daemon-reload
    systemctl enable "vncserver@1.service"
    log "VNC service created and enabled"
}

# Install applications
install_applications() {
    log "Installing applications..."
    
    case $OS in
        ubuntu|debian)
            # Install packages
            local packages=""
            
            # Editors
            for editor in $EDITORS; do
                case $editor in
                    zed)
                        # Install Zed editor
                        install_zed_editor
                        ;;
                    *)
                        packages="$packages $editor"
                        ;;
                esac
            done
            
            # Add other application categories
            packages="$packages $BROWSERS $TERMINALS $DEVELOPMENT $MEDIA $OFFICE $UTILITIES"
            
            if [[ -n "$packages" ]]; then
                apt-get install -y $packages
            fi
            ;;
        centos|rhel|fedora)
            local packages="$BROWSERS $TERMINALS $DEVELOPMENT $MEDIA $OFFICE $UTILITIES"
            
            # Install Zed if in editors list
            if echo "$EDITORS" | grep -q "zed"; then
                install_zed_editor
            fi
            
            # Remove zed from package list and add other editors
            local filtered_editors=$(echo "$EDITORS" | sed 's/zed//g')
            packages="$packages $filtered_editors"
            
            if command -v dnf >/dev/null 2>&1; then
                dnf install -y $packages
            else
                yum install -y $packages
            fi
            ;;
    esac
}

# Install Zed editor
install_zed_editor() {
    log "Installing Zed editor..."
    
    # Download and install Zed
    case $(uname -m) in
        x86_64)
            curl -f https://zed.dev/install.sh | sh
            ;;
        *)
            log "Zed editor not available for $(uname -m) architecture, skipping..."
            ;;
    esac
}

# Configure graphics
configure_graphics() {
    log "Configuring graphics driver: $GRAPHICS_DRIVER"
    
    case $OS in
        ubuntu|debian)
            case $GRAPHICS_DRIVER in
                nvidia)
                    apt-get install -y nvidia-driver-470
                    ;;
                amd)
                    apt-get install -y mesa-vulkan-drivers xserver-xorg-video-amdgpu
                    ;;
                intel)
                    apt-get install -y mesa-vulkan-drivers xserver-xorg-video-intel
                    ;;
                nouveau)
                    apt-get install -y xserver-xorg-video-nouveau
                    ;;
            esac
            ;;
    esac
}

# Setup auto-login if enabled
setup_auto_login() {
    if [[ "$AUTO_LOGIN" == "true" ]]; then
        log "Setting up auto-login for $DESKTOP_USER..."
        
        case $DISPLAY_MANAGER in
            lightdm)
                sed -i "s/#autologin-user=/autologin-user=$DESKTOP_USER/" /etc/lightdm/lightdm.conf
                sed -i "s/#autologin-user-timeout=0/autologin-user-timeout=0/" /etc/lightdm/lightdm.conf
                ;;
            gdm)
                cat > "/etc/gdm3/custom.conf" << EOF
[daemon]
AutomaticLoginEnable=true
AutomaticLogin=$DESKTOP_USER
EOF
                ;;
        esac
    fi
}

# Run remote access setup scripts
setup_remote_access() {
    log "Setting up remote access services..."
    
    # Run SSH setup if enabled
    if [[ "${SSH_ENABLED:-true}" == "true" ]]; then
        log "Running SSH setup..."
        bash /tmp/ssh-setup.sh
    fi
    
    # Run RustDesk setup if enabled
    if [[ "${RUSTDESK_ENABLED:-true}" == "true" ]]; then
        log "Running RustDesk setup..."
        bash /tmp/rustdesk-setup.sh
    fi
    
    # Run Zed setup
    log "Running Zed editor setup..."
    bash /tmp/zed-setup.sh
}

# Display connection summary
display_connection_summary() {
    log ""
    log "=== Desktop Environment Setup Complete ==="
    log ""
    log "Remote Access Options:"
    
    if [[ "${VNC_ENABLED:-true}" == "true" ]]; then
        log "  VNC Server: Port $VNC_PORT"
        log "    Start with: systemctl start vncserver@1.service"
    fi
    
    if [[ "${RUSTDESK_ENABLED:-true}" == "true" ]]; then
        log "  RustDesk: Ports $RUSTDESK_PORT (main), $RUSTDESK_HBBR_PORT (hbbr)"
        log "    Get ID: sudo -u $DESKTOP_USER rustdesk --get-id"
    fi
    
    if [[ "${SSH_ENABLED:-true}" == "true" ]]; then
        log "  SSH Server: Port $SSH_PORT"
        log "    Connect: ssh $DESKTOP_USER@<server-ip> -p $SSH_PORT"
    fi
    
    log ""
    log "Desktop Environment: $DESKTOP_TYPE"
    log "Desktop User: $DESKTOP_USER"
    log "Applications installed: Zed editor and standard desktop apps"
}

# Main installation function
main() {
    log "Starting desktop environment installation..."
    
    detect_os
    update_system
    install_desktop_environment
    install_vnc_server
    install_applications
    configure_graphics
    setup_auto_login
    setup_remote_access
    
    display_connection_summary
    log "Desktop environment installation completed successfully!"
}

# Run main function
main "$@"