#!/bin/bash
################################################################################
# Copyright 2021 ModalAI Inc.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its contributors
#    may be used to endorse or promote products derived from this software
#    without specific prior written permission.
#
# 4. The Software is used solely in conjunction with devices provided by
#    ModalAI Inc.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
################################################################################

UTIL_NAME="voxl-wifi"
DEFAULT_SSID="SSID"
DEFAULT_PASS="1234567890"
PLATFORM=$(voxl-platform)
conf_file=wpa_supplicant.conf

USER=$(whoami)

# Interface variables (set by detect_wifi_interfaces)
AP_INTERFACE=""
STA_INTERFACE=""
SUPPORTS_CONCURRENT=false

#
# Detect available WiFi interfaces and concurrent capability
#
detect_wifi_interfaces() {
	# Check for STA+AP capable interfaces
	if [ -d /sys/class/net/uap0 ] && [ -d /sys/class/net/mlan0 ]; then
		AP_INTERFACE="uap0"
		STA_INTERFACE="mlan0"
		SUPPORTS_CONCURRENT=true
		echo "Detected concurrent-capable WiFi: AP=$AP_INTERFACE, STA=$STA_INTERFACE"
	else
		AP_INTERFACE="wlan0"
		STA_INTERFACE="wlan0"
		SUPPORTS_CONCURRENT=false
		echo "Detected WiFi interface: $AP_INTERFACE"
	fi
}

# sanity checks
if [ "${USER}" != "root" ]; then
	echo "Run this script as the root user"
	exit 1
fi


#
# Display the usage of this utility
#
print_usage () {
	echo "Description:"
	echo "        Configure the Wi-Fi mode."
	echo ""
	echo "Usage:"
	echo "  ${UTIL_NAME} getmode"
	echo "        Print the current mode (softap, station, sta+ap, or disabled)"
	echo ""
	echo "  ${UTIL_NAME} enable-station <ssid> [password]"
	echo "        Enable station mode (concurrent with existing AP if hardware supports it)"
	echo "        If password is omitted, assumes open network"
	echo ""
	echo "  ${UTIL_NAME} enable-softap <ssid> [password] [2ghz|5ghz]"
	echo "        Enable softap mode (concurrent with existing station if hardware supports it)"
	echo "        Band defaults to 5ghz. Password defaults to '${DEFAULT_PASS}'"
	echo ""
	echo "  ${UTIL_NAME} disable-station"
	echo "        Disable station mode only"
	echo ""
	echo "  ${UTIL_NAME} disable-softap"
	echo "        Disable softap mode only"
	echo ""
	echo "  ${UTIL_NAME} disable-all"
	echo "        Disable all WiFi modes"
	echo ""
	exit 0
}


#
# Before we choose which mode to use, start from a known state
#
function stop_wlan () {
	# Stop all WiFi services
	stop_station
	stop_softap
}

function stop_station () {
	# Stop station mode services (both old and new style)
	systemctl disable wpa_supplicant 2>/dev/null >/dev/null
	systemctl stop wpa_supplicant 2>/dev/null >/dev/null
	
	systemctl disable wpa_supplicant@${STA_INTERFACE}.service 2>/dev/null >/dev/null
	systemctl stop wpa_supplicant@${STA_INTERFACE}.service 2>/dev/null >/dev/null

	systemctl disable systemd-resolved 2>/dev/null >/dev/null
	systemctl stop systemd-resolved 2>/dev/null >/dev/null

	# Clean up station files
	if [ -f $station_ssid_file ]; then
		rm $station_ssid_file
	fi
	if [ -f $station_band_file ]; then
		rm $station_band_file
	fi
	if [ -f $station_interface_file ]; then
		rm $station_interface_file
	fi

	# Flush station interface if it exists
	if [ -d /sys/class/net/$STA_INTERFACE ]; then
		ip link set $STA_INTERFACE down 2>/dev/null
		ip addr flush dev $STA_INTERFACE 2>/dev/null
		ip link set $STA_INTERFACE up 2>/dev/null
	fi
}

function stop_softap () {
	# Stop softap mode services  
	systemctl disable dnsmasq 2>/dev/null >/dev/null
	systemctl stop dnsmasq 2>/dev/null >/dev/null

	systemctl disable hostapd 2>/dev/null >/dev/null
	systemctl stop hostapd 2>/dev/null >/dev/null

	systemctl stop voxl-softap 2>/dev/null >/dev/null
	systemctl disable voxl-softap 2>/dev/null >/dev/null

	# remove any old hostapd interface files
	if [ -f /data/misc/wifi/hostapd/wlan0 ]; then
		rm /data/misc/wifi/hostapd/wlan0
	fi
	if [ -f /data/misc/wifi/hostapd/uap0 ]; then
		rm /data/misc/wifi/hostapd/uap0
	fi
	if [ -f $softap_interface_file ]; then
		rm $softap_interface_file
	fi

	# cleanup system-resolve leftover symlinks
	if [ -f /etc/resolv.conf ]; then
        rm /etc/resolv.conf
    fi

	# Flush AP interface if it exists
	if [ -d /sys/class/net/$AP_INTERFACE ]; then
		ip link set $AP_INTERFACE down 2>/dev/null
		ip addr flush dev $AP_INTERFACE 2>/dev/null
		ip link set $AP_INTERFACE up 2>/dev/null
	fi
}


#
# modify conf files for softap
#
function setup_softap_conf_files () {

	if [ "$1" == "softap" ]; then
		cp -f /usr/share/modalai/system-tweaks/hostapd_2ghz ${hostapd_conf_file}
	else
		# Use minimal 5GHz config for uap0 (no advanced features)
		if [ "$SUPPORTS_CONCURRENT" == "true" ]; then
			cp -f /usr/share/modalai/system-tweaks/hostapd_5ghz_uap0 ${hostapd_conf_file}
		else
			cp -f /usr/share/modalai/system-tweaks/hostapd_5ghz_ac ${hostapd_conf_file}
		fi
	fi

	# configure hostapd conf file with ssid and interface
	echo "editing ${hostapd_conf_file} for softap mode, ssid is: ${ssid}, interface: ${AP_INTERFACE}"
	sed -i "/^ssid=/c\ssid=${ssid}" ${hostapd_conf_file}
	sed -i "/^wpa_passphrase=/c\wpa_passphrase=${softap_pass}" ${hostapd_conf_file}
	sed -i "/^interface=/c\interface=${AP_INTERFACE}" ${hostapd_conf_file}
	
	# Store softap interface for mismatch detection
	echo "${AP_INTERFACE}" > $softap_interface_file

	sync
}


################################################################################
## actual start of execution
################################################################################

# Detect WiFi interfaces first
detect_wifi_interfaces

if [ $# -eq 0 ]; then
	echo "Illegal usage"
	print_usage
	exit 1
fi

# print help message if requested
if [ $1 == "-h" ] ||  [ $1 == "--help" ]; then
	print_usage
	exit 0
fi

# Validate command arguments
VALID_COMMANDS="getmode enable-station enable-softap disable-station disable-softap disable-all"
if [[ ! " $VALID_COMMANDS " =~ " $1 " ]]; then
	echo "Illegal command: $1"
	print_usage
	exit 1
fi

# WiFi configuration file paths
wifi_setup_dir=/data/misc/wifi
hostapd_conf_file=/data/misc/wifi/hostapd.conf
wpa_supplicant_conf_file=/data/misc/wifi/wpa_supplicant.conf
station_ssid_file=/data/misc/wifi/station_ssid
station_band_file=/data/misc/wifi/station_band
station_interface_file=/data/misc/wifi/station_interface
softap_interface_file=/data/misc/wifi/softap_interface
mode_file=$wifi_setup_dir/wlan_mode

# Mode commands
if [ $1 == "getmode" ]; then
	if [ -f $mode_file ]; then
		echo $(cat $mode_file)
	else
		echo "softap"
	fi

# New concurrent mode commands
elif [ "$1" == "enable-station" ]; then
	if [[ $# -lt 2 ]]; then
		echo "${UTIL_NAME} enable-station: missing ssid parameter"
		exit 1
	fi

	stop_wlan
	
	station_ssid="$2"
	station_key="${3:-}"
	
	# Scan for target network and save band info before connecting
	echo "Scanning for network '$station_ssid'..."
	
	# Test scan first
	scan_result=$(iw dev $STA_INTERFACE scan 2>&1)
	scan_exit_code=$?
	
	if [ $scan_exit_code -ne 0 ]; then
		echo "Scan failed with exit code $scan_exit_code: $scan_result"
		echo "Could not scan for network '$station_ssid' - proceeding anyway"
	else
		echo "Scan succeeded, looking for network '$station_ssid'"
		
		# Simple grep approach - look for exact SSID match and get freq from same BSS block
		freq=$(echo "$scan_result" | awk -v ssid="$station_ssid" '
		/^BSS / { in_bss = 1; current_freq = "" }
		in_bss && /freq:/ { current_freq = $2 }
		in_bss && /SSID:/ { 
			gsub(/.*SSID: /, "")
			if ($0 == ssid && current_freq != "") {
				print current_freq
				exit
			}
		}
		')
		
		if [ ! -z "$freq" ] && [ "$freq" -gt 0 ] 2>/dev/null; then
			if [ $freq -ge 2400 ] && [ $freq -le 2500 ]; then
				echo "2ghz" > $station_band_file
				echo "Network '$station_ssid' detected on 2ghz (freq: $freq)"
			elif [ $freq -ge 5000 ] && [ $freq -le 6000 ]; then
				echo "5ghz" > $station_band_file
				echo "Network '$station_ssid' detected on 5ghz (freq: $freq)"
			else
				echo "Network '$station_ssid' found but frequency $freq not recognized"
			fi
		else
			echo "Could not find network '$station_ssid' in scan results - proceeding anyway"
		fi
	fi
	
	echo "Enabling station mode on $STA_INTERFACE..."
	# Don't stop existing services in concurrent mode
	
	echo "creating wpa_supplicant configuration for station mode"
	cat > ${wpa_supplicant_conf_file} << EOF2
# Only WPA-PSK is used. Any valid cipher combination is accepted.
ctrl_interface=/var/run/wpa_supplicant
update_config=1

network={
#Open
#        ssid="example open network"
#        key_mgmt=NONE
#WPA-PSK
	ssid="${station_ssid}"
	#proto=WPA
	key_mgmt=WPA-PSK
	pairwise=TKIP CCMP
	group=TKIP CCMP
	psk="${station_key}"
#WEP
#	ssid="example wep network"
#	key_mgmt=NONE
#	wep_key0="abcde"
#	wep_key1=0102030405
#	wep_tx_keyidx=0
}
wowlan_triggers=magic_pkt

EOF2


	# Store station SSID and interface for status display and mismatch detection
	echo "${station_ssid}" > ${station_ssid_file}
	echo "${STA_INTERFACE}" > $station_interface_file

	sync
	
	# Use interface-specific wpa_supplicant service
	systemctl disable wpa_supplicant 2>/dev/null >/dev/null
	systemctl stop wpa_supplicant 2>/dev/null >/dev/null
	
	systemctl enable wpa_supplicant@${STA_INTERFACE}.service 2>/dev/null >/dev/null
	systemctl start wpa_supplicant@${STA_INTERFACE}.service 2>/dev/null >/dev/null
	
	systemctl enable systemd-resolved 2>/dev/null >/dev/null  
	systemctl start systemd-resolved 2>/dev/null >/dev/null
	
	# Update mode tracking
	if [ -f $mode_file ] && grep -q "softap" $mode_file; then
		echo "sta+ap" > $mode_file
	else
		echo "station" > $mode_file
	fi
	
	echo "Station mode enabled on $STA_INTERFACE"

elif [ "$1" == "enable-softap" ]; then
	if [[ $# -lt 2 ]]; then
		echo "${UTIL_NAME} enable-softap: missing ssid parameter"
		exit 1
	fi
	
	ssid="$2"
	softap_pass="${3:-$DEFAULT_PASS}"
	band="${4:-5ghz}"
	
	echo "Enabling softap mode on $AP_INTERFACE..."
	
	# Disable and flush existing softap setup
	stop_softap
	
	if [ "$band" == "2ghz" ]; then
		setup_softap_conf_files "softap"
	else
		setup_softap_conf_files "softap5"  
	fi
	
	echo "Enabling voxl-softap service, this may take a bit"
	systemctl enable voxl-softap 2>/dev/null >/dev/null
	systemctl start voxl-softap 2>/dev/null >/dev/null
	
	# Update mode tracking
	if [ -f $wifi_setup_dir/wlan_mode ] && grep -q "station" $wifi_setup_dir/wlan_mode; then
		echo "sta+ap" > $wifi_setup_dir/wlan_mode
	else
		echo "softap" > $wifi_setup_dir/wlan_mode
	fi
	
	echo "SoftAP mode enabled on $AP_INTERFACE"

elif [ "$1" == "disable-station" ]; then
	echo "Disabling station mode..."
	stop_station
	
	# Remove station SSID file and band file
	if [ -f $station_ssid_file ]; then
		rm $station_ssid_file
	fi
	if [ -f /data/misc/wifi/station_band ]; then
		rm /data/misc/wifi/station_band
	fi
	
	# Update mode tracking
	if [ -f $wifi_setup_dir/wlan_mode ] && grep -q "sta+ap" $wifi_setup_dir/wlan_mode; then
		echo "softap" > $wifi_setup_dir/wlan_mode
	else
		echo "disabled" > $wifi_setup_dir/wlan_mode
	fi
	
	echo "Station mode disabled"

elif [ "$1" == "disable-softap" ]; then
	echo "Disabling softap mode..."
	stop_softap
	
	# Update mode tracking  
	if [ -f $wifi_setup_dir/wlan_mode ] && grep -q "sta+ap" $wifi_setup_dir/wlan_mode; then
		echo "station" > $wifi_setup_dir/wlan_mode
	else
		echo "disabled" > $wifi_setup_dir/wlan_mode
	fi
	
	echo "SoftAP mode disabled"

elif [ "$1" == "disable-all" ]; then
	echo "Disabling all WiFi modes..."
	stop_wlan
	
	# Remove station SSID file and band file
	if [ -f $station_ssid_file ]; then
		rm $station_ssid_file
	fi
	if [ -f /data/misc/wifi/station_band ]; then
		rm /data/misc/wifi/station_band
	fi
	
	echo "disabled" > $wifi_setup_dir/wlan_mode
else
	echo "Illegal command: $1"
	print_usage
	exit 1
fi

sync