#!/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

# Handles delay between plugging in bootmode device and host device registering it
detect_device() {
	echo "Waiting for device to be detected..."
	INITIAL_USB_COUNT="$(lsusb 2>/dev/null | wc -l || true)"

	# Loop until a new device is detected
	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 # Exit the loop
		elif (( CURRENT_USB_COUNT < INITIAL_USB_COUNT )); then
			echo "A USB device has disconnected, updating initial USB count. Please do not disconnect/connect any of your own devices whilst the script is running" >&2
			INITIAL_USB_COUNT="$CURRENT_USB_COUNT"
		else
			echo "No new device yet. Current count: $CURRENT_USB_COUNT. Initial count: $INITIAL_USB_COUNT. Retrying in 1 second..."
			sleep 1
		fi
	done
}

# 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
	cd /
	sudo rpiboot -v -d usr/share/rpiboot/mass-storage-gadget64

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 -v -d usr/share/rpiboot/mass-storage-gadget64 &
	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 -v -d usr/share/rpiboot/mass-storage-gadget64
fi

detect_device

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