#!/bin/bash
#
# CUPS 打印服务一键配置脚本
# 适用于 Armbian / Debian / Ubuntu
# 作者: 自动生成
# 用法: chmod +x setup_cups.sh && ./setup_cups.sh
#

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 打印带颜色的信息
info() { echo -e "${BLUE}[信息]${NC} $1"; }
success() { echo -e "${GREEN}[成功]${NC} $1"; }
warn() { echo -e "${YELLOW}[警告]${NC} $1"; }
error() { echo -e "${RED}[错误]${NC} $1"; exit 1; }

# 检查是否为root用户
check_root() {
    if [ "$(id -u)" != "0" ]; then
        error "请使用 root 用户运行此脚本，或使用 sudo"
    fi
}

# 检测系统版本和架构
detect_os() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        OS_ID=$ID
        OS_VERSION=$VERSION_CODENAME
    else
        OS_ID="unknown"
        OS_VERSION="unknown"
    fi

    # 检测CPU架构
    ARCH=$(uname -m)
    case $ARCH in
        x86_64|amd64)
            ARCH_TYPE="amd64"
            ;;
        aarch64|arm64)
            ARCH_TYPE="arm64"
            ;;
        armv7l|armhf)
            ARCH_TYPE="armhf"
            ;;
        *)
            ARCH_TYPE="unknown"
            ;;
    esac
}

# 更换为国内源
change_to_china_mirror() {
    info "检测当前软件源..."

    detect_os

    # 检查源是否正确
    local need_change=0

    if [ "$OS_ID" = "ubuntu" ]; then
        if [ "$ARCH_TYPE" = "amd64" ]; then
            # x86_64 不应该用 ubuntu-ports
            if grep -q "ubuntu-ports" /etc/apt/sources.list 2>/dev/null; then
                warn "检测到错误的源配置 (ubuntu-ports 用于 ARM，当前是 x86_64)"
                need_change=1
            elif grep -q "mirrors.tuna.tsinghua.edu.cn/ubuntu\|mirrors.aliyun.com/ubuntu\|mirrors.ustc.edu.cn/ubuntu" /etc/apt/sources.list 2>/dev/null; then
                success "当前已是正确的国内源，跳过换源"
                return 0
            else
                need_change=1
            fi
        else
            # ARM 应该用 ubuntu-ports
            if grep -q "mirrors.tuna.tsinghua.edu.cn/ubuntu-ports\|mirrors.aliyun.com/ubuntu-ports" /etc/apt/sources.list 2>/dev/null; then
                success "当前已是正确的国内源，跳过换源"
                return 0
            else
                need_change=1
            fi
        fi
    elif [ "$OS_ID" = "debian" ]; then
        if grep -q "mirrors.tuna.tsinghua.edu.cn/debian\|mirrors.aliyun.com/debian\|mirrors.ustc.edu.cn/debian" /etc/apt/sources.list 2>/dev/null; then
            success "当前已是正确的国内源，跳过换源"
            return 0
        else
            need_change=1
        fi
    else
        need_change=1
    fi

    if [ "$need_change" = "0" ]; then
        return 0
    fi

    # 备份原有源
    if [ -f /etc/apt/sources.list ]; then
        cp /etc/apt/sources.list /etc/apt/sources.list.bak
        info "已备份原有源到 /etc/apt/sources.list.bak"
    fi

    # 根据系统类型选择源
    case $OS_ID in
        debian)
            info "检测到 Debian 系统 ($OS_VERSION)，更换为清华源..."
            cat > /etc/apt/sources.list << EOF
# 清华大学 Debian 镜像源
deb https://mirrors.tuna.tsinghua.edu.cn/debian/ $OS_VERSION main contrib non-free non-free-firmware
deb https://mirrors.tuna.tsinghua.edu.cn/debian/ $OS_VERSION-updates main contrib non-free non-free-firmware
deb https://mirrors.tuna.tsinghua.edu.cn/debian/ $OS_VERSION-backports main contrib non-free non-free-firmware
deb https://mirrors.tuna.tsinghua.edu.cn/debian-security $OS_VERSION-security main contrib non-free non-free-firmware
EOF
            ;;
        ubuntu)
            # 根据架构选择正确的源
            if [ "$ARCH_TYPE" = "amd64" ]; then
                UBUNTU_MIRROR="mirrors.tuna.tsinghua.edu.cn/ubuntu"
                info "检测到 Ubuntu 系统 ($OS_VERSION) [x86_64]，更换为清华源..."
            else
                UBUNTU_MIRROR="mirrors.tuna.tsinghua.edu.cn/ubuntu-ports"
                info "检测到 Ubuntu 系统 ($OS_VERSION) [ARM]，更换为清华源..."
            fi
            cat > /etc/apt/sources.list << EOF
# 清华大学 Ubuntu 镜像源
deb https://$UBUNTU_MIRROR/ $OS_VERSION main restricted universe multiverse
deb https://$UBUNTU_MIRROR/ $OS_VERSION-updates main restricted universe multiverse
deb https://$UBUNTU_MIRROR/ $OS_VERSION-backports main restricted universe multiverse
deb https://$UBUNTU_MIRROR/ $OS_VERSION-security main restricted universe multiverse
EOF
            ;;
        *)
            warn "未识别的系统类型: $OS_ID，跳过换源"
            return 1
            ;;
    esac

    # 处理 Armbian 源
    if [ -f /etc/apt/sources.list.d/armbian.list ]; then
        if ! grep -q "mirrors.tuna.tsinghua.edu.cn" /etc/apt/sources.list.d/armbian.list; then
            cp /etc/apt/sources.list.d/armbian.list /etc/apt/sources.list.d/armbian.list.bak
            sed -i 's|apt.armbian.com|mirrors.tuna.tsinghua.edu.cn/armbian|g' /etc/apt/sources.list.d/armbian.list
            info "Armbian 源已更换为清华镜像"
        fi
    fi

    success "软件源已更换为国内镜像（清华大学）"
}

# 获取本机IP地址
get_ip() {
    ip addr show | grep -oP '(?<=inet\s)\d+(\.\d+){3}' | grep -v '127.0.0.1' | head -n1
}

# 显示欢迎信息
show_banner() {
    clear
    echo -e "${GREEN}"
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║         CUPS 打印服务一键配置脚本                        ║"
    echo "║         适用于 Armbian / Debian / Ubuntu                 ║"
    echo "╚══════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
    echo ""
}

# 选择打印机驱动
select_drivers() {
    echo -e "${YELLOW}请选择要安装的打印机驱动 (可多选，用空格或逗号分隔):${NC}"
    echo ""
    echo "  1) 通用驱动 (推荐，支持大多数打印机)"
    echo "  2) HP 打印机 (hplip)"
    echo "  3) 爱普生 (gutenprint)"
    echo "  4) 兄弟 (brlaser)"
    echo "  5) 奔图 Pantum (从官网下载)"
    echo "  6) 全部安装 (不含奔图)"
    echo "  0) 仅安装CUPS (不安装额外驱动)"
    echo ""
    echo -e "  ${BLUE}示例: 1 5 或 1,5 表示同时安装通用驱动和奔图驱动${NC}"
    echo ""
    read -p "请输入选项 (默认: 1): " driver_choice < /dev/tty
    driver_choice=${driver_choice:-1}
    # 统一分隔符
    driver_choice=$(echo "$driver_choice" | tr ',' ' ')
}

# 安装奔图驱动
install_pantum_driver() {
    info "正在安装奔图 Pantum 打印机驱动..."

    # 检测系统架构
    local arch=$(dpkg --print-architecture)
    info "检测到系统架构: $arch"

    # 奔图驱动下载地址 (通用Linux驱动)
    local pantum_url=""
    local pantum_file=""

    case $arch in
        arm64|aarch64)
            pantum_url="https://drivers.pantum.com/userfiles/files/download/drive/Pantum%20Ubuntu%20Driver%20V1_1_123_1.zip"
            ;;
        armhf|arm)
            pantum_url="https://drivers.pantum.com/userfiles/files/download/drive/Pantum%20Ubuntu%20Driver%20V1_1_123_1.zip"
            ;;
        amd64|x86_64)
            pantum_url="https://drivers.pantum.com/userfiles/files/download/drive/Pantum%20Ubuntu%20Driver%20V1_1_123_1.zip"
            ;;
        *)
            warn "不支持的架构: $arch，跳过奔图驱动安装"
            return 1
            ;;
    esac

    # 创建临时目录
    local tmp_dir=$(mktemp -d)
    cd "$tmp_dir"

    info "下载奔图驱动..."
    if ! wget -q --show-progress -O pantum_driver.zip "$pantum_url"; then
        warn "下载奔图驱动失败，尝试备用方案..."
        # 备用: 直接安装依赖，用户手动下载驱动
        apt install -y libcups2 libcupsimage2
        warn "请手动从 https://www.pantum.com 下载驱动并安装"
        cd - > /dev/null
        rm -rf "$tmp_dir"
        return 1
    fi

    info "解压驱动包..."
    unzip -q pantum_driver.zip

    # 查找并安装deb包
    local deb_file=""
    if [ "$arch" = "arm64" ] || [ "$arch" = "aarch64" ]; then
        deb_file=$(find . -name "*arm64*.deb" -o -name "*aarch64*.deb" 2>/dev/null | head -1)
    elif [ "$arch" = "armhf" ] || [ "$arch" = "arm" ]; then
        deb_file=$(find . -name "*armhf*.deb" -o -name "*arm*.deb" 2>/dev/null | grep -v arm64 | head -1)
    else
        deb_file=$(find . -name "*amd64*.deb" -o -name "*x86_64*.deb" 2>/dev/null | head -1)
    fi

    # 如果找不到对应架构，尝试通用包
    if [ -z "$deb_file" ]; then
        deb_file=$(find . -name "*.deb" 2>/dev/null | head -1)
    fi

    if [ -n "$deb_file" ]; then
        info "安装驱动包: $deb_file"
        dpkg -i "$deb_file" || apt install -f -y
        success "奔图驱动安装完成"
    else
        warn "未找到适合的驱动包"
        warn "请手动从 https://www.pantum.com 下载驱动"
    fi

    # 清理
    cd - > /dev/null
    rm -rf "$tmp_dir"
}

# 安装驱动
install_drivers() {
    local packages=""
    local install_pantum=false

    for choice in $driver_choice; do
        case $choice in
            1)
                info "添加通用驱动..."
                packages="$packages printer-driver-all"
                ;;
            2)
                info "添加 HP 驱动..."
                packages="$packages hplip"
                ;;
            3)
                info "添加爱普生驱动..."
                packages="$packages printer-driver-gutenprint"
                ;;
            4)
                info "添加兄弟驱动..."
                packages="$packages printer-driver-brlaser"
                ;;
            5)
                install_pantum=true
                ;;
            6)
                info "添加全部驱动..."
                packages="printer-driver-all hplip printer-driver-gutenprint printer-driver-brlaser"
                ;;
            0)
                info "跳过驱动安装"
                return 0
                ;;
        esac
    done

    if [ -n "$packages" ]; then
        info "正在安装驱动: $packages"
        apt install -y $packages
    fi

    # 安装奔图驱动
    if [ "$install_pantum" = true ]; then
        install_pantum_driver
    fi
}

# 配置CUPS
configure_cups() {
    info "配置 CUPS 允许远程访问..."

    # 备份原配置
    cp /etc/cups/cupsd.conf /etc/cups/cupsd.conf.bak

    # 修改监听地址
    sed -i 's/^Listen localhost:631$/Listen 0.0.0.0:631/' /etc/cups/cupsd.conf

    # 如果没有找到 Listen localhost:631，则添加
    if ! grep -q "Listen 0.0.0.0:631" /etc/cups/cupsd.conf; then
        sed -i '/^Listen/d' /etc/cups/cupsd.conf
        sed -i '1a Listen 0.0.0.0:631\nListen /run/cups/cups.sock' /etc/cups/cupsd.conf
    fi

    # 添加 ServerAlias * 允许任何主机名访问
    if ! grep -q "^ServerAlias" /etc/cups/cupsd.conf; then
        sed -i '/^Listen.*631/a ServerAlias *' /etc/cups/cupsd.conf
    fi

    # 禁用强制加密（允许 HTTP 访问）
    if ! grep -q "^DefaultEncryption" /etc/cups/cupsd.conf; then
        sed -i '/^ServerAlias/a DefaultEncryption Never' /etc/cups/cupsd.conf
    fi

    # 启用网络浏览
    if grep -q "^Browsing" /etc/cups/cupsd.conf; then
        sed -i 's/^Browsing.*/Browsing Yes/' /etc/cups/cupsd.conf
    else
        echo "Browsing Yes" >> /etc/cups/cupsd.conf
    fi

    # 修改 Policy default 允许匿名打印
    # 找到 <Limit Create-Job Print-Job Print-URI Validate-Job> 块并修改
    sed -i '/<Limit Create-Job Print-Job Print-URI Validate-Job>/,/<\/Limit>/{
        s/Order deny,allow/Order allow,deny/
        /Require user/d
        /Order/a\    Allow all
    }' /etc/cups/cupsd.conf

    # 配置访问权限 - 允许所有网络访问
    # 先删除现有的 Location 块，然后重新添加
    cat > /tmp/cups_locations.conf << 'EOF'

# 允许所有网络访问（内网环境使用）
<Location />
  Order allow,deny
  Allow all
</Location>

<Location /admin>
  Order allow,deny
  Allow all
</Location>

<Location /admin/conf>
  AuthType Default
  Require user @SYSTEM
  Order allow,deny
  Allow all
</Location>

# 允许远程打印
<Location /printers>
  Order allow,deny
  Allow all
</Location>

<Location /printers/*>
  Order allow,deny
  Allow all
</Location>
EOF

    # 移除旧的 Location 块
    sed -i '/<Location/,/<\/Location>/d' /etc/cups/cupsd.conf

    # 添加新的 Location 配置
    cat /tmp/cups_locations.conf >> /etc/cups/cupsd.conf
    rm /tmp/cups_locations.conf

    success "CUPS 配置完成"
}

# 主函数
main() {
    show_banner
    check_root

    LOCAL_IP=$(get_ip)
    info "检测到本机IP: ${LOCAL_IP}"
    echo ""

    # 选择驱动
    select_drivers
    echo ""

    # 询问是否换源
    echo -e "${YELLOW}是否将软件源更换为国内镜像（清华大学）？${NC}"
    echo "  推荐：国内用户选择 Y，可大幅加快下载速度"
    echo ""
    read -p "更换国内源? [Y/n]: " change_mirror < /dev/tty
    change_mirror=${change_mirror:-Y}
    echo ""

    # 确认安装
    echo -e "${YELLOW}即将执行以下操作:${NC}"
    if [[ "$change_mirror" =~ ^[Yy]$ ]]; then
        echo "  1. 更换为国内软件源（清华大学）"
        echo "  2. 更新系统软件源"
    else
        echo "  1. 更新系统软件源"
    fi
    echo "  3. 安装 CUPS 和 Avahi"
    echo "  4. 安装打印机驱动"
    echo "  5. 配置 CUPS 允许局域网访问"
    echo "  6. 将当前用户添加到 lpadmin 组"
    echo "  7. 启动并设置开机自启"
    echo ""
    read -p "是否继续? [Y/n]: " confirm < /dev/tty
    confirm=${confirm:-Y}

    if [[ ! "$confirm" =~ ^[Yy]$ ]]; then
        echo "已取消"
        exit 0
    fi

    echo ""
    info "开始安装..."
    echo ""

    # 1. 换源（如果选择）
    if [[ "$change_mirror" =~ ^[Yy]$ ]]; then
        change_to_china_mirror
    fi

    # 2. 更新系统
    info "更新软件源..."
    apt update
    success "软件源更新完成"

    # 2. 安装 CUPS 和 Avahi
    info "安装 CUPS 和 Avahi..."
    apt install -y cups avahi-daemon wget unzip
    success "CUPS 和 Avahi 安装完成"

    # 3. 安装中文语言包（CUPS汉化）
    info "安装中文语言包..."
    apt install -y locales language-pack-zh-hans 2>/dev/null || apt install -y locales 2>/dev/null
    # 生成中文 locale
    sed -i 's/# zh_CN.UTF-8 UTF-8/zh_CN.UTF-8 UTF-8/' /etc/locale.gen 2>/dev/null
    locale-gen zh_CN.UTF-8 2>/dev/null
    # 设置 CUPS 默认语言为中文
    if ! grep -q "DefaultLanguage" /etc/cups/cupsd.conf 2>/dev/null; then
        echo 'DefaultLanguage zh_CN.UTF-8' >> /etc/cups/cupsd.conf
    fi

    # 安装中文 Web 界面模板
    SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
    if [ -d "$SCRIPT_DIR/cups-templates-zh_CN" ]; then
        info "安装 CUPS 中文界面模板..."
        mkdir -p /usr/share/cups/templates-zh_CN
        cp -f "$SCRIPT_DIR/cups-templates-zh_CN/"*.tmpl /usr/share/cups/templates-zh_CN/
        success "中文界面模板安装完成"
    else
        warn "未找到中文模板目录 (cups-templates-zh_CN)，跳过界面汉化"
    fi
    success "中文语言包安装完成"

    # 3. 安装驱动
    install_drivers
    success "驱动安装完成"

    # 4. 配置 CUPS
    configure_cups

    # 5. 添加用户到 lpadmin 组
    info "将 root 用户添加到 lpadmin 组..."
    usermod -aG lpadmin root

    # 如果有其他用户，也添加
    CURRENT_USER=$(logname 2>/dev/null || echo "")
    if [ -n "$CURRENT_USER" ] && [ "$CURRENT_USER" != "root" ]; then
        usermod -aG lpadmin "$CURRENT_USER"
        info "将 $CURRENT_USER 用户添加到 lpadmin 组"
    fi
    success "用户组配置完成"

    # 6. 启动服务
    info "启动 CUPS 服务..."
    systemctl restart cups
    systemctl enable cups
    success "CUPS 服务已启动并设置开机自启"

    info "启动 Avahi 服务..."
    systemctl restart avahi-daemon
    systemctl enable avahi-daemon
    success "Avahi 服务已启动并设置开机自启"

    # 完成提示
    echo ""
    echo -e "${GREEN}╔══════════════════════════════════════════════════════════╗${NC}"
    echo -e "${GREEN}║                    安装完成!                             ║${NC}"
    echo -e "${GREEN}╚══════════════════════════════════════════════════════════╝${NC}"
    echo ""
    echo -e "  ${BLUE}CUPS 管理界面:${NC} https://${LOCAL_IP}:631"
    echo -e "  ${BLUE}用户名:${NC} root"
    echo -e "  ${BLUE}密码:${NC} 你的 SSH 密码"
    echo ""
    echo -e "  ${YELLOW}下一步:${NC}"
    echo "  1. 用浏览器打开上面的地址"
    echo "  2. 点击 Administration → Add Printer"
    echo "  3. 选择你的 USB 打印机"
    echo "  4. 勾选 Share This Printer"
    echo ""
    echo -e "  ${YELLOW}客户端连接:${NC}"
    echo "  - Windows: 设置 → 打印机 → 添加打印机"
    echo "  - macOS/iOS: 自动通过 AirPrint 发现"
    echo "  - Android: 安装 Mopria Print Service"
    echo ""

    # 检查打印机
    echo -e "${BLUE}[检测]${NC} 正在检测 USB 打印机..."
    if lsusb | grep -i -E "printer|print" > /dev/null 2>&1; then
        success "检测到 USB 打印机:"
        lsusb | grep -i -E "printer|print"
    else
        warn "未检测到 USB 打印机，请确认打印机已连接"
    fi
    echo ""
}

# 运行主函数
main