#!/bin/bash
#
# usb-gadget.sh
# Raspberry Pi 4B USB Gadget - 将 USB-A 口上的大容量存储设备通过 USB-C 口透传给 Windows PC。
#
# 用法: usb-gadget.sh {start|stop|restart|status}
#

set -uo pipefail

# ===================== 可配置项 =====================
# 工作模式：
#   auto   - 自动检测第一个 USB 总线接入的 U 盘/移动硬盘（默认，满足“U 盘透传”需求）
#   device - 使用下面 USB_DISK 变量指定的块设备
#   file   - 使用 /piusb.bin 文件作为 backing（兼容旧脚本行为）
USB_DISK_MODE="${USB_DISK_MODE:-auto}"

# 当 USB_DISK_MODE=device 时使用；例如 /dev/sda（整盘）或 /dev/sda1（单个分区）
USB_DISK="${USB_DISK:-}"

# 当 USB_DISK_MODE=file 时使用
BACKING_FILE="${BACKING_FILE:-/piusb.bin}"
BACKING_SIZE_MB="${BACKING_SIZE_MB:-1024}"

# 是否只读：0=可读写，1=只读
READ_ONLY="${READ_ONLY:-0}"

# Windows 设备管理器中显示的信息
ID_VENDOR="${ID_VENDOR:-0x1d6b}"       # Linux Foundation
ID_PRODUCT="${ID_PRODUCT:-0x0104}"     # Multifunction Composite Gadget
SERIAL="${SERIAL:-1234567890ABCDEF}"
MANUFACTURER="${MANUFACTURER:-DeskPi}"
PRODUCT_NAME="${PRODUCT_NAME:-Shared UDisk}"

# Gadget 资源名
GADGET_NAME="${GADGET_NAME:-g1}"
CONFIG_NAME="${CONFIG_NAME:-c.1}"
FUNCTION_NAME="${FUNCTION_NAME:-mass_storage.usb0}"
LUN_NAME="${LUN_NAME:-lun.0}"

# MaxPower 单位是 2mA，250 = 500mA
MAX_POWER="${MAX_POWER:-250}"
# bmAttributes: 0x80 = bus-powered, 0xC0 = self-powered
BM_ATTRIBUTES="${BM_ATTRIBUTES:-0x80}"

# ===================== 工具函数 =====================

err() { echo "[ERROR] $*" >&2; }
info() { echo "[INFO] $*"; }
warn() { echo "[WARN] $*"; }

# 查找 USB 大容量存储块设备（返回整盘路径），排除 SD 卡/loop/zram 等
find_usb_disk() {
    local dev

    # 方法1：lsblk 识别 TRAN=usb 的块设备（最可靠）
    dev=$(lsblk -d -n -p -o NAME,TRAN 2>/dev/null \
          | awk '$2 == "usb" { print $1; exit }')

    if [ -n "$dev" ]; then
        echo "$dev"
        return 0
    fi

    # 方法2：/dev/disk/by-path 中匹配 usb 的符号链接
    if [ -d /dev/disk/by-path ]; then
        dev=$(find /dev/disk/by-path -maxdepth 1 -type l \
              -name '*usb*' -printf '%l\n' 2>/dev/null \
              | sed 's|^\.\./\.\./|/dev/|' | head -n1)
        if [ -n "$dev" ]; then
            echo "$dev"
            return 0
        fi
    fi

    return 1
}

# 卸载指定块设备的所有分区，防止 Pi 与 Windows 同时写入导致文件系统损坏
unmount_disk() {
    local dev="$1"
    local mp

    # 先同步缓存，再尝试卸载
    sync
    sleep 0.2

    mount | awk -v dev="$dev" '$1 ~ ("^" dev) { print $1 }' \
          | while read -r mp; do
        info "Unmounting $mp from Pi side ..."
        umount "$mp" 2>/dev/null || umount -l "$mp" 2>/dev/null || true
    done

    # 确认已卸载干净；若仍有挂载，给出警告但继续（避免阻塞 gadget 启动）
    local still_mounted
    still_mounted=$(mount | awk -v dev="$dev" '$1 ~ ("^" dev) { print $1 }' | wc -l)
    if [ "$still_mounted" -gt 0 ]; then
        warn "Some partitions of $dev are still mounted; Windows and Pi may both access the disk."
    fi
}

# 检查是否有进程正在占用该设备
is_disk_busy() {
    local dev="$1"
    if command -v fuser >/dev/null 2>&1; then
        if fuser "$dev" >/dev/null 2>&1; then
            return 0
        fi
    fi
    # 对分区也做检查
    local part
    for part in "${dev}"*; do
        [ -b "$part" ] || continue
        if command -v fuser >/dev/null 2>&1 && fuser "$part" >/dev/null 2>&1; then
            return 0
        fi
    done
    return 1
}

# 检查设备是否是系统盘，防止误暴露根分区/SD 卡
is_system_disk() {
    local dev="$1"
    local root_dev root_disk dev_disk

    root_dev=$(findmnt -n -o SOURCE / 2>/dev/null || true)
    [ -z "$root_dev" ] && return 1

    # 把 UUID/PARTUUID 等解析成真实块设备
    [ -b "$root_dev" ] || root_dev=$(readlink -f "$root_dev" 2>/dev/null || true)
    [ -b "$root_dev" ] || return 1

    # 获取根分区所在的整盘
    root_disk=$(lsblk -n -o PKNAME "$root_dev" 2>/dev/null | head -n1)
    # 如果 PKNAME 为空（根直接在整盘上），则使用设备自身
    [ -z "$root_disk" ] && root_disk=$(lsblk -d -n -o NAME "$root_dev" 2>/dev/null)

    dev_disk=$(lsblk -d -n -o NAME "$dev" 2>/dev/null)

    [ -n "$root_disk" ] && [ -n "$dev_disk" ] && [ "$root_disk" = "$dev_disk" ]
}

# 根据模式解析要使用的 backing 设备/文件
resolve_backing() {
    case "$USB_DISK_MODE" in
        auto)
            find_usb_disk
            ;;
        device)
            if [ -b "$USB_DISK" ]; then
                echo "$USB_DISK"
            else
                err "USB_DISK=$USB_DISK is not a block device"
                return 1
            fi
            ;;
        file)
            echo "$BACKING_FILE"
            ;;
        *)
            err "Unknown USB_DISK_MODE: $USB_DISK_MODE"
            return 1
            ;;
    esac
}

# file 模式下，若 backing 文件不存在则创建并格式化为 FAT32
ensure_backing_file() {
    if [ ! -f "$BACKING_FILE" ]; then
        info "Creating backing file $BACKING_FILE (${BACKING_SIZE_MB}MB) ..."
        dd if=/dev/zero of="$BACKING_FILE" bs=1M count="$BACKING_SIZE_MB" status=progress
        mkfs.vfat -F 32 "$BACKING_FILE"
    fi
}

# ===================== Gadget 生命周期 =====================

gadget_path() {
    echo "/sys/kernel/config/usb_gadget/$GADGET_NAME"
}

udc_name() {
    ls /sys/class/udc/ | head -n1
}

load_modules() {
    if ! modprobe libcomposite; then
        err "Failed to load libcomposite module. Check that USB gadget support is enabled in kernel."
        return 1
    fi
    if ! modprobe usb_f_mass_storage; then
        err "Failed to load usb_f_mass_storage module."
        return 1
    fi
}

stop_gadget() {
    local gp
    gp=$(gadget_path)

    if [ ! -d "$gp" ]; then
        info "Gadget not active."
        return 0
    fi

    info "Stopping gadget ..."
    # 1. 先解绑 UDC，否则无法删除 function
    echo "" > "$gp/UDC" 2>/dev/null || true
    sleep 0.5

    # 2. 删除 function 到 config 的符号链接
    rm -f "$gp/configs/$CONFIG_NAME/$FUNCTION_NAME"

    # 3. 删除 LUN 与 function
    rmdir "$gp/functions/$FUNCTION_NAME/$LUN_NAME" 2>/dev/null || true
    rmdir "$gp/functions/$FUNCTION_NAME" 2>/dev/null || true

    # 4. 删除 config 字符串与 config
    rmdir "$gp/configs/$CONFIG_NAME/strings/0x409" 2>/dev/null || true
    rmdir "$gp/configs/$CONFIG_NAME" 2>/dev/null || true

    # 5. 删除字符串
    rmdir "$gp/strings/0x409" 2>/dev/null || true

    # 6. 删除 gadget
    rmdir "$gp" 2>/dev/null || true

    info "Gadget stopped."
}

start_gadget() {
    local gp udc disk

    gp=$(gadget_path)

    # 如果已存在，先停止（实现幂等）
    [ -d "$gp" ] && stop_gadget

    # 解析 backing
    disk=$(resolve_backing)
    if [ -z "$disk" ]; then
        info "No USB disk detected. Gadget will not be activated."
        info "Plug a USB drive into the USB-A port, then run: systemctl restart usb-gadget"
        return 0
    fi

    if [ "$USB_DISK_MODE" != "file" ]; then
        if is_system_disk "$disk"; then
            err "Refusing to expose system disk $disk"
            return 1
        fi
        if is_disk_busy "$disk"; then
            warn "Disk $disk appears to be in use by another process."
        fi
        unmount_disk "$disk"
    else
        ensure_backing_file
    fi

    info "Using backing storage: $disk"

    load_modules || return 1

    info "Creating USB gadget '$GADGET_NAME' ..."
    mkdir -p "$gp"

    # 设备描述符
    echo "$ID_VENDOR"   > "$gp/idVendor"
    echo "$ID_PRODUCT"  > "$gp/idProduct"
    echo 0x0100         > "$gp/bcdDevice"      # 设备版本 v1.0.0
    echo 0x0200         > "$gp/bcdUSB"         # USB 2.0

    # 字符串描述符
    mkdir -p "$gp/strings/0x409"
    echo "$SERIAL"       > "$gp/strings/0x409/serialnumber"
    echo "$MANUFACTURER" > "$gp/strings/0x409/manufacturer"
    echo "$PRODUCT_NAME" > "$gp/strings/0x409/product"

    # 配置描述符
    mkdir -p "$gp/configs/$CONFIG_NAME"
    mkdir -p "$gp/configs/$CONFIG_NAME/strings/0x409"
    echo "Mass Storage" > "$gp/configs/$CONFIG_NAME/strings/0x409/configuration"
    echo "$MAX_POWER"   > "$gp/configs/$CONFIG_NAME/MaxPower"
    echo "$BM_ATTRIBUTES" > "$gp/configs/$CONFIG_NAME/bmAttributes"

    # Mass Storage Function
    mkdir -p "$gp/functions/$FUNCTION_NAME"

    # Windows 兼容性：开启 STALL，提高大容量存储在 Windows 下的稳定性
    echo 1 > "$gp/functions/$FUNCTION_NAME/stall"

    mkdir -p "$gp/functions/$FUNCTION_NAME/$LUN_NAME"

    # 必须先设置 LUN 属性，最后再绑定 file；否则 file 绑定后部分属性会变为只读/EBUSY
    echo 1            > "$gp/functions/$FUNCTION_NAME/$LUN_NAME/removable"
    echo "$READ_ONLY" > "$gp/functions/$FUNCTION_NAME/$LUN_NAME/ro"
    echo 0            > "$gp/functions/$FUNCTION_NAME/$LUN_NAME/cdrom"
    echo "$disk"      > "$gp/functions/$FUNCTION_NAME/$LUN_NAME/file"

    # 绑定 function 到 config
    ln -s "$gp/functions/$FUNCTION_NAME" "$gp/configs/$CONFIG_NAME/"

    # 激活 UDC
    udc=$(udc_name)
    if [ -z "$udc" ]; then
        err "No USB Device Controller (UDC) found."
        err "Make sure 'dtoverlay=dwc2,dr_mode=peripheral' is set in /boot/firmware/config.txt"
        err "and the Pi 4B USB-C port is wired for device/peripheral mode."
        stop_gadget
        return 1
    fi

    echo "$udc" > "$gp/UDC"
    info "OK: Gadget '$PRODUCT_NAME' bound to UDC $udc with backing $disk"
    info "Connect USB-C to Windows PC; the drive should appear shortly."
}

status_gadget() {
    local gp udc disk ro
    gp=$(gadget_path)

    if [ -d "$gp" ]; then
        udc=$(cat "$gp/UDC" 2>/dev/null || true)
        if [ -n "$udc" ]; then
            info "Gadget is active on UDC: $udc"
            disk=$(cat "$gp/functions/$FUNCTION_NAME/$LUN_NAME/file" 2>/dev/null || true)
            ro=$(cat "$gp/functions/$FUNCTION_NAME/$LUN_NAME/ro" 2>/dev/null || true)
            info "  Backing: $disk (ro=$ro)"
        else
            info "Gadget directory exists but is not bound to a UDC."
        fi
    else
        info "Gadget is not active."
    fi
}

# ===================== 主入口 =====================

case "${1:-start}" in
    start)
        start_gadget
        ;;
    stop)
        stop_gadget
        ;;
    restart|reload)
        stop_gadget
        start_gadget
        ;;
    status)
        status_gadget
        ;;
    *)
        echo "Usage: $0 {start|stop|restart|status}"
        exit 1
        ;;
esac
