#!/bin/bash

# Flash OS images to devices

# Paths (package installation only)
CONFIG_FILE="/etc/rpi-power-hat/ports.conf"
NETWORK_CONFIG_FILE="/etc/rpi-power-hat/network.conf"
SHARE_DIR="/usr/share/rpi-power-hat"
SCRIPTS_DIR="$SHARE_DIR/scripts"
TEMPLATES_DIR="$SHARE_DIR/templates/firstrun"
UTILS_DIR="$SHARE_DIR/utils"
VAR_DIR="/var/lib/rpi-power-hat"
IMAGES_DIR="$VAR_DIR/images"
WORKING_FIRSTRUN_DIR=""  # Set per-run in main()
WORKING_CLOUDINIT_DIR=""  # Set per-run in main()
HISTORY_DIR="$VAR_DIR/history"
KEYS_DIR="$VAR_DIR/keys"
PORTS_CONF_LOCK="/var/lock/rpi-power-hat-ports-conf.lock"

locked_config_update() {
    local lockfd
    exec {lockfd}>"$PORTS_CONF_LOCK"
    flock -w 30 "$lockfd" || { echo "Error: Could not acquire config lock" >&2; return 1; }
    "$@"
    local rc=$?
    exec {lockfd}>&-
    return $rc
}

if [ ! -f "$CONFIG_FILE" ]; then
    echo "Error: Could not find ports.conf configuration file at $CONFIG_FILE" >&2
    exit 1
fi

# Download helpers
resolve_latest_image_url() {
	local release="${1:-raspios_arm64}"
	local base_url="https://downloads.raspberrypi.com/"
	local most_recent_dir
	most_recent_dir="$(curl --silent "${base_url}${release}/images/" \
		| grep -oP 'a href="\K'"${release}"'-\d{4}-\d{2}-\d{2}(?=/")' \
		| sort | tail -1)"
	if [ -z "$most_recent_dir" ]; then
		echo "Error: Could not find latest image directory for $release" >&2
		return 1
	fi
	local img_xz_file
	img_xz_file="$(curl --silent "${base_url}${release}/images/${most_recent_dir}/" \
		| grep -oP 'a href="\K\d{4}-\d{2}-\d{2}.*?\.img\.xz(?=")' \
		| tail -1)"
	if [ -z "$img_xz_file" ]; then
		echo "Error: Could not find .img.xz file in $most_recent_dir" >&2
		return 1
	fi
	echo "${base_url}${release}/images/${most_recent_dir}/${img_xz_file}"
}

download_image_to_dir() {
	local url="$1"
	local dest_dir="$2"
	mkdir -p "$dest_dir"
	local filename
	filename="$(basename "$url")"
	local dest_path="$dest_dir/$filename"
	# Reuse if already downloaded
	if [ -f "$dest_path" ]; then
		echo "$dest_path"
		return 0
	fi
	# Acquire per-file lock to prevent concurrent downloads of the same image
	local dl_lockfd
	exec {dl_lockfd}>"${dest_path}.lock"
	flock -w 600 "$dl_lockfd" || { echo "Error: Could not acquire download lock for $filename" >&2; return 1; }
	# Double-check: another process may have finished while we waited
	if [ -f "$dest_path" ]; then
		exec {dl_lockfd}>&-
		echo "$dest_path"
		return 0
	fi
	local tmp_path="${dest_path}.part"
	if command -v curl >/dev/null 2>&1; then
		curl -fL --progress-bar -o "$tmp_path" "$url" || { exec {dl_lockfd}>&-; return 1; }
	elif command -v wget >/dev/null 2>&1; then
		wget --progress=bar:force -O "$tmp_path" "$url" || { exec {dl_lockfd}>&-; return 1; }
	else
		echo "Error: Neither curl nor wget is available for downloading." >&2
		exec {dl_lockfd}>&-
		return 1
	fi
	mv "$tmp_path" "$dest_path"
	exec {dl_lockfd}>&-
	echo "$dest_path"
}

# Function to read configuration value for a given port
read_config() {
    local port="$1"
    local key="$2"
    local config_file="$3"
    
    # Read the value from the specified port section
    local value=$(awk -v port="[$port]" -v key="$key" '
        BEGIN { in_section = 0 }
        $0 ~ "^\\[" { in_section = ($0 == port) }
        in_section && $0 ~ "^" key "=" { 
            gsub("^" key "=", ""); 
            print; 
            exit 
        }
    ' "$config_file")
    
    echo "$value"
}

# Wrapper to read Wi-Fi configuration (SSID/password) from network.conf,
# with backward-compatible fallback to ports.conf if network.conf is absent.
read_network_config() {
    local port="$1"
    local key="$2"
    local value=""

    # 1) network.conf [PORT]
    if [ -f "$NETWORK_CONFIG_FILE" ]; then
        value=$(read_config "$port" "$key" "$NETWORK_CONFIG_FILE")
    fi
    # 2) ports.conf [PORT]
    if [ -z "$value" ]; then
        value=$(read_config "$port" "$key" "$CONFIG_FILE")
    fi
    # 3) network.conf [DEFAULT]
    if [ -z "$value" ] && [ -f "$NETWORK_CONFIG_FILE" ]; then
        value=$(read_config "DEFAULT" "$key" "$NETWORK_CONFIG_FILE")
    fi
    # 4) ports.conf [DEFAULT]
    if [ -z "$value" ]; then
        value=$(read_config "DEFAULT" "$key" "$CONFIG_FILE")
    fi
    echo "$value"
}

# Function to determine device path based on port using port-map
get_device_path() {
    local port="$1"
    
    # Use packaged port-map.sh script
    local port_map_script="$UTILS_DIR/port-map.sh"
    
    # Use port-map.sh to get actual block device (auto-detects Pi model)
    if [[ -x "$port_map_script" ]]; then
        local device=$("$port_map_script" -d auto "$port" 2>/dev/null)
        if [[ -n "$device" && -e "$device" ]]; then
            echo "$device"
            return 0
        fi
    fi
    
    echo "Error: Could not determine device for port $port via USB topology." >&2
    echo "Ensure the device is connected and in mass-storage mode." >&2
    return 1
}

# Function to generate unique hostname (does not update ports.conf)
generate_hostname() {
    local port="$1"
    echo "rpi-${port}-$(head /dev/urandom | tr -dc a-z0-9 | head -c 8)"
}

# Write hostname and ssh_host to ports.conf for a given port
save_hostname_to_config() {
    local port="$1"
    local config_file="$2"
    local new_hostname="$3"

    if grep -q "^\[${port}\]" "$config_file"; then
        if sed -n "/^\[${port}\]/,/^\[/p" "$config_file" | grep -q "^hostname="; then
            sed -i "/^\[${port}\]/,/^\[/ s/^hostname=.*/hostname=${new_hostname}/" "$config_file"
        else
            sed -i "/^\[${port}\]$/a hostname=${new_hostname}" "$config_file"
        fi

        if sed -n "/^\[${port}\]/,/^\[/p" "$config_file" | grep -q "^ssh_host="; then
            sed -i "/^\[${port}\]/,/^\[/ s/^ssh_host=.*/ssh_host=${new_hostname}/" "$config_file"
        else
            sed -i "/^\[${port}\]$/a ssh_host=${new_hostname}" "$config_file"
        fi
    fi
}

# Function to generate unique SSH key pair (does not update ports.conf)
generate_ssh_keys() {
    local port="$1"

    # Key storage path (package installation)
    local key_dir="$KEYS_DIR"

    # Create keys directory if it doesn't exist
    mkdir -p "$key_dir"

    # Generate unique key name with timestamp and random suffix
    local key_name="rpi-${port}-$(date +%Y%m%d)-$(head /dev/urandom | tr -dc a-z0-9 | head -c 8)"
    local private_key_path="${key_dir}/${key_name}"
    local public_key_path="${key_dir}/${key_name}.pub"

    # Generate SSH key pair (ed25519 for better security and smaller size)
    ssh-keygen -t ed25519 -f "$private_key_path" -N "" -C "rpi-power-hat ${port} $(date +%Y-%m-%d)" >/dev/null 2>&1

    if [[ $? -ne 0 ]]; then
        echo "Error: Failed to generate SSH key pair for port $port" >&2
        return 1
    fi

    # Read the public key
    local public_key=$(cat "$public_key_path" 2>/dev/null)
    if [[ -z "$public_key" ]]; then
        echo "Error: Failed to read generated public key for port $port" >&2
        return 1
    fi

    # Make generated key available later in this run
    LAST_GENERATED_PUBLIC_KEY="$public_key"
    LAST_PRIVATE_KEY_PATH="$private_key_path"

    # Set appropriate permissions
    chmod 600 "$private_key_path"
    chmod 644 "$public_key_path"

    echo "Generated SSH key pair: $key_name"
    echo "  Private key: $private_key_path"
    echo "  Public key: $public_key_path"

    return 0
}

# Write SSH key and config to ports.conf and SSH config snippet
save_ssh_keys_to_config() {
    local port="$1"
    local config_file="$2"
    local public_key="$LAST_GENERATED_PUBLIC_KEY"
    local private_key_path="$LAST_PRIVATE_KEY_PATH"

    if grep -q "^\[${port}\]" "$config_file"; then
        # Update ssh_key
        if sed -n "/^\[${port}\]/,/^\[/p" "$config_file" | grep -q "^ssh_key="; then
            sed -i "/^\[${port}\]/,/^\[/ s|^ssh_key=.*|ssh_key=${public_key}|" "$config_file"
        else
            sed -i "/^\[${port}\]$/a ssh_key=${public_key}" "$config_file"
        fi

        # Update ssh_private_key_path
        if sed -n "/^\[${port}\]/,/^\[/p" "$config_file" | grep -q "^ssh_private_key_path="; then
            sed -i "/^\[${port}\]/,/^\[/ s|^ssh_private_key_path=.*|ssh_private_key_path=${private_key_path}|" "$config_file"
        else
            sed -i "/^\[${port}\]$/a ssh_private_key_path=${private_key_path}" "$config_file"
        fi
    fi

    # Write SSH client config snippet for convenience (non-fatal)
    write_ssh_config_snippet_for_port "$port" "$config_file" "$private_key_path"
}

# Setup SSH configuration structure for the user
setup_user_ssh_config() {
    # Determine target user and home
    # Note: We use $USER (not SUDO_USER) because the script is meant to be
    # run as the target user directly (e.g. sudo -u gitlab-runner).
    # Running as root via plain sudo is not supported.
    local target_user="$USER"
    local home_dir
    home_dir="$(getent passwd "$target_user" | cut -d: -f6 2>/dev/null)"
    [ -n "$home_dir" ] || home_dir="$HOME"
    [ -n "$home_dir" ] || return 0

    local ssh_dir="$home_dir/.ssh"
    local config_file="$ssh_dir/config"
    local pi_config_file="$ssh_dir/pi_config"
    local config_d_dir="$ssh_dir/config.d"
    local snippet_dir="$ssh_dir/rpi-power-hat.d"

    # Skip if already configured
    if [ -f "$pi_config_file" ] && grep -qE '^[Ii]nclude[[:space:]]+~/.ssh/rpi-power-hat\.d/\*' "$pi_config_file" 2>/dev/null; then
        echo "SSH configuration already setup for user $target_user"
        return 0
    fi

    echo "Setting up SSH configuration for user $target_user"

    # Create SSH directories with proper permissions
    mkdir -p "$ssh_dir" "$config_d_dir" "$snippet_dir" || {
        echo "Error: Failed to create SSH directories" >&2
        return 1
    }

    # Create pi_config if not present
    if [ ! -f "$pi_config_file" ]; then
        touch "$pi_config_file" || {
            echo "Error: Failed to create pi_config file" >&2
            return 1
        }
    fi

    # Create main config if not present
    if [ ! -f "$config_file" ]; then
        touch "$config_file" || {
            echo "Error: Failed to create SSH config file" >&2
            return 1
        }
    fi

    # Add Include directives to main config if missing
    if ! grep -qE '^[Ii]nclude[[:space:]]+~/.ssh/config.d/\*' "$config_file" 2>/dev/null; then
        local tmp_file=$(mktemp)
        {
            echo "Include ~/.ssh/config.d/*"
            echo "Include ~/.ssh/pi_config"
            cat "$config_file"
        } > "$tmp_file"
        cp "$tmp_file" "$config_file" || {
            echo "Error: Failed to update SSH config file" >&2
            rm "$tmp_file"
            return 1
        }
        rm "$tmp_file"
    fi

    # Add rpi-power-hat.d include to pi_config if missing
    if ! grep -qE '^[Ii]nclude[[:space:]]+~/.ssh/rpi-power-hat\.d/\*' "$pi_config_file" 2>/dev/null; then
        local tmp_file2=$(mktemp)
        {
            echo "Include ~/.ssh/rpi-power-hat.d/*"
            cat "$pi_config_file"
        } > "$tmp_file2"
        cp "$tmp_file2" "$pi_config_file" || {
            echo "Error: Failed to update pi_config file" >&2
            rm "$tmp_file2"
            return 1
        }
        rm "$tmp_file2"
    fi

    # Set proper permissions (SSH requires specific permissions)
    chmod 700 "$ssh_dir" "$config_d_dir" "$snippet_dir" || {
        echo "Warning: Failed to set directory permissions" >&2
    }
    chmod 600 "$config_file" "$pi_config_file" || {
        echo "Warning: Failed to set file permissions" >&2
    }

    echo "SSH configuration setup completed for user $target_user"
}

# Create SSH config files per port, place it in ~/.ssh/rpi-power-hat.d/
# During setup an include line has been added to ~/.ssh/config to include all of these files.
write_ssh_config_snippet_for_port() {
    local port="$1"
    local config_file="$2"
    local private_key_path="$3"

    # Ensure SSH config structure is setup first
    setup_user_ssh_config

    # Determine target user and home (see setup_user_ssh_config for rationale)
    local target_user="$USER"
    local home_dir
    home_dir="$(getent passwd "$target_user" | cut -d: -f6 2>/dev/null)"
    [ -n "$home_dir" ] || home_dir="$HOME"
    [ -n "$home_dir" ] || return 0

    local ssh_dir="$home_dir/.ssh"
    local snippet_dir="$ssh_dir/rpi-power-hat.d"

    # Read current username and hostname from config
    local username=$(read_config "$port" "username" "$config_file")
    local hostname=$(read_config "$port" "hostname" "$config_file")
    [ -n "$username" ] || username="pi"

    # Path for this port's snippet; keep latest mapping per port
    local snippet_file="$snippet_dir/${port}.conf"

    # Compose host patterns: specific and wildcard by port
    local host_pattern="rpi-${port}-*"
    if [ -n "$hostname" ]; then
        host_pattern="$host_pattern $hostname ${hostname}.local"
    fi

    # Create snippet file as target user with proper content
    local tmp_file=$(mktemp)
    cat > "$tmp_file" << EOF
Host $host_pattern
    User $username
    IdentityFile $private_key_path
    IdentitiesOnly yes
    StrictHostKeyChecking no
    UserKnownHostsFile /dev/null
    LogLevel ERROR
EOF

    # Check if snippet directory exists before copying
    if [ ! -d "$snippet_dir" ]; then
        echo "Warning: SSH snippet directory $snippet_dir does not exist, creating it..."
        mkdir -p "$snippet_dir" || {
            echo "Error: Failed to create SSH snippet directory $snippet_dir" >&2
            rm "$tmp_file"
            return 1
        }
        chmod 700 "$snippet_dir"
    fi
    
    # Copy with proper permissions
    cp "$tmp_file" "$snippet_file" || {
        echo "Error: Failed to create SSH config snippet $snippet_file" >&2
        rm "$tmp_file"
        return 1
    }
    chmod 600 "$snippet_file"
    rm "$tmp_file"
}

# Function to save imaging history
save_imaging_history() {
    local port="$1"
    local hostname="$2"
    local device_path="$3"
    local image_path="$4"
    local model="$5"
    local config_file="$6"
    
    # Create timestamp for filename and record
    local timestamp=$(date '+%Y-%m-%d_%H-%M-%S')
    local date_readable=$(date '+%Y-%m-%d %H:%M:%S')
    
    # History directory (package installation)
    local history_dir="$HISTORY_DIR"
    mkdir -p "$history_dir"
    
    # Create history filename
    local history_file="${history_dir}/${timestamp}_${port}_${hostname}.log"
    
    # Read additional config values
    local username=$(read_config "$port" "username" "$config_file")
    local wlan_ssid=$(read_network_config "$port" "wlan_ssid")
    local ssh_host=$(read_config "$port" "ssh_host" "$config_file")
    
    # Write imaging history
    cat > "$history_file" << EOF
# rpi-power-hat Imaging History
# Generated: $date_readable

[IMAGING_SESSION]
timestamp=$timestamp
date=$date_readable
port=$port
hostname=$hostname
device=$device_path
image=$image_path
model=$model

[DEVICE_CONFIG]
username=$username
ssh_host=$ssh_host
wlan_ssid=$wlan_ssid

[PATHS]
config_file=$config_file
history_file=$history_file

# SSH Connection:
# ssh $username@$hostname.local

EOF
    
    echo "Imaging history saved: $history_file"
}

# Function to generate firstrun.sh based on configuration using template
generate_firstrun() {
    local port="$1"
    local config_file="$2"
    
    # Determine template and output paths (package installation)
    local template_file="$TEMPLATES_DIR/firstrun_template.sh"
    local output_file="$WORKING_FIRSTRUN_DIR/firstrun_${port}.sh"
    local working_dir="$WORKING_FIRSTRUN_DIR"
    
    # Check if template exists
    if [[ ! -f "$template_file" ]]; then
        echo "Error: Template file '$template_file' not found!" >&2
        exit 1
    fi
    
    # Read configuration values
    local username=$(read_config "$port" "username" "$config_file")
    local ssh_key="$LAST_GENERATED_PUBLIC_KEY"
    local wlan_ssid=$(read_network_config "$port" "wlan_ssid")
    local wlan_password=$(read_network_config "$port" "wlan_password")
    local hostname="$LAST_GENERATED_HOSTNAME"
    
    # Use defaults if values are empty
    username=${username:-"pi"}
    ssh_key=${ssh_key:-""}
    wlan_ssid=${wlan_ssid:-""}
    wlan_password=${wlan_password:-""}
    hostname=${hostname:-"cake"}
    
    # Ensure working directory exists
    mkdir -p "$working_dir"
    
    # Read template and substitute placeholders
    local template_content=$(cat "$template_file")
    
    # Replace placeholders with actual values
    template_content="${template_content//\{\{HOSTNAME\}\}/$hostname}"
    template_content="${template_content//\{\{USERNAME\}\}/$username}"
    template_content="${template_content//\{\{SSH_KEY\}\}/$ssh_key}"
    template_content="${template_content//\{\{WLAN_SSID\}\}/$wlan_ssid}"
    template_content="${template_content//\{\{WLAN_PASSWORD\}\}/$wlan_password}"
    
    # Write the generated script
    echo "$template_content" > "$output_file" || { echo "Error: Failed to write $output_file" >&2; exit 1; }
    chmod +x "$output_file"

    echo "$output_file"
}

# Function to generate cloud-init network-config based on configuration
generate_network_config() {
    local port="$1"
    local config_file="$2"

    local working_dir="$WORKING_CLOUDINIT_DIR"
    mkdir -p "$working_dir"

    local wlan_ssid=$(read_network_config "$port" "wlan_ssid")
    local wlan_password=$(read_network_config "$port" "wlan_password")

    local output_file="$working_dir/network-config"

    if [[ -n "$wlan_ssid" ]]; then
        cat > "$output_file" << EOF
network:
  version: 2
  renderer: NetworkManager
  wifis:
    wlan0:
      dhcp4: true
      access-points:
        "${wlan_ssid}":
          password: "${wlan_password}"
          hidden: true
      optional: true
EOF
        [[ $? -ne 0 ]] && { echo "Error: Failed to write $output_file" >&2; exit 1; }
    else
        # Minimal valid netplan to leave Wi‑Fi unconfigured
        cat > "$output_file" << EOF
network:
  version: 2
  renderer: NetworkManager
EOF
        [[ $? -ne 0 ]] && { echo "Error: Failed to write $output_file" >&2; exit 1; }
    fi

    echo "$output_file"
}

# Function to generate cloud-init user-data based on configuration
generate_userdata() {
    local port="$1"
    local config_file="$2"

    local working_dir="$WORKING_CLOUDINIT_DIR"
    mkdir -p "$working_dir"

    local username=$(read_config "$port" "username" "$config_file")
    local hostname="$LAST_GENERATED_HOSTNAME"
    local ssh_key="$LAST_GENERATED_PUBLIC_KEY"

    username=${username:-"pi"}
    hostname=${hostname:-"rpi-${port}"}

    local output_file="$working_dir/user-data"

    # Hash for default password "raspberry"
    local password_hash='$6$bJN6rc42S4by69dC$O1RXFgeQdhyKXh0As8fwP8QbKf4NCwUgAee0P9SowTdOG3bNxSdypU0.a0mX2yXLuJqTHHGm1HNr.Pm1UftEW1'

    cat > "$output_file" << EOF
#cloud-config
hostname: ${hostname}
manage_etc_hosts: true
manage_resolv_conf: false
packages:
- avahi-daemon
apt:
  preserve_sources_list: true
  conf: |
    Acquire {
      Check-Date "false";
    };
timezone: GB
keyboard:
  model: pc105
  layout: gb
users:
- name: ${username}
  groups: users,adm,dialout,audio,netdev,video,plugdev,cdrom,games,input,gpio,spi,i2c,render,sudo
  shell: /bin/bash
  lock_passwd: false
  passwd: "${password_hash}"
  ssh_authorized_keys:
    - ${ssh_key}
  sudo: ALL=(ALL) NOPASSWD:ALL
enable_ssh: true
ssh_pwauth: false
EOF
    [[ $? -ne 0 ]] && { echo "Error: Failed to write $output_file" >&2; exit 1; }

    echo "$output_file"
}

update_image_path_in_config() {
    local port="$1"
    local value="$2"
    local cfg="$CONFIG_FILE"
    if sed -n "/^\[$port\]/,/^\[/p" "$cfg" | grep -q "^image_path="; then
        sed -i "/^\[$port\]/,/^\[/ s|^image_path=.*|image_path=$value|" "$cfg"
    else
        sed -i "/^\[$port\]$/a image_path=$value" "$cfg"
    fi
}

# Main script logic
main() {
    local port="$1"
    local image_arg="$2"  # optional: URL | latest | nightly

    # Create isolated temp directory for this run (avoids stale files from other users)
    WORKING_DIR=$(mktemp -d /tmp/rpi-power-hat.XXXXXXXXXX)
    WORKING_FIRSTRUN_DIR="$WORKING_DIR/firstrun"
    WORKING_CLOUDINIT_DIR="$WORKING_DIR/cloudinit"
    trap 'rm -rf "$WORKING_DIR"' EXIT

    # Set default port
    port=${port:-"DEFAULT"}
    
    echo "Using configuration for port: $port"

    # Generate unique hostname and SSH key pair (config updated after successful flash)
    local new_hostname=$(generate_hostname "$port")
    LAST_GENERATED_HOSTNAME="$new_hostname"
    echo "Generated unique hostname: $new_hostname"

    generate_ssh_keys "$port"
    
    local firstrun_script=$(generate_firstrun "$port" "$CONFIG_FILE")

    echo "Generated firstrun script: $firstrun_script"
    
    # Determine image to use
    local image_path=""
    case "$image_arg" in
        latest)
            echo "Resolving latest image URL..."
            local latest_url
            latest_url="$(resolve_latest_image_url "raspios_arm64")" || { echo "Error: Failed to resolve latest image URL" >&2; exit 1; }
            echo "Downloading latest image: $latest_url"
            image_path="$(download_image_to_dir "$latest_url" "$IMAGES_DIR" || true)"
            [ -n "$image_path" ] && locked_config_update update_image_path_in_config "$port" "$image_path"
            ;;
        nightly)
            echo "Downloading nightly image..."
            nightly_date="$(date +%Y-%m-%d)"
            nightly_filename="image_${nightly_date}-raspios-trixie-arm64.img.xz"
            nightly_url="https://pi-gen.pitowers.org/pi-gen/latest/raspios_arm64/${nightly_filename}"
            image_path="$(download_image_to_dir "$nightly_url" "$IMAGES_DIR" || true)"
            [ -n "$image_path" ] && locked_config_update update_image_path_in_config "$port" "$image_path"
            ;;
        http://*|https://*)
            echo "Downloading image from $image_arg..."
            image_path="$(download_image_to_dir "$image_arg" "$IMAGES_DIR" || true)"
            [ -n "$image_path" ] && locked_config_update update_image_path_in_config "$port" "$image_path"
            ;;
        *)
            echo "No image argument, using image from config..."
            # Fallback to config if arg isn't a URL keyword
            image_path="$(read_config "$port" "image_path" "$CONFIG_FILE")"
            ;;
    esac

    # Get device model for display
    local model=$(read_config "$port" "model" "$CONFIG_FILE")
    model=${model:-"5"}

    # If no /dev/sdX devices are present, try to boot the device into mass-storage mode
    shopt -s nullglob
    local sd_devs=(/dev/sd[a-z])
    shopt -u nullglob
    if [[ ${#sd_devs[@]} -eq 0 ]]; then
        echo "No /dev/sdX devices found. Attempting to boot device on port $port into mass-storage mode..."
        if command -v rpi-power-hat >/dev/null 2>&1; then
            rpi-power-hat mass-storage "$port" "$model"
        elif [[ -x "$SCRIPTS_DIR/mass-storage" ]]; then
            "$SCRIPTS_DIR/mass-storage" "$port" "$model"
        else
            echo "Warning: mass-storage not found; cannot auto-boot device" >&2
        fi
        sleep 5
    fi

    # Determine device path from port (after possible boot)
    local device_path=$(get_device_path "$port")
    
    # Display configuration being used
    echo "Configuration:"
    echo "  Port: $port"
    echo "  Username: $(read_config "$port" "username" "$CONFIG_FILE")"
    echo "  WLAN SSID: $(read_network_config "$port" "wlan_ssid")"
    echo "  Model: $model"
    echo "  Image: $image_path"
    echo "  Device: $device_path"
    echo
    echo "Starting image flash process..."
    echo "Image: $image_path"
    echo "Device: $device_path"
    echo "Firstrun script: $firstrun_script"
    echo

    # For trixie images, use cloud-init. Generate files from config and provide paths.
    if [[ "$image_path" == *trixie* ]]; then
        echo "Imaging using cloud-init..."
        local network_config_file=$(generate_network_config "$port" "$CONFIG_FILE")
        local user_data_file=$(generate_userdata "$port" "$CONFIG_FILE")
        echo "Generated network-config: $network_config_file"
        echo "Generated user-data: $user_data_file"
        sudo rpi-imager --cli "$image_path" "$device_path" --cloudinit-networkconfig "$network_config_file" --cloudinit-userdata "$user_data_file" \
            || { echo "Error: rpi-imager failed" >&2; exit 1; }
    else
        echo "Imaging using firstrun script..."
        sudo rpi-imager --cli "$image_path" "$device_path" --first-run-script "$firstrun_script" \
            || { echo "Error: rpi-imager failed" >&2; exit 1; }
    fi
    
    # Flash succeeded — now update ports.conf and SSH config
    _save_all_config_for_port() {
        save_hostname_to_config "$1" "$2" "$3"
        save_ssh_keys_to_config "$1" "$2"
    }
    locked_config_update _save_all_config_for_port "$port" "$CONFIG_FILE" "$new_hostname"

    # Save imaging history
    save_imaging_history "$port" "$new_hostname" "$device_path" "$image_path" "$model" "$CONFIG_FILE"
    
    sleep 2
    
    # Cut power to the device
    echo "Restarting device..."
    rpi-power-hat set "$port" 0

    sleep 2

    # Restore power to the device
    rpi-power-hat set "$port" 1
    
}

# Show usage if no arguments or help requested
if [[ "$1" == "-h" || "$1" == "--help" ]]; then
	echo "Image flashing script with port-based configuration"
	echo "Usage: rpi-power-hat flash-image [port] [latest|nightly|URL]"
	echo
	echo "Arguments:"
	echo "  port    - Configuration port (B1, B2, T1, T2) [default: DEFAULT]"
	echo "  latest  - Download latest Raspberry Pi OS arm64 image and use it"
	echo "  nightly - Download latest Raspberry Pi OS nightly (arm64) and use it"
	echo "  URL     - Download image from given URL and use it"
	echo
	echo "Examples:"
	echo "  rpi-power-hat flash-image              # Use DEFAULT config"
	echo "  rpi-power-hat flash-image B1           # Use image specified for B1 in ports.conf"
	echo "  rpi-power-hat flash-image B1 latest    # Auto-download latest image"
	echo "  rpi-power-hat flash-image B2 nightly   # Auto-download latest nightly image"
	echo "  rpi-power-hat flash-image T2 https://example.com/image.img.xz"
	echo
	echo "Configuration file: $CONFIG_FILE"
	exit 0
fi

# Call main function with all arguments
main "$@"