#!/usr/bin/env python3
import os
import sys
import subprocess


reset  = "\033[0m"
red    = "\033[91m"
yellow = "\033[33m"
green  = "\033[32m"
cyan   = "\033[36m"
bold   = "\033[1m"
dim    = "\033[2m"


def print_usage():
    interfaces = get_wifi_interfaces()
    hw_type = "Concurrent STA+AP capable" if interfaces['concurrent'] else "Single-mode STA or AP"
    hw_interfaces = f"(AP:{interfaces['ap']}, STA:{interfaces['sta']})" if interfaces['concurrent'] else f"({interfaces['ap']})"
    
    print(f"""Usage -- {red}voxl-wifi
{bold+green}Description:{reset}
    Configure WiFi modes on {hw_type} hardware {hw_interfaces}
    
{bold+green}Usage:{reset}
    voxl-wifi getmode
        {dim}Print the current Wi-Fi mode{reset}

    voxl-wifi station <ssid> [password]
        {dim}Set Wi-Fi configuration to factory defaults. This will set the Wi-Fi to {reset}
        {dim}SSID and password (if provided).{reset}

    voxl-wifi softap <ssid> [password]
        {dim}Set Wi-Fi to 2.4ghz access point mode (ch6). Default password is '1234567890'{reset}
    
    voxl-wifi softap5 <ssid> [password]
        {dim}Set Wi-Fi to 5ghz access point mode (ch149). Default password is '1234567890'{reset}

    voxl-wifi disable-station
        {dim}Disable station mode (if currently active){reset}
    
    voxl-wifi disable-softap
        {dim}Disable softap mode (if currently active){reset}
    
    voxl-wifi disable-all
        {dim}Disable all active WiFi modes{reset}
    
    voxl-wifi factory
        {dim}Set Wi-Fi configuration to factory defaults. This will set the Wi-Fi to{reset}
        {dim}5ghz access point mode, with SSID 'VOXL-[MAC address]' and password '1234567890'.{reset}

    voxl-wifi
        {dim}Launch interactive configuration wizard{reset}
""")

    # Show USB options only on QCS6490
    if is_qcs6490() and has_usb_interface():
        print(f"""
{bold+green}USB Interface (QCS6490 only):{reset}
    voxl-wifi usb-enable
        {dim}Enable USB network interface with DHCP server (192.168.100.x){reset}

    voxl-wifi usb-disable
        {dim}Disable USB network interface DHCP server{reset}
""")


def print_error(err):
    print(f"{bold+red}Error:{reset} {err}")

def print_success(msg):
    print(f"{green}{msg}{reset}")

def print_debug(msg):
    if "VOXL_WIFI_DEBUG" in os.environ:
        print(f"{yellow}[DEBUG]\t{msg}{reset}")


def call_voxl_wifi_legacy(args):
    cmd = ["/usr/bin/voxl-wifi-legacy"]
    cmd.extend(args)

    print_debug(f"running {cmd}")

    try:
        subprocess.run(cmd)
    except:
        print_error("voxl-wifi-legacy failed\n")
        sys.exit(1)


SYS_NET = "/sys/class/net"


def supports_concurrent_modes():
    """Check if hardware supports concurrent STA+AP modes"""

    # Botsu-sparrow support STA+AP when in mode 3
    if os.path.exists('/sys/module/botsu_sparrow') and (os.path.exists('/sys/class/net/uap0') and os.path.exists('/sys/class/net/mlan0')):
        return True

    return False


def is_qcs6490():
    """Check if running on QCS6490 platform"""
    try:
        result = subprocess.run(['voxl-chip'], capture_output=True, text=True)
        return 'qcs6490' in result.stdout.lower()
    except:
        return False


def has_usb_interface():
    """Check if usb0 interface exists"""
    return os.path.isdir(f"{SYS_NET}/usb0")


def get_usb_status():
    """Get current USB interface status"""
    if not has_usb_interface():
        return {'exists': False, 'enabled': False, 'ip': None, 'up': False}

    # Check if systemd-networkd config exists and has DHCPServer enabled
    config_file = "/etc/systemd/network/10-usb0.network"
    enabled = False
    if os.path.exists(config_file):
        try:
            with open(config_file, 'r') as f:
                content = f.read()
                enabled = 'DHCPServer=yes' in content
        except:
            pass

    # Check if interface is up
    up = False
    try:
        with open(f"{SYS_NET}/usb0/operstate", 'r') as f:
            state = f.read().strip()
            up = state in ['up', 'unknown']  # 'unknown' often means up for gadget
    except:
        pass

    # Get current IP
    ip = None
    try:
        result = subprocess.run(['ip', '-4', 'addr', 'show', 'usb0'],
                                capture_output=True, text=True)
        for line in result.stdout.split('\n'):
            if 'inet ' in line:
                ip = line.strip().split()[1].split('/')[0]
                break
    except:
        pass

    return {'exists': True, 'enabled': enabled, 'ip': ip, 'up': up}


def enable_usb_interface():
    """Enable USB interface with DHCP server"""
    config_content = """[Match]
Name=usb0

[Network]
Address=192.168.100.1/24
DHCPServer=yes
IPForward=yes

[DHCPServer]
PoolOffset=10
PoolSize=90
EmitDNS=yes
DNS=8.8.8.8
DefaultLeaseTimeSec=3600
"""
    config_file = "/etc/systemd/network/10-usb0.network"

    try:
        with open(config_file, 'w') as f:
            f.write(config_content)

        # Bring interface up
        subprocess.run(['ip', 'link', 'set', 'usb0', 'up'],
                       capture_output=True)

        # Enable and restart systemd-networkd
        subprocess.run(['systemctl', 'enable', 'systemd-networkd'],
                       capture_output=True)
        subprocess.run(['systemctl', 'restart', 'systemd-networkd'],
                       capture_output=True)
        return True
    except Exception as e:
        print_error(f"Failed to enable USB interface: {e}")
        return False


def disable_usb_interface():
    """Disable USB interface completely"""
    config_file = "/etc/systemd/network/10-usb0.network"

    try:
        # Create a config that keeps the interface down
        config_content = """[Match]
Name=usb0

[Link]
ActivationPolicy=down
Unmanaged=yes
"""
        with open(config_file, 'w') as f:
            f.write(config_content)

        # Bring interface down immediately
        subprocess.run(['ip', 'link', 'set', 'usb0', 'down'],
                       capture_output=True)

        # Restart systemd-networkd to apply
        subprocess.run(['systemctl', 'restart', 'systemd-networkd'],
                       capture_output=True)
        return True
    except Exception as e:
        print_error(f"Failed to disable USB interface: {e}")
        return False


def is_wireless(ifname):
    """Return True if interface is a real WiFi interface"""
    return (
        os.path.isdir(f"{SYS_NET}/{ifname}/wireless") or
        os.path.exists(f"{SYS_NET}/{ifname}/phy80211")
    )

def get_wifi_interfaces():
    """Detect available WiFi interfaces"""

    # 1) Concurrent AP + STA (common on some Qualcomm / Marvell setups)
    if (
        os.path.isdir(f"{SYS_NET}/uap0") and
        os.path.isdir(f"{SYS_NET}/mlan0")
    ):
        return {
            'ap': 'uap0',
            'sta': 'mlan0',
            'concurrent': True
        }

    # 2) Prefer wlan0 if it exists and is wireless
    if os.path.isdir(f"{SYS_NET}/wlan0") and is_wireless("wlan0"):
        return {
            'ap': 'wlan0',
            'sta': 'wlan0',
            'concurrent': False
        }

    # 3) Otherwise pick the first real wireless interface (wlx*, wlp*, etc.)
    for ifname in sorted(os.listdir(SYS_NET)):
        if not ifname.startswith("wl"):
            continue
        if not is_wireless(ifname):
            continue

        return {
            'ap': ifname,
            'sta': ifname,
            'concurrent': False
        }

    # Nothing found
    return {
        'ap': None,
        'sta': None,
        'concurrent': False
    }

def get_mac_address(interface=None):
    if interface is None:
        interfaces = get_wifi_interfaces()
        interface = interfaces['ap']
    
    cmd = "ifconfig " + interface + " | grep ether | cut -d' ' -f 10"
    result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if result.stderr:
        print(f"Error: {result.stderr.decode('utf-8')}")
    else:
        return result.stdout.decode('utf-8').strip()


def get_input(prompt, default=None, optional=False):
    default_str = f" [default: {default}]" if default is not None else ""
    prompt_str = f"{bold+yellow}{prompt}{reset+yellow}{default_str}{bold}: {reset}"

    res = input(prompt_str)
    if len(res) > 0:
        return res
    elif default != None:
        return default
    elif optional:
        return None
    else:
        print_error("Input is required")
        return get_input(prompt, default)


def get_input_options(prompt, options, default=None):
    print("Options:")
    for i, option in enumerate(options):
        print(f"\t{i+1}) {option}")

    # ensure string to appease pyright
    res = str(get_input(prompt, default)).strip()
    res_int = None
    if res.isnumeric():
        res_int = int(res)
    else:
        res_int = 0

    if res in options:
        return res
    elif res.isnumeric() and 0 < res_int <= len(options):
        return options[res_int - 1]
    else:
        print_error(f"Input must be one of: {options}, or 1 to {len(options)}")
        return get_input_options(prompt, options, default)


def run_station(ssid, passwd=""):
    call_voxl_wifi_legacy(["enable-station", ssid, passwd])
    print_success("Station Mode Configuration succeeded!")


def run_softap(ssid, passwd="1234567890"):
    call_voxl_wifi_legacy(["enable-softap", ssid, passwd, "2ghz"])
    print_success("SoftAP Configuration succeeded!")

def run_softap5(ssid, passwd="1234567890"):
    call_voxl_wifi_legacy(["enable-softap", ssid, passwd, "5ghz"])
    print_success("SoftAP Configuration succeeded!")

def get_current_modes():
    """Get current WiFi modes (station and/or softap status)"""
    mode = None
    try:
        with open("/data/misc/wifi/wlan_mode", 'r') as mode_file:
            mode = "".join(mode_file.readlines()).strip()
    except:
        mode = "undefined"
    
    # Parse mode into individual components
    modes = {
        'station': mode in ['station', 'sta+ap'],
        'softap': mode in ['softap', 'sta+ap'],
        'disabled': mode == 'disabled'
    }
    return modes, mode

def get_current_ssids():
    """Get current SSID for station and softap"""
    station_ssid = "undefined"
    softap_ssid = "undefined"
    
    # Get station SSID from dedicated station_ssid file
    try:
        with open("/data/misc/wifi/station_ssid", 'r') as station_file:
            station_ssid = station_file.read().strip()
    except:
        pass
    
    # Get softap SSID from hostapd.conf
    try:
        with open("/data/misc/wifi/hostapd.conf", 'r') as hostapd_conf_file:
            lines = hostapd_conf_file.readlines()
            ssid_lines = [line for line in lines if line.startswith("ssid=")]
            if ssid_lines:
                softap_ssid = ssid_lines[0].strip().split("=")[1]
    except:
        pass
        
    return station_ssid, softap_ssid

def run_getmode():
    modes, mode_str = get_current_modes()
    station_ssid, softap_ssid = get_current_ssids()
    interfaces = get_wifi_interfaces()
    
    print("WiFi is currently set up as follows:")
    print(f"\tMode: {dim}{mode_str}{reset}")
    
    any_mismatch = False
    
    if interfaces['concurrent']:
        print(f"\tHardware: {dim}Concurrent STA+AP capable (AP:{interfaces['ap']}, STA:{interfaces['sta']}){reset}")
        if modes['station']:
            has_mismatch, configured_interface = check_station_interface_mismatch()
            mismatch_text = f" {yellow}(configured for {configured_interface}, expected {interfaces['sta']}){reset}" if has_mismatch else ""
            print(f"\tStation: {dim}Active on {configured_interface} - \"{station_ssid}\"{reset}{mismatch_text}")
            any_mismatch = any_mismatch or has_mismatch
        else:
            print(f"\tStation: {dim}Disabled{reset}")
            
        if modes['softap']:
            has_mismatch, configured_interface = check_softap_interface_mismatch()
            mismatch_text = f" {yellow}(configured for {configured_interface}, expected {interfaces['ap']}){reset}" if has_mismatch else ""
            print(f"\tSoftAP: {dim}Active on {configured_interface} - \"{softap_ssid}\"{reset}{mismatch_text}")
            any_mismatch = any_mismatch or has_mismatch
        else:
            print(f"\tSoftAP: {dim}Disabled{reset}")
    else:
        print(f"\tHardware: {dim}single-mode ({interfaces['ap']}){reset}")
        if modes['station']:
            has_mismatch, configured_interface = check_station_interface_mismatch()
            mismatch_text = f" {yellow}(configured for {configured_interface}, expected {interfaces['sta']}){reset}" if has_mismatch else ""
            print(f"\tStation on {configured_interface}: {dim}{station_ssid}{reset}{mismatch_text}")
            any_mismatch = any_mismatch or has_mismatch
        elif modes['softap']:
            has_mismatch, configured_interface = check_softap_interface_mismatch()
            mismatch_text = f" {yellow}(configured for {configured_interface}, expected {interfaces['ap']}){reset}" if has_mismatch else ""
            print(f"\tSoftAP on {configured_interface}: {dim}{softap_ssid}{reset}{mismatch_text}")
            any_mismatch = any_mismatch or has_mismatch
    
    # Show summary warning if any mismatches detected
    if any_mismatch:
        print(f"\n{bold+yellow}WARNING:{reset} Interface mismatch detected!")
        print(f"{yellow}Configuration was set up for different hardware. Run 'voxl-wifi disable-all' and reconfigure.{reset}")

    # Show USB interface status on QCS6490
    if is_qcs6490() and has_usb_interface():
        usb_status = get_usb_status()
        if usb_status['enabled']:
            print(f"\n\tUSB Interface: {dim}Enabled (DHCP server on {usb_status['ip']}){reset}")
        else:
            print(f"\n\tUSB Interface: {dim}Disabled (no DHCP server){reset}")

def run_factory():
    # Generate default SSID using MAC address
    ssid="VOXL-SOFTAP"
    
    # use serial number
    if os.path.exists('/sys/devices/soc0/serial_number'):
        # Fallback to serial number if MAC address fails
        with open('/sys/devices/soc0/serial_number', 'r') as file:
            ssid = "VOXL-" + file.read().strip()
    else:
        mac = get_mac_address()
        if mac:
            ssid = "VOXL-" + mac
    
    # Disable all first, then enable 5GHz softap
    call_voxl_wifi_legacy(["disable-all"])
    call_voxl_wifi_legacy(["enable-softap", ssid, "1234567890", "5ghz"])
    print_success("Factory Mode Configuration succeeded!")

def get_current_softap_band():
    """Detect current softap band from hostapd.conf"""
    try:
        with open("/data/misc/wifi/hostapd.conf", 'r') as hostapd_file:
            content = hostapd_file.read()
            
            # Check hw_mode to determine band
            if "hw_mode=g" in content:
                return "2ghz"
            elif "hw_mode=a" in content:
                return "5ghz"
    except:
        pass
    return None

def run_station_config_wizard():
    """Unified station configuration wizard"""
    print(f"\n{cyan}Please enter the SSID to connect to{reset}")
    ssid = get_input("SSID")
    print(f"\n{cyan}Please enter the Wi-Fi password, or leave blank for no authentication.{reset}")
    passwd = get_input("Password", optional=True)
    
    print(f"\n{cyan}Configuring station mode...{reset}")
    
    # Use the smart station function with the collected inputs
    run_smart_station(ssid, passwd if passwd is not None else "")

def run_softap_config_wizard(band):
    """Unified SoftAP configuration wizard with band validation"""
    print(f"\n{cyan}Please enter the SSID to broadcast{reset}")
    ssid = get_input("SSID")
    print(f"\n{cyan}Please enter a password, or leave blank to use the default password.{reset}")

    passwd = str(get_input("Password", default='1234567890'))
    while len(passwd) < 8:
        print_error("Password must be at least 8 characters")
        passwd = str(get_input("Password", default='1234567890'))

    print(f"\n{cyan}Configuring {band} SoftAP mode...{reset}")
    
    # Use the smart softap function with the collected inputs
    run_smart_softap(ssid, passwd, band)

def get_station_band():
    """Get station band from saved file"""
    try:
        with open("/data/misc/wifi/station_band", 'r') as f:
            saved_band = f.read().strip()
            if saved_band in ["2ghz", "5ghz"]:
                return saved_band
    except:
        pass
    return None

def check_station_interface_mismatch():
    """Check if station interface configuration mismatches current hardware
    Returns: (has_mismatch: bool, configured_interface: str)"""
    interfaces = get_wifi_interfaces()
    try:
        with open("/data/misc/wifi/station_interface", 'r') as f:
            configured_interface = f.read().strip()
            has_mismatch = configured_interface != interfaces['sta']
            return has_mismatch, configured_interface
    except:
        # If no file exists, assume current interface
        return False, interfaces['sta']

def check_softap_interface_mismatch():
    """Check if softap interface configuration mismatches current hardware
    Returns: (has_mismatch: bool, configured_interface: str)"""
    interfaces = get_wifi_interfaces()
    try:
        with open("/data/misc/wifi/softap_interface", 'r') as f:
            configured_interface = f.read().strip()
            has_mismatch = configured_interface != interfaces['ap']
            return has_mismatch, configured_interface
    except:
        # If no file exists, assume current interface
        return False, interfaces['ap']


def run_wizard():
    """Unified wizard that adapts to hardware capabilities"""
    modes, _ = get_current_modes()
    interfaces = get_wifi_interfaces()
    
    # Show current status
    run_getmode()
    
    
    print("")
    print("Station mode connects to an existing WiFi network.")
    print("SoftAP mode creates a WiFi hotspot that other devices can connect to.")
    print("")
    
    if interfaces['concurrent']:
        print(f"{bold}Hardware: Concurrent STA+AP capable{reset} (can run both modes simultaneously)")
    else:
        print(f"{bold}Hardware: Single-mode STA or AP{reset} (exclusive modes)")
    
    print("")
    print(f"{cyan}Please select WiFi operation:{reset}")
    
    # Build dynamic menu based on hardware and current state
    options = []
    actions = {}
    
    # Configure options (always available for reconfiguration)
    options.append("Configure Station")
    actions["Configure Station"] = "config_station"
    
    options.append("Configure SoftAP (2.4GHz)")
    options.append("Configure SoftAP (5GHz)")
    actions["Configure SoftAP (2.4GHz)"] = "config_softap_2ghz"
    actions["Configure SoftAP (5GHz)"] = "config_softap_5ghz"
    
    # Disable options - individual options only for concurrent hardware
    if interfaces['concurrent']:
        if modes['station']:
            options.append("Disable Station")
            actions["Disable Station"] = "disable_station"
        
        if modes['softap']:
            options.append("Disable SoftAP")
            actions["Disable SoftAP"] = "disable_softap"
    
    # Always show disable wifi option regardless of hardware or current modes
    options.append("Disable WiFi (disables all active modes)")
    actions["Disable WiFi (disables all active modes)"] = "disable_all"
    
    # USB options (QCS6490 only)
    if is_qcs6490() and has_usb_interface():
        usb_status = get_usb_status()
        if usb_status['enabled']:
            options.append("Disable USB Interface")
            actions["Disable USB Interface"] = "usb_disable"
        else:
            options.append("Enable USB Interface")
            actions["Enable USB Interface"] = "usb_enable"

    # Always available options
    options.append("Factory Reset")
    options.append("Exit")
    actions["Factory Reset"] = "factory"
    actions["Exit"] = "exit"
    
    choice = get_input_options("Operation", options)
    action = actions.get(choice)
    
    # Execute action
    if action == "config_station":
        run_station_config_wizard()
    elif action == "config_softap_2ghz":
        run_softap_config_wizard("2ghz")
    elif action == "config_softap_5ghz":
        run_softap_config_wizard("5ghz")
    elif action == "disable_station":
        call_voxl_wifi_legacy(["disable-station"])
        print_success("Station mode disabled!")
    elif action == "disable_softap":
        call_voxl_wifi_legacy(["disable-softap"])
        print_success("SoftAP mode disabled!")
    elif action == "disable_all":
        call_voxl_wifi_legacy(["disable-all"])
        print_success("All WiFi modes disabled!")
    elif action == "factory":
        run_factory()
    elif action == "usb_enable":
        if enable_usb_interface():
            print_success("USB interface enabled with DHCP server on 192.168.100.1")
        else:
            sys.exit(1)
    elif action == "usb_disable":
        if disable_usb_interface():
            print_success("USB interface DHCP server disabled")
        else:
            sys.exit(1)
    elif action == "exit":
        sys.exit(0)
    else:
        print_error("Unrecognized option.")
        sys.exit(1)

def run_smart_station(ssid, password=""):
    """Smart station configuration that adapts to hardware"""
    interfaces = get_wifi_interfaces()
    
    # Configure station
    if password:
        call_voxl_wifi_legacy(["enable-station", ssid, password])
    else:
        call_voxl_wifi_legacy(["enable-station", ssid])
    
    print_success("Station mode configured!")

def run_smart_softap(ssid, password="1234567890", band="2ghz"):
    """Smart SoftAP configuration that adapts to hardware"""
    interfaces = get_wifi_interfaces()
    modes, _ = get_current_modes()
    
    # For single-mode hardware, disable other modes first
    if interfaces['concurrent']:
        if modes['station']:
            station_band = get_station_band()
            
            if station_band:
                if station_band != band:
                    print_error(f"Cannot enable {band} SoftAP while station is on {station_band}")
                    print_error("Both must be on the same band for concurrent operation.")
                    sys.exit(1)
                else:
                    print(f"Station is on {station_band} - compatible with {band} SoftAP")
    
    # Configure SoftAP
    call_voxl_wifi_legacy(["enable-softap", ssid, password, band])
    print_success(f"SoftAP mode configured on {band}!")

def main(args):
    # check for help
    if "--help" in args or "-h" in args:
        print_usage()
        return

    # default to wizard
    if len(args) == 0:
        run_wizard()
    
    elif args[0] == "getmode":
        run_getmode()
    
    elif args[0] == "station":
        if len(args) < 2:
            print_error("station command requires SSID")
            print("Usage: voxl-wifi station <ssid> [password]")
            sys.exit(1)
        ssid = args[1]
        password = args[2] if len(args) > 2 else ""
        run_smart_station(ssid, password)
    
    elif args[0] == "softap":
        if len(args) < 2:
            print_error("softap command requires SSID")
            print("Usage: voxl-wifi softap <ssid> [password]")
            sys.exit(1)
        ssid = args[1]
        password = args[2] if len(args) > 2 else "1234567890"
        run_smart_softap(ssid, password, "2ghz")
    
    elif args[0] == "softap5":
        if len(args) < 2:
            print_error("softap5 command requires SSID")
            print("Usage: voxl-wifi softap5 <ssid> [password]")
            sys.exit(1)
        ssid = args[1]
        password = args[2] if len(args) > 2 else "1234567890"
        run_smart_softap(ssid, password, "5ghz")
    
    elif args[0] == "disable-station":
        call_voxl_wifi_legacy(["disable-station"])
        print_success("Station mode disabled!")
    
    elif args[0] == "disable-softap":
        call_voxl_wifi_legacy(["disable-softap"])
        print_success("SoftAP mode disabled!")
    
    elif args[0] == "disable-all":
        call_voxl_wifi_legacy(["disable-all"])
        print_success("All WiFi modes disabled!")
    
    elif args[0] == "factory" or args[0] == "-f":
        run_factory()
    
    elif args[0] == "wizard":
        run_wizard()

    elif args[0] == "usb-enable":
        if not is_qcs6490():
            print_error("USB interface control is only available on QCS6490 platform")
            sys.exit(1)
        if not has_usb_interface():
            print_error("USB interface (usb0) not found")
            sys.exit(1)
        if enable_usb_interface():
            print_success("USB interface enabled with DHCP server on 192.168.100.1")
        else:
            sys.exit(1)

    elif args[0] == "usb-disable":
        if not is_qcs6490():
            print_error("USB interface control is only available on QCS6490 platform")
            sys.exit(1)
        if not has_usb_interface():
            print_error("USB interface (usb0) not found")
            sys.exit(1)
        if disable_usb_interface():
            print_success("USB interface DHCP server disabled")
        else:
            sys.exit(1)

    else:
        print_error(f"Unrecognized command '{args[0]}'")
        print("Run 'voxl-wifi --help' for usage information")
        sys.exit(1)

if __name__ ==  "__main__":
    try:
        main(sys.argv[1:])
        print()
        sys.exit(0)
    except KeyboardInterrupt:
        print_error("Caught keyboard interrupt, exiting...")
        sys.exit(1)
