#!/usr/bin/env bash
# fix-screen.sh — 修复 GNU screen 颜色、编码和鼠标滚轮配置
# 支持 --dry-run 预览模式，不会修改任何文件

set -euo pipefail

# ── 常量 ────────────────────────────────────────────────────────────────────
SCREENRC="$HOME/.screenrc"
BASHRC="$HOME/.bashrc"
BACKUP_SUFFIX="$(date +%Y%m%d-%H%M%S)"
BACKUP_KEEP=3          # 每个文件最多保留的备份数
BASHRC_BEGIN='# BEGIN screen-fix'
BASHRC_END='# END screen-fix'

# ── dry-run 开关 ─────────────────────────────────────────────────────────────
DRY_RUN=0
for arg in "$@"; do
    case "$arg" in
        --dry-run|-n) DRY_RUN=1 ;;
        --help|-h)
            printf 'Usage: %s [--dry-run]\n' "$(basename "$0")"
            printf '  --dry-run  预览将要做的变更，不修改任何文件\n'
            exit 0
            ;;
        *)
            printf '未知参数: %s\n' "$arg" >&2
            exit 1
            ;;
    esac
done

# ── 辅助函数 ─────────────────────────────────────────────────────────────────

# 打印信息（dry-run 时加前缀）
log() {
    if [ "$DRY_RUN" -eq 1 ]; then
        printf '[dry-run] %s\n' "$*"
    else
        printf '%s\n' "$*"
    fi
}

warn() { printf '[警告] %s\n' "$*" >&2; }

die() { printf '[错误] %s\n' "$*" >&2; exit 1; }

# 备份文件，并清理超过 BACKUP_KEEP 份的旧备份（B + A）
backup_if_needed() {
    local file="$1"
    local changed="$2"   # 1 = 本次会修改该文件
    [ -f "$file" ] || return 0
    [ "$changed" -eq 1 ] || return 0

    if [ "$DRY_RUN" -eq 0 ]; then
        cp -p "$file" "$file.bak.$BACKUP_SUFFIX" \
            || die "备份失败: $file"
        log "已备份 $file -> $file.bak.$BACKUP_SUFFIX"

        # 只保留最近 BACKUP_KEEP 份，删除更旧的（A）
        local pattern="${file}.bak.*"
        local count
        count=$(ls -1 ${pattern} 2>/dev/null | wc -l)
        if [ "$count" -gt "$BACKUP_KEEP" ]; then
            ls -1t ${pattern} 2>/dev/null \
                | tail -n +"$((BACKUP_KEEP + 1))" \
                | xargs rm -f --
            log "旧备份已清理，仅保留最近 ${BACKUP_KEEP} 份"
        fi
    else
        log "将备份 $file -> $file.bak.<时间戳>"
    fi
}

# 检查 screenrc 中某条指令是否需要更新（C）
# 返回 0 = 需要变更，1 = 已是目标值无需变更
screenrc_needs_change() {
    local keyword="$1"   # 指令名，如 term / defscrollback
    local target="$2"    # 目标完整行，如 'term "screen-256color"'
    [ -f "$SCREENRC" ] || return 0
    # 如果精确匹配目标行，无需变更
    grep -Fqx -- "$target" "$SCREENRC" && return 1
    return 0
}

# 设置 screenrc 中某条指令为目标值（C）
# 如有旧值先删除再追加，避免冲突
set_screenrc_line() {
    local keyword="$1"
    local target="$2"
    if [ "$DRY_RUN" -eq 1 ]; then
        log "  .screenrc: 设置 '$target'"
        return
    fi
    # 删除同指令的所有旧行（包括带不同参数的）
    sed -i "/^[[:space:]]*${keyword}[[:space:]]/d" "$SCREENRC"
    printf '%s\n' "$target" >> "$SCREENRC"
}

# 检查 .bashrc 的 screen-fix 标记块是否需要更新（E）
bashrc_block_needs_change() {
    [ -f "$BASHRC" ] || return 0
    # 如果标记块已存在且内容完全一致则无需变更
    grep -Fq "$BASHRC_BEGIN" "$BASHRC" && return 1
    return 0
}

# 写入 .bashrc 标记块（E）
write_bashrc_block() {
    if [ "$DRY_RUN" -eq 1 ]; then
        log "  .bashrc: 写入 screen-fix 标记块（彩色提示符支持）"
        return
    fi
    # 先移除旧的 Debian 单行 patch（若存在）
    sed -i 's/xterm-color|xterm-256color|screen|screen-256color) color_prompt=yes;;/xterm-color) color_prompt=yes;;/' \
        "$BASHRC" 2>/dev/null || true
    # 移除已有标记块（幂等）
    sed -i "/${BASHRC_BEGIN}/,/${BASHRC_END}/d" "$BASHRC"
    # 追加新标记块
    cat >> "$BASHRC" <<'EOF'
# BEGIN screen-fix
# 在 screen / xterm-256color 内启用彩色提示符
case "$TERM" in
    xterm-256color|screen|screen-256color) color_prompt=yes ;;
esac
# END screen-fix
EOF
}

# 向活跃 screen 会话应用单条命令，失败时警告（F）
apply_to_session() {
    local session="$1"
    local cmd="$2"
    if ! screen -S "$session" -X $cmd >/dev/null 2>&1; then
        warn "会话 $session: 无法执行 '$cmd'，请手动运行"
    fi
}

# ── 检测哪些文件需要变更（B）────────────────────────────────────────────────
touch "$SCREENRC" "$BASHRC"   # 确保文件存在

screenrc_dirty=0
declare -A SCREENRC_SETTINGS=(
    [term]='term "screen-256color"'
    [defbce]='defbce on'
    [defutf8]='defutf8 on'
    [utf8]='utf8 on'
    [defscrollback]='defscrollback 10000'
    [defmousetrack]='defmousetrack on'
    [mousetrack]='mousetrack on'
)
for kw in term defbce defutf8 utf8 defscrollback defmousetrack mousetrack; do
    if screenrc_needs_change "$kw" "${SCREENRC_SETTINGS[$kw]}"; then
        screenrc_dirty=1
        break
    fi
done

bashrc_dirty=0
bashrc_block_needs_change && bashrc_dirty=1

# ── 主流程 ────────────────────────────────────────────────────────────────────
printf '正在修复 GNU screen 颜色、编码和鼠标滚轮配置'
[ "$DRY_RUN" -eq 1 ] && printf '（dry-run 预览模式）'
printf '...\n\n'

# 备份（仅在有实际变更时）
backup_if_needed "$SCREENRC" "$screenrc_dirty"
backup_if_needed "$BASHRC"   "$bashrc_dirty"

# .screenrc 变更
if [ "$screenrc_dirty" -eq 1 ]; then
    log "[.screenrc] 应用配置："
    for kw in term defbce defutf8 utf8 defscrollback defmousetrack mousetrack; do
        set_screenrc_line "$kw" "${SCREENRC_SETTINGS[$kw]}"
    done
else
    log "[.screenrc] 已是最新，无需修改"
fi

# .bashrc 变更
if [ "$bashrc_dirty" -eq 1 ]; then
    log "[.bashrc] 写入彩色提示符支持块"
    write_bashrc_block
else
    log "[.bashrc] 已是最新，无需修改"
fi

# 将设置应用到当前运行中的 screen 会话（F）
if command -v screen >/dev/null 2>&1; then
    session_count=0
    while read -r session rest; do
        case "$session" in
            [0-9]*.*)
                session_count=$((session_count + 1))
                if [ "$DRY_RUN" -eq 1 ]; then
                    log "[会话] 将应用到: $session"
                else
                    apply_to_session "$session" "mousetrack on"
                    apply_to_session "$session" "scrollback 10000"
                    log "[会话] 已应用到: $session"
                fi
                ;;
        esac
    done < <(screen -ls 2>/dev/null)
    [ "$session_count" -eq 0 ] && log "[会话] 无运行中的 screen 会话"
fi

printf '\n修复完成。新建 screen 会话会自动生效。\n'
[ "$DRY_RUN" -eq 0 ] && printf '当前 shell 请执行: source ~/.bashrc\n'