#!/bin/bash

# -------------------------------
# adbdump.sh - Unified partition dumper
# Supports legacy, new, and dynamic/mapper layouts
# -------------------------------

# Default partition base (new layout)
PART_BASE="/dev/block/by-name"

# Parse flags and remove them from arguments
NEW_ARGS=()
for arg in "$@"; do
    case "$arg" in
        --legacy)
            PART_BASE="/dev/block/bootdevice/by-name"
            ;;
        --mapper)
            PART_BASE="/dev/block/mapper"
            ;;
        *)
            NEW_ARGS+=("$arg")
            ;;
    esac
done
set -- "${NEW_ARGS[@]}"

# -------------------------------
# Helper functions
# -------------------------------

get_single_device() {
    devices=($(adb devices | grep -E "device$|recovery$" | awk '{print $1}'))
    if [ ${#devices[@]} -eq 1 ]; then
        echo "${devices[0]}"
    else
        echo ""
    fi
}

get_partition_path() {
    local part="$1"
    if [[ "$part" == */* ]]; then
        echo "$part"
    else
        echo "$PART_BASE/$part"
    fi
}

# -------------------------------
# Help text (always works)
# -------------------------------
if [ "$1" == "-h" ] || [ "$1" == "help" ] || [ "$1" == "usage" ]; then
    echo "Usage:"
    echo "  $0 [-ld]                          List connected devices"
    echo "  $0 [-lp [serial] <partition>]    List partitions on device"
    echo "  $0 [-S [serial] <partition>]     Save partition as .img"
    echo "  $0 [-fd [serial]]                Full dump all partitions"
    echo "Options:"
    echo "  --legacy                          Use legacy partition layout (/dev/block/bootdevice/by-name)"
    echo "  --mapper                          Use dynamic partitions (/dev/block/mapper)"
    echo "  -h, help, usage                   Show this help message"
    exit 0
fi

OPTION="$1"
shift  # remove option

# -------------------------------
# Main command handling
# -------------------------------
case "$OPTION" in
    -ld)
        adb devices
        exit 0
        ;;

    -lp)
        serial="$1"
        partition="$2"

        if [ -z "$serial" ]; then
            serial=$(get_single_device)
            if [ -z "$serial" ]; then
                echo "Error: Serial number required because multiple or no devices connected."
                exit 1
            else
                echo "Using device: $serial"
            fi
        fi

        if [ -z "$partition" ]; then
            echo "Error: Partition name required."
            exit 1
        fi

        partition_path=$(get_partition_path "$partition")
        adb -s "$serial" shell "ls $partition_path"
        exit 0
        ;;

    -S)
        if [ $# -eq 1 ]; then
            serial=$(get_single_device)
            if [ -z "$serial" ]; then
                echo "Error: Serial number required because multiple or no devices connected."
                exit 1
            fi
            partition="$1"
        elif [ $# -eq 2 ]; then
            serial="$1"
            partition="$2"
        else
            echo "Usage: $0 -S [serial] <partition>"
            exit 1
        fi

        echo "Using device: $serial"
        partition_path=$(get_partition_path "$partition")
        outfile="$partition.img"

        if command -v pv >/dev/null 2>&1; then
            echo "Using pv to show progress..."
            adb -s "$serial" shell "dd if=$partition_path" | pv > "$outfile"
        else
            echo "pv not found, dumping without progress..."
            adb -s "$serial" shell "dd if=$partition_path" > "$outfile"
        fi
        exit 0
        ;;

    -fd)
        if [ -z "$1" ]; then
            serial=$(get_single_device)
            if [ -z "$serial" ]; then
                echo "Error: Serial number required because multiple or no devices connected."
                exit 1
            fi
        else
            serial="$1"
        fi

        echo "Using device: $serial"
        echo "Listing partitions..."

        partitions=$(adb -s "$serial" shell "ls $PART_BASE" | tr -d '\r')

        for partition in $partitions; do
            echo "Dumping $partition..."
            partition_path=$(get_partition_path "$partition")
            outfile="${partition}.img"

            if command -v pv >/dev/null 2>&1; then
                adb -s "$serial" shell "dd if=$partition_path" | pv > "$outfile"
            else
                adb -s "$serial" shell "dd if=$partition_path" > "$outfile"
            fi
        done

        echo "Full dump complete."
        exit 0
        ;;

    *)
        echo "Unknown option: $OPTION"
        echo "Try '$0 -h' for help."
        exit 1
        ;;
esac
