#!/bin/bash
set -euo pipefail

usage() {
	cat <<'EOF'
Usage:
  apstainit.sh <PHYNAME> <DEVNAME> <SSID> <CHANNEL> <TXPOWER> <NOISEFLOOR> <TOPO> <AP_NODE>
  apstainit.sh <ap|sta> <PHYNAME> <DEVNAME> <SSID> <CHANNEL> <TXPOWER> <NOISEFLOOR> [AP_NODE]

Examples:
  apstainit.sh phy0 wlan0 infra0 36 700 -110 2hop apu00
  apstainit.sh phy1 wlan1 infra1 161 700 -110 5x5 0
  apstainit.sh ap  phy0 wlan0 infra0 36 700 -110

Notes:
  - In topology mode, the local node becomes `ap` iff its hostname matches `AP_NODE`.
    All other nodes inside the topology become `sta`.
  - Nodes outside the topology skip initialization.
  - `AP_NODE` accepts either `apu00` style names or numeric IDs like `0`.
	 - Infrastructure mode uses WPA2-PSK/CCMP with the fixed shared secret
	   `meshman11s`.
EOF
}

if [[ $# -lt 7 ]]; then
	usage
	exit 1
fi

MODE="topology"
ROLE=""
TOPOLOGY=""
AP_NODE=""
WPA_PASSPHRASE="meshman11s"

if [[ "${1}" == "ap" || "${1}" == "sta" ]]; then
	MODE="manual"
	ROLE=$1
	PHYNAME=$2
	DEVNAME=$3
	SSID=$4
	CHANNEL_NUMBER=$5
	TXPOWER=$6
	NOISEFLOOR=$7
	AP_NODE=${8:-}
else
	if [[ $# -lt 8 ]]; then
		usage
		exit 1
	fi

	PHYNAME=$1
	DEVNAME=$2
	SSID=$3
	CHANNEL_NUMBER=$4
	TXPOWER=$5
	NOISEFLOOR=$6
	TOPOLOGY=$7
	AP_NODE=$8
fi

source /opt/scripts/mesh/freq.sh
source /opt/scripts/mesh/node_addresses.sh

FREQ=${chan[$CHANNEL_NUMBER]:-}
if [[ -z "${FREQ}" ]]; then
	echo "Channel ${CHANNEL_NUMBER} not found in freq.sh."
	exit 1
fi

normalize_node_name() {
	local node="$1"

	node=${node%%.*}

	if [[ -z "$node" || "$node" == "-" ]]; then
		echo ""
		return 0
	fi

	if [[ "$node" =~ ^[0-9]+$ ]]; then
		printf "apu%02d\n" "$((10#$node))"
		return 0
	fi

	if [[ "$node" =~ ^apu([0-9]+)$ ]]; then
		printf "apu%02d\n" "$((10#${BASH_REMATCH[1]}))"
		return 0
	fi

	echo "$node"
}

node_name_to_id() {
	local node="$1"

	node=$(normalize_node_name "$node")
	if [[ "$node" =~ ^apu([0-9]+)$ ]]; then
		echo "$((10#${BASH_REMATCH[1]}))"
		return 0
	fi

	return 1
}

node_id_in_topology() {
	local node_id="$1"
	local topo_node

	for topo_node in "${NODES[@]}"; do
		[[ "$topo_node" -eq "$node_id" ]] && return 0
	done

	return 1
}

HOSTNAME=$(uname -n | cut -d. -f1)
if [[ " ${neighbors["all"]} " != *" ${HOSTNAME} "* ]]; then
	echo "Hostname ${HOSTNAME} is not listed in node_addresses.sh."
	exit 1
fi

HOSTNAME=$(normalize_node_name "$HOSTNAME")
AP_NODE=$(normalize_node_name "$AP_NODE")

if [[ "$MODE" == "topology" ]]; then
	CURRENT_NODE_ID=$(node_name_to_id "$HOSTNAME") || {
		echo "Could not parse numeric node ID from hostname ${HOSTNAME}."
		exit 1
	}

	AP_NODE_ID=$(node_name_to_id "$AP_NODE") || {
		echo "Could not parse numeric node ID from AP node ${AP_NODE}."
		exit 1
	}

	source /opt/scripts/mesh/topology.sh
	TOPO="$TOPOLOGY"
	set_topology "$TOPOLOGY"

	if ! node_id_in_topology "$AP_NODE_ID"; then
		echo "AP node ${AP_NODE} is not part of topology ${TOPOLOGY}."
		exit 1
	fi

	if ! node_id_in_topology "$CURRENT_NODE_ID"; then
		echo "Hostname ${HOSTNAME} is not part of topology ${TOPOLOGY}. Skipping ap/sta init."
		exit 0
	fi

	if [[ "$HOSTNAME" == "$AP_NODE" ]]; then
		ROLE="ap"
	else
		ROLE="sta"
	fi
fi

if [[ "$ROLE" != "ap" && "$ROLE" != "sta" ]]; then
	echo "Role must be 'ap' or 'sta'."
	exit 1
fi

PHY_NUMBER="${PHYNAME//[!0-9]/}"
DEV_NUMBER="${DEVNAME//[!0-9]/}"
PHY_PATH="/sys/kernel/debug/ieee80211/${PHYNAME}"
RUNTIME_DIR="/run/apu-ap-sta"
PID_FILE="${RUNTIME_DIR}/hostapd-${DEVNAME}.pid"
CONF_FILE="${RUNTIME_DIR}/hostapd-${DEVNAME}.conf"
LOG_FILE="${RUNTIME_DIR}/hostapd-${DEVNAME}.log"
HOSTAPD_UNIT="ap-sta-hostapd-${DEVNAME}.service"
SUPPLICANT_PID_FILE="${RUNTIME_DIR}/wpa_supplicant-${DEVNAME}.pid"
SUPPLICANT_CONF_FILE="${RUNTIME_DIR}/wpa_supplicant-${DEVNAME}.conf"
SUPPLICANT_LOG_FILE="${RUNTIME_DIR}/wpa_supplicant-${DEVNAME}.log"
SUPPLICANT_UNIT="ap-sta-wpa-supplicant-${DEVNAME}.service"

mkdir -p "$RUNTIME_DIR"

apply_system_tuning() {
	sysctl -w net.core.rmem_max=8388608
	sysctl -w net.core.wmem_max=8388608
	sysctl -w net.core.rmem_default=8388608
	sysctl -w net.core.wmem_default=8388608
	sysctl -w net.core.optmem_max=8388608
	sysctl -w net.ipv4.tcp_mem='8388608 8388608 8388608'
	sysctl -w net.ipv4.tcp_rmem='4096 87380 8388608'
	sysctl -w net.ipv4.tcp_wmem='4096 65536 8388608'
	sysctl -w net.ipv4.udp_mem='8388608 8388608 8388608'
	sysctl -w net.ipv4.udp_rmem_min=4096
	sysctl -w net.ipv4.udp_wmem_min=4096
	sysctl -w net.core.netdev_max_backlog=5000
	sysctl -w net.ipv4.tcp_max_syn_backlog=5000
	sysctl -w net.ipv4.tcp_syn_retries=9
	sysctl -w net.ipv4.tcp_synack_retries=9
	sysctl -w net.ipv4.tcp_no_metrics_save=1
	sysctl -w net.ipv4.route.flush=1
	echo 0 > /sys/module/tcp_cubic/parameters/hystart
	echo 0 > /sys/module/tcp_cubic/parameters/hystart_detect
	sysctl -w net.ipv4.neigh.default.base_reachable_time_ms=1200000
	sysctl -w net.ipv4.neigh.default.mcast_solicit=25
	sysctl -w net.ipv4.neigh.default.ucast_solicit=25
}

stop_previous_hostapd() {
	systemctl stop "$HOSTAPD_UNIT" 2>/dev/null || true
	systemctl reset-failed "$HOSTAPD_UNIT" 2>/dev/null || true

	if [[ -f "$PID_FILE" ]]; then
		OLD_PID=$(cat "$PID_FILE")
		if kill -0 "$OLD_PID" 2>/dev/null; then
			kill "$OLD_PID"
			sleep 1
		fi
		rm -f "$PID_FILE"
	fi

	rm -f "$CONF_FILE" "$LOG_FILE"
}

stop_previous_wpa_supplicant() {
	systemctl stop "$SUPPLICANT_UNIT" 2>/dev/null || true
	systemctl reset-failed "$SUPPLICANT_UNIT" 2>/dev/null || true

	if [[ -f "$SUPPLICANT_PID_FILE" ]]; then
		OLD_PID=$(cat "$SUPPLICANT_PID_FILE")
		if kill -0 "$OLD_PID" 2>/dev/null; then
			kill "$OLD_PID"
			sleep 1
		fi
		rm -f "$SUPPLICANT_PID_FILE"
	fi

	rm -f "$SUPPLICANT_CONF_FILE" "$SUPPLICANT_LOG_FILE"
}

cleanup_phy_interfaces() {
	local iface

	stop_previous_hostapd
	stop_previous_wpa_supplicant

	while read -r iface; do
		[[ -z "$iface" ]] && continue
		echo "Removing existing interface ${iface} from phy${PHY_NUMBER}"
		ip addr flush dev "$iface" 2>/dev/null || true
		ip link set "$iface" down 2>/dev/null || true
		iw dev "$iface" disconnect 2>/dev/null || true
		iw dev "$iface" del 2>/dev/null || true
	done < <(iw dev | awk -v phy="$PHY_NUMBER" '
		$1 == "phy#"phy { in_phy=1; next }
		in_phy && $1 ~ /^phy#/ { in_phy=0 }
		in_phy && $1 == "Interface" { print $2 }
	')
}

create_interface() {
	if [[ ! -d "$PHY_PATH" ]]; then
		echo "${PHYNAME} does not exist in ${PHY_PATH}."
		exit 1
	fi

	iw reg set US
	iw phy "$PHYNAME" interface add "$DEVNAME" type managed
	if ! iw phy "$PHYNAME" set antenna 3 3 2>/dev/null; then
		echo "Skipping unsupported antenna configuration on ${PHYNAME}"
	fi
	iw phy "$PHYNAME" set retry short 7 long 4
	iw phy "$PHYNAME" set rts off
	ip link set "$DEVNAME" up
	iw dev "$DEVNAME" set power_save off
	iw dev "$DEVNAME" set txpower fixed "$TXPOWER"

	sysctl -w "net.ipv4.conf.${DEVNAME}.arp_ignore=1"
	sysctl -w "net.ipv4.conf.${DEVNAME}.arp_announce=2"
}

apply_radio_experiment_settings() {
	if [[ -w "${PHY_PATH}/ath9k/ani" ]]; then
		echo 0 > "${PHY_PATH}/ath9k/ani"
	fi

	if [[ -w "${PHY_PATH}/ath9k/nf_override" ]]; then
		echo "${NOISEFLOOR}" > "${PHY_PATH}/ath9k/nf_override"
	fi

	if [[ -w "${PHY_PATH}/rc/fixed_rate_idx" ]]; then
		echo 198 > "${PHY_PATH}/rc/fixed_rate_idx"
	fi

	if [[ "$CHANNEL_NUMBER" -lt 15 ]]; then
		iw dev "$DEVNAME" set bitrates ht-mcs-2.4 6 sgi-2.4
	else
		iw dev "$DEVNAME" set bitrates ht-mcs-5 6 sgi-5
	fi
}

assign_static_ip() {
	local node_ipv4=""
	local node_ipv6=""
	local broadcast_ip=""

	ip addr flush dev "$DEVNAME"

	case "$DEV_NUMBER" in
		0)
			node_ipv4=${phy0IP["$HOSTNAME"]:-}
			node_ipv6=${phy0IPv6["$HOSTNAME"]:-}
			broadcast_ip=${phy0IP["broadcast"]}
			;;
		1)
			node_ipv4=${phy1IP["$HOSTNAME"]:-}
			node_ipv6=${phy1IPv6["$HOSTNAME"]:-}
			broadcast_ip=${phy1IP["broadcast"]}
			;;
		*)
			echo "Unsupported device number parsed from ${DEVNAME}."
			exit 1
			;;
	esac

	if [[ -z "$node_ipv4" ]]; then
		echo "No IPv4 mapping found for ${HOSTNAME}/${DEVNAME}."
		exit 1
	fi

	ip -4 addr add "${node_ipv4}/24" broadcast "$broadcast_ip" dev "$DEVNAME"
	if [[ -n "$node_ipv6" ]]; then
		ip -6 addr add "${node_ipv6}/64" dev "$DEVNAME" 2>/dev/null || true
	fi

	echo "Assigned ${node_ipv4} to ${DEVNAME}"
}

write_hostapd_conf() {
	local hw_mode="a"

	if [[ "$CHANNEL_NUMBER" -lt 15 ]]; then
		hw_mode="g"
	fi

	cat > "$CONF_FILE" <<EOF
interface=${DEVNAME}
driver=nl80211
ssid=${SSID}
hw_mode=${hw_mode}
channel=${CHANNEL_NUMBER}
country_code=US
beacon_int=100
wmm_enabled=1
ieee80211n=1
ht_capab=[HT40+]
wpa=2
wpa_key_mgmt=WPA-PSK
rsn_pairwise=CCMP
wpa_passphrase=${WPA_PASSPHRASE}
EOF
}

start_ap() {
	local hostapd_pid=""

	command -v hostapd >/dev/null 2>&1 || {
		echo "hostapd is required for AP mode."
		exit 1
	}

	write_hostapd_conf
	systemd-run \
		--unit "$HOSTAPD_UNIT" \
		--collect \
		--property Restart=no \
		/bin/sh -lc "exec hostapd '$CONF_FILE' >>'$LOG_FILE' 2>&1" >/dev/null
	sleep 2

	if ! systemctl is-active --quiet "$HOSTAPD_UNIT"; then
		echo "hostapd failed to start. Log:"
		journalctl -u "$HOSTAPD_UNIT" -n 50 --no-pager || true
		cat "$LOG_FILE"
		exit 1
	fi

	hostapd_pid=$(systemctl show -p MainPID --value "$HOSTAPD_UNIT" 2>/dev/null || true)
	if [[ -n "$hostapd_pid" && "$hostapd_pid" != "0" ]]; then
		echo "$hostapd_pid" > "$PID_FILE"
	fi

	echo "AP mode active on ${DEVNAME} (${SSID}, channel ${CHANNEL_NUMBER})"
}

write_wpa_supplicant_conf() {
	local ap_bssid="$1"

	cat > "$SUPPLICANT_CONF_FILE" <<EOF
ctrl_interface=/run/wpa_supplicant
ap_scan=1
network={
	ssid="${SSID}"
	psk="${WPA_PASSPHRASE}"
	key_mgmt=WPA-PSK
	proto=RSN
	pairwise=CCMP
	group=CCMP
	scan_ssid=1
	freq_list=${FREQ}
EOF

	if [[ -n "$ap_bssid" ]]; then
		cat >> "$SUPPLICANT_CONF_FILE" <<EOF
	bssid=${ap_bssid}
EOF
	fi

	cat >> "$SUPPLICANT_CONF_FILE" <<'EOF'
}
EOF
}

start_wpa_supplicant() {
	local ap_bssid="$1"
	local supp_pid=""

	command -v wpa_supplicant >/dev/null 2>&1 || {
		echo "wpa_supplicant is required for STA mode with WPA2."
		exit 1
	}

	write_wpa_supplicant_conf "$ap_bssid"
	systemd-run \
		--unit "$SUPPLICANT_UNIT" \
		--collect \
		--property Restart=no \
		/bin/sh -lc "exec wpa_supplicant -D nl80211 -i '$DEVNAME' -c '$SUPPLICANT_CONF_FILE' >>'$SUPPLICANT_LOG_FILE' 2>&1" >/dev/null
	sleep 2

	if ! systemctl is-active --quiet "$SUPPLICANT_UNIT"; then
		echo "wpa_supplicant failed to start. Log:"
		journalctl -u "$SUPPLICANT_UNIT" -n 50 --no-pager || true
		cat "$SUPPLICANT_LOG_FILE"
		exit 1
	fi

	supp_pid=$(systemctl show -p MainPID --value "$SUPPLICANT_UNIT" 2>/dev/null || true)
	if [[ -n "$supp_pid" && "$supp_pid" != "0" ]]; then
		echo "$supp_pid" > "$SUPPLICANT_PID_FILE"
	fi
}

resolve_ap_bssid() {
	local ap_bssid=""

	[[ -z "$AP_NODE" ]] && return 0

	case "$DEV_NUMBER" in
		0) ap_bssid=${phy0MAC["$AP_NODE"]:-} ;;
		1) ap_bssid=${phy1MAC["$AP_NODE"]:-} ;;
	esac

	if [[ -z "$ap_bssid" ]]; then
		echo "Could not resolve AP MAC for ${AP_NODE}/${DEVNAME}."
		exit 1
	fi

	echo "$ap_bssid"
}

sta_is_connected() {
	iw dev "$DEVNAME" link | grep -q '^Connected to '
}

wait_for_sta_connected() {
	local max_wait_s=${1:-10}
	local waited=0

	while (( waited < max_wait_s )); do
		if sta_is_connected; then
			return 0
		fi
		sleep 1
		waited=$((waited + 1))
	done

	return 1
}

connect_sta() {
	local ap_bssid=""
	local attempt

	ap_bssid=$(resolve_ap_bssid)

	for attempt in $(seq 1 10); do
		echo "STA connect attempt ${attempt}/10 on ${DEVNAME}"
		iw dev "$DEVNAME" disconnect 2>/dev/null || true
		stop_previous_wpa_supplicant
		start_wpa_supplicant "$ap_bssid"
		if wait_for_sta_connected; then
			iw dev "$DEVNAME" link
			echo "STA connected to ${SSID}"
			return 0
		fi

		iw dev "$DEVNAME" link || true
		sleep 2
	done

	echo "STA failed to connect to ${SSID} on ${DEVNAME}."
	exit 1
}

echo "Configuring ${HOSTNAME}: role=${ROLE} dev=${DEVNAME} phy=${PHYNAME} channel=${CHANNEL_NUMBER} mode=${MODE}${TOPOLOGY:+ topo=${TOPOLOGY}}${AP_NODE:+ ap=${AP_NODE}} auth=wpa2-psk"

apply_system_tuning
systemctl stop systemd-timesyncd.service || true
cleanup_phy_interfaces
create_interface
apply_radio_experiment_settings

case "$ROLE" in
	ap) start_ap ;;
	sta) connect_sta ;;
esac

assign_static_ip

echo "Infrastructure mode setup complete for ${DEVNAME}"