#!/bin/bash

# 截屏脚本 - 适配多种桌面环境和发行版

# 配置变量
readonly SCREENSHOT_TARGET_DIR="/usr/lib/huiontablet/res"
readonly SCREENSHOT_TARGET_FILE="$SCREENSHOT_TARGET_DIR/screenshot.png"
readonly LOG_FILE="/tmp/screenshot_script.log"

# 创建日志函数
log_message() {
    # echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" >> "$LOG_FILE"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
}

# 检测桌面环境
detect_desktop_environment() {
    if [ -n "$XDG_CURRENT_DESKTOP" ]; then
        echo "$XDG_CURRENT_DESKTOP" | tr '[:upper:]' '[:lower:]'
    elif [ -n "$GDMSESSION" ]; then
        echo "$GDMSESSION" | tr '[:upper:]' '[:lower:]'
    elif [ -n "$DESKTOP_SESSION" ]; then
        echo "$DESKTOP_SESSION" | tr '[:upper:]' '[:lower:]'
    else
        # 通过进程检测
        if pgrep -x "gnome-shell" >/dev/null; then
            echo "gnome"
        elif pgrep -x "plasmashell" >/dev/null; then
            echo "kde"
        elif pgrep -x "xfce4-session" >/dev/null; then
            echo "xfce"
        elif pgrep -x "cinnamon" >/dev/null; then
            echo "cinnamon"
        elif pgrep -x "mate-session" >/dev/null; then
            echo "mate"
        elif pgrep -x "deepin-session" >/dev/null; then
            echo "deepin"
        elif pgrep -x "budgie-panel" >/dev/null; then
            echo "budgie"
        else
            echo "unknown"
        fi
    fi
}

# 检测可用的截屏工具
detect_screenshot_tool() {
    local tools=("gnome-screenshot" "spectacle" "ksnapshot" "xfce4-screenshooter" 
                "scrot" "maim" "import" "deepin-screenshot" "mate-screenshot" "kylin-screenshot")
    
    for tool in "${tools[@]}"; do
        if command -v "$tool" >/dev/null 2>&1; then
            echo "$tool"
            return 0
        fi
    done
    echo "none"
}

# 获取当前时间戳
get_current_timestamp() {
    date +%s
}

# 获取当前登录的图形用户
get_graphical_user() {
    # 方法1: 通过loginctl查找图形会话
    local user=$(loginctl list-sessions --no-legend | grep -E '(wayland|x11)' | awk '{print $3}' | head -1)
    
    # 方法2: 通过who命令
    if [ -z "$user" ]; then
        user=$(who | awk '$2 ~ /:[0-9]/ {print $1}' | head -1)
    fi
    
    # 方法3: 通过进程查找
    if [ -z "$user" ]; then
        user=$(ps -o user= -C gnome-shell -C xfce4-session -C plasmashell 2>/dev/null | head -1)
    fi
    
    echo "${user:-$USER}"
}

# 完整的国际化图片目录检测函数
get_localized_pictures_dir() {
    local detected_dir=""
    
    # 1. 首先尝试 xdg-user-dir 命令（最准确）
    if command -v xdg-user-dir >/dev/null 2>&1; then
        detected_dir=$(xdg-user-dir PICTURES 2>/dev/null)
        if [ -n "$detected_dir" ] && [ -d "$detected_dir" ]; then
            echo "$detected_dir"
            return 0
        fi
    fi
    
    # 2. 尝试读取 XDG 配置文件
    if [ -f "${HOME}/.config/user-dirs.dirs" ]; then
        detected_dir=$(grep "^XDG_PICTURES_DIR" "${HOME}/.config/user-dirs.dirs" | cut -d'=' -f2- | tr -d '"' | sed "s#\\$HOME#${HOME}#" | sed "s#\\${HOME}#${HOME}#")
        if [ -n "$detected_dir" ] && [ -d "$detected_dir" ]; then
            echo "$detected_dir"
            return 0
        fi
    fi
    
    # 3. 检测常见语言特定的目录名称
    local common_dirs=(
        "${HOME}/Pictures"    # English
        "${HOME}/图片"        # Chinese
        "${HOME}/图片"        # Japanese (same as Chinese)
        "${HOME}/Bilder"      # German
        "${HOME}/Imágenes"    # Spanish
        "${HOME}/Images"      # French
        "${HOME}/Immagini"    # Italian
        "${HOME}/Afbeeldingen" # Dutch
        "${HOME}/Рисунки"     # Russian
        "${HOME}/사진"        # Korean
        "${HOME}/ภาพ"         # Thai
        "${HOME}/Resimler"    # Turkish
        "${HOME}/Obrázky"     # Czech
        "${HOME}/Obrázků"     # Slovak
    )
    
    for dir in "${common_dirs[@]}"; do
        if [ -d "$dir" ]; then
            echo "$dir"
            return 0
        fi
    done
    
    # 4. 最后尝试创建或使用默认英文目录
    echo "${HOME}/Pictures"
}

kylin_capture_and_move_screenshot() {
    local pictures_dir=$(get_localized_pictures_dir)
    local max_wait=10
    
    # 记录开始时间
    local start_time=$(date +%s)
    kylin-screenshot full
    
    # 等待并查找新文件
    local wait_time=0
    local new_file=""
    
    while [ $wait_time -lt $max_wait ]; do
        # 查找在截屏开始后创建的png文件
        new_file=$(find "$pictures_dir" -maxdepth 1 -name "*.png" -newermt "@$start_time" -type f -printf "%T@ %p\n" 2>/dev/null | \
                  sort -nr | head -1 | cut -d' ' -f2-)
        
        if [ -n "$new_file" ] && [ -f "$new_file" ]; then
            # echo "发现新文件: $new_file"
            
            # 移动文件
            if mv "$new_file" "$SCREENSHOT_TARGET_FILE" 2>/dev/null; then
                # echo "文件已移动: $SCREENSHOT_TARGET_FILE"
                return 0
            else
                # echo "错误: 文件移动失败" >&2
                echo "FAIL"
                return 1
            fi
        fi
        
        sleep 1
        wait_time=$((wait_time + 1))
    done
    
    # echo "错误: 等待超时，未发现新图片" >&2
    echo "FAIL"
    return 1
}

# 将~/图片/中的截屏图片移动到/tmp/screenshot.png
move_screenshot_to_tmp() {
    local pictures_dir=$(get_localized_pictures_dir)

    # 检查目录
    if [ ! -d "$pictures_dir" ]; then
        # echo "错误: 图片目录不存在: $pictures_dir" >&2
        echo "FAIL"
        exit 1
    fi

    # 查找最新的 png 文件
    latest_png=$(find "$pictures_dir" -maxdepth 1 -name "*.png" -type f -printf "%T@ %p\n" 2>/dev/null | sort -nr | head -1 | cut -d' ' -f2-)
    if [ -n "$latest_png" ] && [ -f "$latest_png" ]; then
        cp "$latest_png" "$SCREENSHOT_TARGET_FILE"
    else
        echo "FAIL"
        exit 1
    fi
}

# 在当前登录用户的图形桌面环境中执行截图指令
execute_screenshot_cmd_on_wayland() {
    local target_user=$(get_graphical_user)
    # 必要的Wayland环境变量
    local wayland_display="${WAYLAND_DISPLAY:-wayland-0}"
    local xdg_runtime_dir="${XDG_RUNTIME_DIR:-/run/user/$(id -u "$target_user")}"
    
    # 截图工具
    local screenshot_tool="$1"
    
    # 执行截图
    if [ "$(id -u)" -eq 0 ] && [ -n "$target_user" ]; then
        # 以目标用户身份运行
        sudo -u "$target_user" \
            env WAYLAND_DISPLAY="$wayland_display" \
                XDG_RUNTIME_DIR="$xdg_runtime_dir" \
                DBUS_SESSION_BUS_ADDRESS="unix:path=$xdg_runtime_dir/bus" \
                bash -c "$screenshot_tool"
        return 0
    else
        # 直接运行
        env WAYLAND_DISPLAY="$wayland_display" \
            XDG_RUNTIME_DIR="$xdg_runtime_dir" \
            bash -c "$screenshot_tool"
        return 0
    fi
    return 1
}

# UKUI 麒麟系统 桌面截屏
kylin_screenshot() {
    # echo "使用 kylin-screenshot"
    if kylin_capture_and_move_screenshot; then
        # cp "$(ls -t ~/图片/*.png | head -1)" /tmp/screenshot.png
        # move_screenshot_to_tmp
        return 0
    fi
    return 1
}

# GNOME 桌面截屏
gnome_screenshot() {
    # echo "使用 gnome-screenshot"
    # 截图工具
    local screenshot_tool="gnome-screenshot -f $1 >/dev/null 2>&1"
    # 执行截图
    if execute_screenshot_cmd_on_wayland "$screenshot_tool"; then
        return 0
    fi
    
    return 1
}

# KDE 桌面截屏
kde_screenshot() {
    # echo "使用 KDE 工具"
    # 尝试 spectacle (新版本KDE)
    if command -v spectacle >/dev/null 2>&1; then
        # 脚本用sudo权限运行的话，会一直卡在这一步，因为报错说无法保存内容为空的图片 
        local screenshot_tool="spectacle -b -n -o $1 >/dev/null 2>&1"
        if execute_screenshot_cmd_on_wayland "$screenshot_tool"; then
            return 0
        fi
    fi
    
    # 尝试 ksnapshot (旧版本KDE)
    if command -v ksnapshot >/dev/null 2>&1; then
        local screenshot_tool="ksnapshot -fullscreen -filename $1 >/dev/null 2>&1"
        if execute_screenshot_cmd_on_wayland "$screenshot_tool"; then
            return 0
        fi
    fi
    return 1
}

# XFCE 桌面截屏
xfce_screenshot() {
    # echo "使用 xfce4-screenshooter"
    local screenshot_tool="xfce4-screenshooter -f -s $1 >/dev/null 2>&1"
    if execute_screenshot_cmd_on_wayland "$screenshot_tool"; then
        return 0
    fi
    return 1
}

# 使用 scrot 工具
scrot_screenshot() {
    # echo "使用 scrot"
    local screenshot_tool="scrot $1 >/dev/null 2>&1"
    if execute_screenshot_cmd_on_wayland "$screenshot_tool"; then
        return 0
    fi
    return 1
}

# 使用 maim 工具
maim_screenshot() {
    # echo "使用 maim"
    local screenshot_tool="maim $1 >/dev/null 2>&1"
    if execute_screenshot_cmd_on_wayland "$screenshot_tool"; then
        return 0
    fi
    return 1
}

# 使用 ImageMagick 的 import
import_screenshot() {
    # echo "使用 ImageMagick import"
    local screenshot_tool="import -window root $1 >/dev/null 2>&1"
    if execute_screenshot_cmd_on_wayland "$screenshot_tool"; then
        return 0
    fi
    return 1
}

# Deepin 桌面截屏
deepin_screenshot() {
    # echo "使用 deepin-screenshot"
    local screenshot_tool="deepin-screenshot -s $1 >/dev/null 2>&1"
    if execute_screenshot_cmd_on_wayland "$screenshot_tool"; then
        return 0
    fi
    return 1
}

# MATE 桌面截屏
mate_screenshot() {
    # echo "使用 mate-screenshot"
    local screenshot_tool="mate-screenshot -f $1 >/dev/null 2>&1"
    if execute_screenshot_cmd_on_wayland "$screenshot_tool"; then
        return 0
    fi
    return 1
}

take_screenshot() {
    local de="$1"
    local tool="$2"
    local output_file="$3"
    
    # echo "桌面环境: $de, 工具: $tool"

    case "$tool" in
        "gnome-screenshot")
            if gnome_screenshot "$output_file"; then
                return 0
            else
                return 1
            fi
            ;;
        "kylin-screenshot")
            if kylin_screenshot "$output_file"; then
                return 0
            else
                return 1
            fi
            ;;
        "spectacle"|"ksnapshot")
            if kde_screenshot "$output_file"; then
                return 0
            else
                return 1
            fi
            ;;
        "xfce4-screenshooter")
            if xfce_screenshot "$output_file"; then
                return 0
            else
                return 1
            fi
            ;;
        "scrot")
            if scrot_screenshot "$output_file"; then
                return 0
            else
                return 1
            fi
            ;;
        "maim")
            if maim_screenshot "$output_file"; then
                return 0
            else
                return 1
            fi
            ;;
        "import")
            if import_screenshot "$output_file"; then
                return 0
            else
                return 1
            fi
            ;;
        "deepin-screenshot")
            if deepin_screenshot "$output_file"; then
                return 0
            else
                return 1
            fi
            ;;
        "mate-screenshot")
            if mate_screenshot "$output_file"; then
                return 0
            else
                return 1
            fi
            ;;
        *)
            return 1
            ;;
    esac
}

# 主函数
main() {
    # 开始截屏脚本
    
    # 检测桌面环境
    local de=$(detect_desktop_environment)
    # 检测可用的截屏工具
    local tool=$(detect_screenshot_tool)
    
    if [ "$tool" = "none" ]; then
        # echo "错误: 未找到可用的截屏工具" >&2
        # echo "请安装以下工具之一: gnome-screenshot, spectacle, scrot, maim, ImageMagick" >&2
        echo "FAIL"
        exit 1
    fi
    
    # 尝试截屏
    if take_screenshot "$de" "$tool" "$SCREENSHOT_TARGET_FILE"; then
        echo "SUCCEED"
        chmod 777 $SCREENSHOT_TARGET_FILE
    else
        echo "FAIL"
        exit 1
    fi
}

# 信号处理
cleanup() {
    # echo "脚本被中断"
    echo "FAIL"
    exit 1
}

trap cleanup INT TERM

# 检查是否在图形环境中
if [ -z "$DISPLAY" ] && [ -z "$WAYLAND_DISPLAY" ]; then
    # echo "错误: 未检测到图形环境" >&2
    echo "FAIL"
    exit 1
fi

# 运行主函数
main
exit 0
