#!/usr/bin/env bash
#
# flash_pcan.sh — One-click script to flash canable/candleLight hardware with PCAN-USB compatible firmware
# Supports Ubuntu / macOS, automatically locates firmware file, no manual path modification required.
#
# Usage:
#   ./flash_pcan.sh                     # Auto search firmware file
#   ./flash_pcan.sh /path/to/fw.bin     # Manually specify firmware file path
#
set -euo pipefail


# ---------- Automatic permission & dependency handling (newly added, original flashing logic unchanged) ----------
# First-time recommendation: bash flash_pcan_ubuntu.sh
# The script will automatically obtain execute permission after running, then you can directly use ./flash_pcan_ubuntu.sh

chmod +x "$0" 2>/dev/null || true

# Auto install usbutils(lsusb) for Ubuntu
if command -v apt >/dev/null 2>&1; then
    if ! command -v lsusb >/dev/null 2>&1; then
        echo "[INFO] lsusb not detected, installing usbutils automatically..."
        sudo apt update
        sudo apt install -y usbutils
    fi
fi


# ---------- Basic Configuration ----------
FLASH_ADDR="0x08000000"
DFU_VID_PID="0483:df11"      # VID:PID of STM32 ROM DFU bootloader
TARGET_VID="0c72"            # Expected PEAK-System VID after flashing
FW_NAME_PATTERN="pcan_canable_hw*.bin"

OS="$(uname -s)"

# ---------- Utility Functions ----------
info()  { printf '\033[1;34m[INFO]\033[0m %s\n' "$1"; }
ok()    { printf '\033[1;32m[✓]\033[0m %s\n' "$1"; }
warn()  { printf '\033[1;33m[!]\033[0m %s\n' "$1"; }
err()   { printf '\033[1;31m[✗]\033[0m %s\n' "$1"; }

# ---------- 1. Locate Firmware File ----------
find_firmware() {
    # 1) Command line argument takes priority
    if [[ $# -ge 1 && -n "${1:-}" ]]; then
        if [[ -f "$1" ]]; then
            echo "$1"
            return 0
        else
            err "Specified path does not exist: $1"
            exit 1
        fi
    fi

    # 2) Auto search in common locations (priority ordered)
    local search_dirs=(
        "$(pwd)"
        "$HOME/Downloads"
        "$HOME/Desktop"
        "/mnt/user-data/uploads"
        "/tmp"
    )

    local found=()
    for d in "${search_dirs[@]}"; do
        [[ -d "$d" ]] || continue
        while IFS= read -r -d '' f; do
            found+=("$f")
        done < <(find "$d" -maxdepth 2 -iname "$FW_NAME_PATTERN" -print0 2>/dev/null)
    done

    # Deduplicate entries
    if [[ ${#found[@]} -gt 0 ]]; then
        mapfile -t found < <(printf '%s\n' "${found[@]}" | sort -u)
    fi

    if [[ ${#found[@]} -eq 0 ]]; then
        return 1
    elif [[ ${#found[@]} -eq 1 ]]; then
        echo "${found[0]}"
        return 0
    else
        # Multiple candidates, prompt user to select
        warn "Multiple matching firmware files found, please select:"
        select f in "${found[@]}"; do
            if [[ -n "$f" ]]; then
                echo "$f"
                return 0
            fi
        done
    fi
}

info "Searching for firmware file (pattern: $FW_NAME_PATTERN)..."
FW_BIN="$(find_firmware "${1:-}" || true)"

if [[ -z "${FW_BIN:-}" ]]; then
    err "Firmware file not found via auto search."
    read -rp "Please manually input full path to .bin firmware file: " FW_BIN
    if [[ ! -f "$FW_BIN" ]]; then
        err "Path does not exist: $FW_BIN"
        exit 1
    fi
fi

ok "Using firmware file: $FW_BIN"
FW_SIZE="$(wc -c < "$FW_BIN" | tr -d ' ')"
info "File size: ${FW_SIZE} bytes"
if command -v sha256sum >/dev/null 2>&1; then
    info "SHA256: $(sha256sum "$FW_BIN" | awk '{print $1}')"
elif command -v shasum >/dev/null 2>&1; then
    info "SHA256: $(shasum -a 256 "$FW_BIN" | awk '{print $1}')"
fi

# ---------- 2. Check / Install dfu-util ----------
if ! command -v dfu-util >/dev/null 2>&1; then
    warn "dfu-util not detected, attempting automatic installation..."
    if [[ "$OS" == "Linux" ]]; then
        sudo apt update && sudo apt install -y dfu-util
    elif [[ "$OS" == "Darwin" ]]; then
        if ! command -v brew >/dev/null 2>&1; then
            err "Homebrew not detected, please install it first: https://brew.sh"
            exit 1
        fi
        brew install dfu-util
    else
        err "Unknown operating system: $OS, please install dfu-util manually"
        exit 1
    fi
fi

if ! command -v dfu-util >/dev/null 2>&1; then
    err "dfu-util installation failed or not in PATH, install manually and retry."
    exit 1
fi
ok "dfu-util ready ($(dfu-util --version 2>&1 | head -1))"

# ---------- 3. Detect Device in DFU Mode ----------
check_dfu_device() {
    if [[ "$OS" == "Linux" ]]; then
        command -v lsusb >/dev/null 2>&1 || { warn "lsusb missing, trying to install usbutils..."; sudo apt install -y usbutils; }
        lsusb | grep -qi "$DFU_VID_PID"
    elif [[ "$OS" == "Darwin" ]]; then
        if command -v lsusb >/dev/null 2>&1; then
            lsusb | grep -qi "$DFU_VID_PID"
        else
            system_profiler SPUSBDataType 2>/dev/null | grep -qi "df11\|STM Device in DFU Mode"
        fi
    fi
}

info "Checking for device in DFU mode (${DFU_VID_PID})..."
if ! check_dfu_device; then
    err "No device in DFU mode detected."
    echo "   Please verify:"
    echo "   1. DIP switch on USB2CAN module set to BOOT position"
    echo "   2. Module plugged into host USB port (try re-plugging)"
    if [[ "$OS" == "Darwin" ]] && ! command -v lsusb >/dev/null 2>&1; then
        echo "   3. For macOS, install usbutils for accurate detection: brew install usbutils"
    fi
    exit 1
fi
ok "DFU device detected"

# ---------- 4. List alt-setting for manual verification ----------
info "DFU partition information of current device:"
sudo dfu-util -l || true
echo ""
read -rp "Confirm output contains alt=0 with name 'Internal Flash', press Enter to continue flashing (Ctrl+C to abort)..." _

# ---------- 5. Flash Firmware ----------
info "Start flashing firmware to ${FLASH_ADDR} ..."
sudo dfu-util -a 0 -s "${FLASH_ADDR}:leave" -D "$FW_BIN"

ok "Flashing completed!"
echo ""
echo "Next steps:"
echo "  1. Unplug USB2CAN module"
echo "  2. Toggle DIP switch back to normal operating position (120R)"
echo "  3. Wait 3 seconds then reinsert into USB port"
echo ""
echo "Run the following command to verify flashing result:"
if [[ "$OS" == "Linux" ]]; then
    echo "  lsusb | grep ${TARGET_VID}          # PEAK-System device should appear"
    echo "  sudo modprobe peak_usb"
    echo "  ip -br link                          # can0 interface should appear"
else
    echo "  lsusb | grep ${TARGET_VID}   # Alternative: system_profiler SPUSBDataType | grep -A3 '${TARGET_VID}'"
fi