#!/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="/tmp/rpi-power-hat/firstrun"
WORKING_CLOUDINIT_DIR="/tmp/rpi-power-hat/cloudinit"
HISTORY_DIR="$VAR_DIR/history"
KEYS_DIR="$VAR_DIR/keys"

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

# Download helpers
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
	local tmp_path="${dest_path}.part"
	if command -v curl >/dev/null 2>&1; then
		curl -fL --progress-bar -o "$tmp_path" "$url" || return 1
	elif command -v wget >/dev/null 2>&1; then
		wget --progress=bar:force -O "$tmp_path" "$url" || return 1
	else
		echo "Error: Neither curl nor wget is available for downloading." >&2
		return 1
	fi
	mv "$tmp_path" "$dest_path"
	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 "Warning: Port $port device not found, using fallback selection" >&2

    # Fallback: try common device names, but ignore 0B devices
    for dev in /dev/sd{a,b,c,d}; do
        if [[ -e "$dev" ]]; then
            size=$(lsblk -b -dn -o SIZE "$dev" 2>/dev/null)
            if [[ -n "$size" && "$size" -gt 0 ]]; then
                echo "$dev"
                return 0
            fi
        fi
    done

    echo "/dev/sda"  # Final fallback
}

# Function to generate unique hostname and update ports.conf
update_port_hostname() {
    local port="$1"
    local config_file="$2"
    
    # Generate unique hostname: rpi-{port}-{8-char-random}
    local new_hostname="rpi-${port}-$(head /dev/urandom | tr -dc a-z0-9 | head -c 8)"
    
    # Add or update hostname field in the specific port section
    if grep -q "^\[${port}\]" "$config_file"; then
        # Check if hostname field exists in this section
        if sed -n "/^\[${port}\]/,/^\[/p" "$config_file" | grep -q "^hostname="; then
            # Update existing hostname field
            sed -i "/^\[${port}\]/,/^\[/ s/^hostname=.*/hostname=${new_hostname}/" "$config_file"
        else
            # Add hostname field after the section header
            sed -i "/^\[${port}\]$/a hostname=${new_hostname}" "$config_file"
        fi
        
        # Update or add ssh_host field
        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
    
    echo "$new_hostname"
}

# Function to generate unique SSH key pair and update ports.conf
update_port_ssh_keys() {
    local port="$1"
    local config_file="$2"
    
    # 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"
    
    # Write SSH client config snippet for convenience (non-fatal)
    write_ssh_config_snippet_for_port "$port" "$config_file" "$private_key_path"
    
    return 0
}

# Setup SSH configuration structure for the user
setup_user_ssh_config() {
    # Determine target user and home
    local target_user="${SUDO_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
    local target_user="${SUDO_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=$(read_config "$port" "hostname" "$config_file")
    
    # 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"
    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
    else
        # Minimal valid netplan to leave Wi‑Fi unconfigured
        cat > "$output_file" << EOF
network:
  version: 2
  renderer: NetworkManager
EOF
    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=$(read_config "$port" "hostname" "$config_file")
    local ssh_key="$LAST_GENERATED_PUBLIC_KEY"

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

    local output_file="$working_dir/user-data"

    cat > "$output_file" << EOF
#cloud-config
hostname: ${hostname}
manage_etc_hosts: true
ssh_pwauth: false
disable_root: true
users:
  - name: ${username}
    shell: /bin/bash
    lock_passwd: true
    sudo: ALL=(ALL) NOPASSWD:ALL
    ssh_authorized_keys:
      - ${ssh_key}
timezone: GB
keyboard:
  layout: gb
runcmd:
  - [ systemctl, enable, --now, ssh ]
EOF

    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

    # Set default port
    port=${port:-"DEFAULT"}
    
    echo "Using configuration for port: $port"
    
    # Generate unique hostname and update config
    local new_hostname=$(update_port_hostname "$port" "$CONFIG_FILE")
    echo "Generated unique hostname: $new_hostname"
    
    # Generate unique SSH key pair and update config
    update_port_ssh_keys "$port" "$CONFIG_FILE"
    
    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 "Downloading latest image..."
            image_path="$(download_image_to_dir "https://downloads.raspberrypi.com/raspios_arm64/images/raspios_arm64-2025-05-13/2025-05-13-raspios-bookworm-arm64.img.xz" "$IMAGES_DIR" || true)"
            [ -n "$image_path" ] && 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" ] && 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" ] && 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"
        rpi-imager --cli "$image_path" "$device_path" --cloudinit-networkconfig "$network_config_file" --cloudinit-userdata "$user_data_file"
    else
        echo "Imaging using firstrun script..."
        rpi-imager --cli "$image_path" "$device_path" --first-run-script "$firstrun_script"
    fi
    
    # Save imaging history
    save_imaging_history "$port" "$(read_config "$port" "hostname" "$CONFIG_FILE")" "$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 "$@"