#!/bin/bash

# Boot device as mass storage gadget

# Source hardware configuration reader
source "/usr/share/rpi-power-hat/scripts/hardware-reader.sh"

if [ -z "$1" ]; then
	echo "Usage: rpi-power-hat mass-storage <PORT> [MODEL]" >&2
	echo "Ports: B1 B2 T1 T2 | Models: 5 [default], 4, 0 (Pi Zero 2 W)" >&2
	exit 1
fi

if [ "$1" = "-h" ] || [ "$1" = "--help" ] || [ "$1" = "-help" ]; then
	echo "Mass-storage gadget helper"
	echo
	echo "Usage: rpi-power-hat mass-storage <PORT> [MODEL]"
	echo
	echo "Arguments:"
	echo "  PORT   - B1, B2, T1, T2 (silkscreened on the HAT)"
	echo "  MODEL  - 5 [default], 4, 0 (Pi Zero 2 W)"
	echo
	echo "Notes:"
	echo "  - Connect a data-capable USB-C cable to the selected port."
	echo "  - Pi 5: press/hold boot button or use GPIO wire per hardware notes."
	echo "  - Zero 2 W: Attach GPIO wire to CMD on SD card interposer."
	echo "  - Wait for 'New USB device detected!' before flashing."
	exit 0
fi

PORT_SELECTION="$1"  # 1st arg
MODEL="$2"           # 2nd arg

# Determine config file (optional; skip silently if not present)
if [ -f "/etc/rpi-power-hat/ports.conf" ]; then
	CONFIG_FILE="/etc/rpi-power-hat/ports.conf"
else
	CONFIG_FILE=""
fi

# If model is provided, persist it to ports.conf under the selected port
update_model_in_config() {
	local port="$1"
	local model="$2"
	local config_file="$3"

	[ -n "$config_file" ] && [ -f "$config_file" ] || return 0

	# Update existing model key or insert if missing within the [PORT] section
	if sed -n "/^\[$port\]/,/^\[/p" "$config_file" | grep -q "^model="; then
		sed -i "/^\[$port\]/,/^\[/ s/^model=.*/model=$model/" "$config_file"
	else
		sed -i "/^\[$port\]$/a model=$model" "$config_file"
	fi
	echo "Model updated in config: $config_file"
}

# Get GPIO line names for the selected port
if ! validate_port "$PORT_SELECTION"; then
	echo "Error: Invalid port selection '$PORT_SELECTION'" >&2
	exit 1
fi

PWR_GPIO_LINE="$(get_gpio_line_name "$PORT_SELECTION")"
BOOT_GPIO_LINE="$(get_boot_gpio_line_name "$PORT_SELECTION")"

if [[ -z "$PWR_GPIO_LINE" || -z "$BOOT_GPIO_LINE" ]]; then
	echo "Error: Unable to get GPIO line names for port $PORT_SELECTION" >&2
	exit 1
fi

# Resolve USB pathname for this port (used to target rpiboot)
PORT_MAP_SCRIPT="/usr/share/rpi-power-hat/utils/port-map.sh"
USB_PATH=""
if [[ -x "$PORT_MAP_SCRIPT" ]]; then
	USB_PATH="$("$PORT_MAP_SCRIPT" auto "$PORT_SELECTION" 2>/dev/null)"
fi

# Build rpiboot arguments (shared across all model branches)
RPIBOOT_ARGS="-v -d usr/share/rpiboot/mass-storage-gadget64"
[[ -n "$USB_PATH" ]] && RPIBOOT_ARGS+=" -p $USB_PATH"

# Handles delay between plugging in bootmode device and host device registering it
detect_device() {
	if [[ -n "$USB_PATH" ]]; then
		echo "Waiting for device on USB path $USB_PATH..."
		while true; do
			if ls /sys/bus/usb/devices/"$USB_PATH" >/dev/null 2>&1; then
				echo "Device detected on USB path $USB_PATH!"
				break
			fi
			sleep 1
		done
	else
		# Fallback: original counting method
		echo "Waiting for device to be detected..."
		INITIAL_USB_COUNT="$(lsusb 2>/dev/null | wc -l || true)"

		while true; do
			CURRENT_USB_COUNT="$(lsusb 2>/dev/null | wc -l || true)"
			if (( CURRENT_USB_COUNT > INITIAL_USB_COUNT )); then
				echo "New USB device detected!"
				break
			elif (( CURRENT_USB_COUNT < INITIAL_USB_COUNT )); then
				INITIAL_USB_COUNT="$CURRENT_USB_COUNT"
			fi
			sleep 1
		done
	fi
}

# Model selection defaults to Pi 5
if [ "$MODEL" = "4" ]; then
	echo "Booting Pi 4 on port ${PORT_SELECTION} as mass-storage-device"
	update_model_in_config "$PORT_SELECTION" "$MODEL" "$CONFIG_FILE"
	gpioset -t0 "${PWR_GPIO_LINE}=0"
	echo "Please ensure wire is fitted between the HAT ${BOOT_GPIO_LINE} and the DUT GPIO8"
	gpioset -t0 "${BOOT_GPIO_LINE}=0"
	sleep 3
	gpioset -t0 "${PWR_GPIO_LINE}=1"
	detect_device
	gpioget "${BOOT_GPIO_LINE}"
	cd /
	sudo rpiboot $RPIBOOT_ARGS

elif [ "$MODEL" = "0" ]; then
	echo "Booting Pi Zero 2 W on port $PORT_SELECTION as mass-storage-device"
	update_model_in_config "$PORT_SELECTION" "$MODEL" "$CONFIG_FILE"
	gpioset -t0 "${PWR_GPIO_LINE}=0"
	echo "If you are not using an SD sniffer, please remove SD card now and reinsert after usbboot has started"
	echo "Disrupting SD card (connect ${BOOT_GPIO_LINE} to CMD)..."
	gpioset -t0 "${BOOT_GPIO_LINE}=1"
	sleep 3
	gpioset -t0 "${PWR_GPIO_LINE}=1"
	detect_device
	gpioget "${BOOT_GPIO_LINE}"
	echo "Re-insert SD card now"
	sleep 2
	cd /
	sudo rpiboot $RPIBOOT_ARGS &
	TOOL_PID=$!
	wait $TOOL_PID
	# sleep 5
else
	# If a model was explicitly provided (e.g. 5), persist it; otherwise do not modify ports.conf
	if [ -n "$MODEL" ]; then
		update_model_in_config "$PORT_SELECTION" "$MODEL" "$CONFIG_FILE"
	fi
	echo "Booting Pi 5 on port ${PORT_SELECTION} as mass-storage-device"
	echo "Please ensure wire is fitted between the HAT ${BOOT_GPIO_LINE} and the target device TP42/bootheader"
	gpioset -t0 "${PWR_GPIO_LINE}=0"
	echo "If no wire fitted, please press boot button"
	gpioset -t0 "${BOOT_GPIO_LINE}=0"
	sleep 3
	gpioset -t0 "${PWR_GPIO_LINE}=1"
	detect_device
	# 'disconnects' bootpin by setting it to gpio input
	gpioget "${BOOT_GPIO_LINE}"
	cd /
	sudo rpiboot $RPIBOOT_ARGS
fi

detect_device

echo "Success, device now recognised as mass-storage-gadget"
