#!/bin/bash

# SSH to device and set hostname; updates ports.conf if successful

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

# Determine config file path, skip if not found
if [ -f "/etc/rpi-power-hat/ports.conf" ]; then
    CONFIG_FILE="/etc/rpi-power-hat/ports.conf"
else
    CONFIG_FILE=""
fi

print_usage() {
	echo "Set the hostname on a device via SSH, then update ports.conf on success"
	echo "Usage: rpi-power-hat set-hostname <port> <new-hostname>"
	echo
	echo "Examples:"
	echo "  rpi-power-hat set-hostname B1 rpi-B1-test"
}

read_config() {
	local section="$1"
	local key="$2"
	awk -F= -v section="[$section]" -v key="$key" '
		$0 == section { in_section = 1; next }
		/^\[/ && in_section { in_section = 0 }
		in_section && $1 == key { print $2; exit }
	' "$CONFIG_FILE"
}

is_valid_hostname() {
	local h="$1"
	# RFC 1123 simplified: labels of letters, digits, hyphens; 1-63 chars per label; total <=253
	# Lowercase for consistency
	local lc
	lc="$(echo "$h" | tr 'A-Z' 'a-z')"
	# quick rejects
	[[ -z "$lc" ]] && return 1
	[[ ${#lc} -gt 253 ]] && return 1
	# must not start or end with hyphen or dot
	[[ "$lc" =~ ^[-.] || "$lc" =~ [-.]$ ]] && return 1
	# allowed chars only
	[[ "$lc" =~ ^[a-z0-9.-]+$ ]] || return 1
	# each label 1-63, no leading/trailing hyphen
	IFS='.' read -r -a parts <<< "$lc"
	for label in "${parts[@]}"; do
		[[ -z "$label" ]] && return 1
		[[ ${#label} -gt 63 ]] && return 1
		[[ "$label" =~ ^- || "$label" =~ -$ ]] && return 1
		[[ "$label" =~ ^[a-z0-9-]+$ ]] || return 1
	done
	return 0
}

update_key_in_config() {
    local port="$1"
    local key="$2"
    local value="$3"
    local file="$4"

    if sed -n "/^\[$port\]/,/^\[/p" "$file" | grep -q "^$key="; then
        sed -i "/^\[$port\]/,/^\[/ s/^$key=.*/$key=$value/" "$file" 2>/dev/null || \
        sudo -n sed -i "/^\[$port\]/,/^\[/ s/^$key=.*/$key=$value/" "$file" 2>/dev/null || return 1
    else
        sed -i "/^\[$port\]$/a $key=$value" "$file" 2>/dev/null || \
        sudo -n sed -i "/^\[$port\]$/a $key=$value" "$file" 2>/dev/null || return 1
    fi
    return 0
}

run_remote_set_hostname() {
	local port="$1"
	local newhost="$2"

	if command -v rpi-power-hat >/dev/null 2>&1; then
		rpi-power-hat ssh "$port" sudo hostnamectl set-hostname "$newhost"
		return $?
	fi

	local local_ssh_script
	local_ssh_script="$(dirname "$0")/ssh"
	if [ -x "$local_ssh_script" ]; then
		"$local_ssh_script" "$port" sudo hostnamectl set-hostname "$newhost"
		return $?
	fi

	echo "Error: could not find 'rpi-power-hat ssh' or local ssh script" >&2
	return 127
}

amend_ssh_snippet() {
    local port="$1"
    local newhost="$2"

    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"
    local snippet_file="$snippet_dir/${port}.conf"

    mkdir -p "$snippet_dir" 2>/dev/null || true

    if [ -f "$snippet_file" ]; then
        local to_add=""
        if ! grep -qE "^Host .*\\b${newhost}\\b" "$snippet_file" 2>/dev/null; then
            to_add="$newhost"
        fi
        if ! grep -qE "^Host .*\\b${newhost}\\.local\\b" "$snippet_file" 2>/dev/null; then
            to_add="${to_add} ${newhost}.local"
        fi
        to_add="${to_add# }"
        if [ -n "$to_add" ]; then
            sed -i "0,/^Host / s/^Host \(.*\)$/Host \1 ${to_add}/" "$snippet_file"
        fi
        chmod 600 "$snippet_file" 2>/dev/null || true
        chown "$target_user":"$target_user" "$snippet_file" 2>/dev/null || true
        return 0
    fi

    # Create a new snippet if none exists
    local username
    username="$(read_config "$port" "username")"
    [ -n "$username" ] || username="$(read_config "DEFAULT" "username")"
    [ -n "$username" ] || username="pi"

    local key_dir="/var/lib/rpi-power-hat/keys"
    local identity_file=""
    if [ -d "$key_dir" ]; then
        identity_file="$(ls -1t "$key_dir"/rpi-${port}-* 2>/dev/null | grep -v '\.pub$' | head -n1)"
    fi

    local host_pattern="rpi-${port}-* ${newhost} ${newhost}.local"

    {
        echo "Host ${host_pattern}"
        echo "\tUser ${username}"
        if [ -n "$identity_file" ]; then
            echo "\tIdentityFile ${identity_file}"
            echo "\tIdentitiesOnly yes"
        fi
        echo "\tStrictHostKeyChecking no"
        echo "\tUserKnownHostsFile /dev/null"
        echo "\tLogLevel ERROR"
    } > "$snippet_file"

    chmod 600 "$snippet_file" 2>/dev/null || true
    chown "$target_user":"$target_user" "$snippet_file" 2>/dev/null || true
}

main() {
	local port="$1"
	local new_hostname="$2"

	if [ -z "$port" ] || [ "$port" = "-h" ] || [ "$port" = "--help" ] || [ "$port" = "help" ]; then
		if [ -z "$port" ] || [ -z "$new_hostname" ]; then
			print_usage
			exit 0
		fi
	fi

	if [ -z "$new_hostname" ]; then
		echo "Error: new hostname is required" >&2
		print_usage
		exit 1
	fi

	if ! validate_port "$port"; then
		echo "Error: Invalid port '$port'. Valid ports are: $(list_ports)" >&2
		exit 1
	fi

	if ! is_valid_hostname "$new_hostname"; then
		echo "Error: invalid hostname '$new_hostname'" >&2
		exit 1
	fi

	if [ -z "$CONFIG_FILE" ] || [ ! -f "$CONFIG_FILE" ]; then
		echo "Error: ports.conf not found" >&2
		exit 1
	fi

	echo "Setting hostname on $port to '$new_hostname'..."
	if run_remote_set_hostname "$port" "$new_hostname"; then
		echo "Remote hostname set successfully. Updating ports.conf..."
		if ! update_key_in_config "$port" "hostname" "$new_hostname" "$CONFIG_FILE"; then
			echo "Error: failed to update hostname in $CONFIG_FILE (try running with sudo)" >&2
			exit 1
		fi
		# Keep ssh_host in sync for future SSH
		if ! update_key_in_config "$port" "ssh_host" "$new_hostname" "$CONFIG_FILE"; then
			echo "Warning: failed to update ssh_host in $CONFIG_FILE" >&2
		fi
		# Amend local SSH client snippet so the new hostname is recognized
		amend_ssh_snippet "$port" "$new_hostname"
		echo "Done."
	else
		echo "Error: SSH command failed; not updating ports.conf" >&2
		exit 1
	fi
}

main "$@"


