#!/bin/bash
# wps-office-integrator – fixes integration issues for WPS Office
# Version 1.6 – preserves correct desktop entry names (wps, wpp, et, pdf)

set -e

# Colors
if [ -t 1 ]; then
    GREEN='\033[0;32m'; BLUE='\033[0;34m'; RED='\033[0;31m'; YELLOW='\033[1;33m'; NC='\033[0m'
else
    GREEN=''; BLUE=''; RED=''; YELLOW=''; NC=''
fi
info() { echo -e "${BLUE}➡${NC} $1"; }
success() { echo -e "${GREEN}✅${NC} $1"; }
error() { echo -e "${RED}❌${NC} $1"; }
warn() { echo -e "${YELLOW}⚠${NC} $1"; }

if [ "$EUID" -eq 0 ]; then
    warn "Running as root. Some operations will affect all users."
fi

# 1. Fix document templates – remove all non‑working, keep only Chinese symlinks
fix_templates() {
    info "Fixing document templates..."
    TEMPLATE_SOURCE="/opt/kingsoft/wps-office/office6/asso_template"
    TEMPLATE_TARGET="/usr/share/templates"
    USER_TEMPLATES="$HOME/Templates"
    
    if [ ! -d "$TEMPLATE_SOURCE" ]; then
        warn "WPS template folder not found: $TEMPLATE_SOURCE"
        return 1
    fi
    
    # Remove old WPS-related files from system templates
    find "$TEMPLATE_TARGET" -maxdepth 1 -type f \( \
        -name "*.wps" -o -name "*.wpt" -o -name "*.dot" -o -name "*.dotx" \
        -o -name "*.xlt" -o -name "*.xltx" -o -name "*.pot" -o -name "*.potx" \
        -o -name "wps文字文档*" -o -name "wps演示文档*" -o -name "wps表格文档*" \
        -o -name "*.doc" -o -name "*.docx" -o -name "*.dps" -o -name "*.pptx" \
        -o -name "*.et" -o -name "*.xlsx" \
    \) -exec rm -f {} \; 2>/dev/null || true
    success "Old WPS template files removed from $TEMPLATE_TARGET"
    
    # Remove from user templates as well
    if [ -d "$USER_TEMPLATES" ]; then
        find "$USER_TEMPLATES" -maxdepth 1 -type f \( \
            -name "*.wps" -o -name "*.wpt" -o -name "*.dot" -o -name "*.dotx" \
            -o -name "*.xlt" -o -name "*.xltx" -o -name "*.pot" -o -name "*.potx" \
            -o -name "wps文字文档*" -o -name "wps演示文档*" -o -name "wps表格文档*" \
            -o -name "*.doc" -o -name "*.docx" -o -name "*.dps" -o -name "*.pptx" \
            -o -name "*.et" -o -name "*.xlsx" \
        \) -exec rm -f {} \; 2>/dev/null || true
        success "WPS template files removed from $USER_TEMPLATES"
    fi
    
    mkdir -p "$TEMPLATE_TARGET"
    
    # Create fresh symlinks only for working Chinese templates
    declare -A templates=(
        ["wps.doc"]="wps文字文档.doc"
        ["wps.docx"]="wps文字文档.docx"
        ["wps.dps"]="wps演示文档.dps"
        ["wps.pptx"]="wps演示文档.pptx"
        ["wps.et"]="wps表格文档.et"
        ["wps.xlsx"]="wps表格文档.xlsx"
    )
    
    for src in "${!templates[@]}"; do
        target="${templates[$src]}"
        if [ -f "$TEMPLATE_SOURCE/$src" ]; then
            ln -sf "$TEMPLATE_SOURCE/$src" "$TEMPLATE_TARGET/$target"
            success "Created symlink: $target"
        else
            warn "Template file missing: $src"
        fi
    done
}

# 2. Ensure user ~/Templates exists
fix_user_templates_dir() {
    info "Ensuring user Templates directory exists..."
    if command -v xdg-user-dirs-update &>/dev/null; then
        xdg-user-dirs-update
        success "User Templates directory created (if missing)."
    else
        warn "xdg-user-dirs-update not found, skipping."
    fi
}

# 3. Update MIME database
fix_mime() {
    info "Updating MIME database..."
    if [ -d /usr/share/mime/packages ]; then
        update-mime-database /usr/share/mime &>/dev/null
        success "MIME database updated."
    else
        warn "/usr/share/mime/packages not found."
    fi
}

# 4. Update desktop database
fix_desktop_database() {
    info "Updating desktop database..."
    update-desktop-database /usr/share/applications &>/dev/null
    success "Desktop database updated."
}

# 5. Safe cleanup: remove only known broken/duplicate .desktop files,
#    preserving the correct ones (wps-office-wps, wpp, et, pdf, and builder)
cleanup_old_desktop_files() {
    info "Cleaning up known broken/duplicate WPS .desktop files..."
    
    # Preserve the real desktop entry names used by WPS Office on Fedora/Arch
    local keep_files=(
        "wps-office-wps.desktop"          # Writer
        "wps-office-wpp.desktop"          # Presentation
        "wps-office-et.desktop"           # Spreadsheets
        "wps-office-pdf.desktop"          # PDF reader
        "wps-office-builder.desktop"      # optional, if present
    )
    local keep_pattern=$(printf "|%s" "${keep_files[@]}")
    keep_pattern="(${keep_pattern:1})"
    
    find /usr/share/applications -maxdepth 1 -type f -name "wps-office*.desktop" | while read -r file; do
        local base=$(basename "$file")
        if ! echo "$base" | grep -qE "$keep_pattern"; then
            rm -f "$file"
            success "Removed duplicate/broken: $base"
        fi
    done
}

# 6. Fix Exec parameter (%U -> %F) and MimeType in .desktop files
fix_desktop_exec() {
    info "Fixing Exec parameters and MimeType in WPS .desktop files..."
    
    local desktop_files=(
        "/usr/share/applications/wps-office-wps.desktop"
        "/usr/share/applications/wps-office-wpp.desktop"
        "/usr/share/applications/wps-office-et.desktop"
        "/usr/share/applications/wps-office-pdf.desktop"
    )
    
    for desk in "${desktop_files[@]}"; do
        if [ -f "$desk" ]; then
            if grep -q '%U' "$desk"; then
                sed -i 's/%U/%F/g' "$desk"
                success "Fixed Exec in $(basename "$desk") (%U -> %F)"
            fi
            
            local base=$(basename "$desk")
            local new_mime=""
            case "$base" in
                wps-office-wps.desktop)
                    new_mime="application/msword;application/vnd.openxmlformats-officedocument.wordprocessingml.document;application/rtf;text/plain;"
                    ;;
                wps-office-et.desktop)
                    new_mime="application/vnd.ms-excel;application/vnd.openxmlformats-officedocument.spreadsheetml.sheet;text/csv;"
                    ;;
                wps-office-wpp.desktop)
                    new_mime="application/vnd.ms-powerpoint;application/vnd.openxmlformats-officedocument.presentationml.presentation;"
                    ;;
                wps-office-pdf.desktop)
                    new_mime="application/pdf;"
                    ;;
            esac
            if [ -n "$new_mime" ] && ! grep -q "MimeType=.*$new_mime" "$desk"; then
                sed -i "s|^MimeType=.*|MimeType=$new_mime|" "$desk"
                success "Updated MimeType in $(basename "$desk")"
            fi
        else
            warn "Desktop file not found: $desk"
        fi
    done
}

# 7. Wayland fix (force X11 for WPS)
fix_wayland_env() {
    info "Setting up Wayland compatibility (force X11 for WPS)..."
    ENV_FILE="/etc/profile.d/wps-wayland-fix.sh"
    if [ ! -f "$ENV_FILE" ]; then
        cat > "$ENV_FILE" << 'EOF'
#!/bin/bash
export QT_QPA_PLATFORM=xcb
EOF
        chmod 644 "$ENV_FILE"
        success "Created environment file: $ENV_FILE"
    else
        info "Environment file already exists."
    fi
}

# 8. Refresh KDE menu cache
refresh_kde_cache() {
    if command -v kbuildsycoca6 &>/dev/null; then
        info "Refreshing KDE menu cache (kbuildsycoca6)..."
        kbuildsycoca6 &>/dev/null
        success "KDE menu cache refreshed."
    elif command -v kbuildsycoca5 &>/dev/null; then
        info "Refreshing KDE menu cache (kbuildsycoca5)..."
        kbuildsycoca5 &>/dev/null
        success "KDE menu cache refreshed."
    else
        warn "kbuildsycoca not found. Menu may not update until next login."
    fi
}

# 9. Restore missing essential WPS desktop files (interactive)
restore_wps_desktop_files() {
    local missing=0
    local required=(
        "/usr/share/applications/wps-office-wps.desktop"
        "/usr/share/applications/wps-office-wpp.desktop"
        "/usr/share/applications/wps-office-et.desktop"
        "/usr/share/applications/wps-office-pdf.desktop"
    )
    for desk in "${required[@]}"; do
        if [ ! -f "$desk" ]; then
            warn "Missing $(basename "$desk")"
            missing=1
        fi
    done
    
    if [ $missing -eq 1 ]; then
        info "Attempting to restore missing WPS desktop files..."
        if command -v dnf &>/dev/null; then
            echo -e "\n⚠ Some WPS menu entries are missing. Would you like to restore them by reinstalling 'wps-office'? (y/n)"
            read -r answer
            if [[ "$answer" =~ ^[Yy]$ ]]; then
                dnf reinstall -y wps-office
                success "Reinstalled wps-office. Desktop files restored."
            else
                warn "Skipping restoration. You may manually run: sudo dnf reinstall wps-office"
            fi
        else
            warn "dnf not found. Please reinstall wps-office manually to restore menu entries."
        fi
    else
        success "All essential WPS desktop files are present."
    fi
}

# Main
main() {
    echo "=== WPS Office Integrator v1.6 ==="
    echo "Applying fixes for WPS Office integration"
    fix_templates
    fix_user_templates_dir
    fix_mime
    fix_desktop_database
    cleanup_old_desktop_files
    fix_desktop_exec
    fix_wayland_env
    refresh_kde_cache
    restore_wps_desktop_files
    echo ""
    success "All fixes applied successfully."
    info "It is recommended to log out and back in (or restart KDE/plasmashell)."
    info "For Plasma 6: kquitapp6 plasmashell && kstart6 plasmashell"
}

if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main
fi