#!/usr/bin/env python3

import os
import subprocess
import sys
import time
from collections import Counter

MAX_TRIES = 10


# Utility function for invoking iw
def run_iw(args: list[str]) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        ["iw", *args],
        check=True,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )


# Get first WiFi interface name
def get_wifi_if_name() -> str | None:
    # Call the iw dev command and capture its output
    try:
        result = run_iw(["dev"])
    except subprocess.CalledProcessError as e:
        print(f"Error: {e.stderr}")
        sys.exit(1)

    lines = result.stdout.splitlines()

    # Iterate over the lines and extract the interface name
    for line in lines:
        if line.lstrip().startswith("Interface "):
            interface_name = line.lstrip().split()[1]
            return interface_name

    return None


# Check if a regulatory domain is set already (not "country 00")
def reg_domain_is_unset() -> bool:
    try:
        result = run_iw(["reg", "get"])
    except subprocess.CalledProcessError as e:
        print(f"Error: {e.stderr}")
        sys.exit(2)

    return "country 00:" in result.stdout


# Get regulatory domains as country code list
def get_wifi_reg_domain(ifname: str) -> list[str] | None:
    print(f"Scan for WiFi networks with interface {ifname}")
    try:
        result = run_iw(["dev", ifname, "scan"])
    except subprocess.CalledProcessError as e:
        if "resource busy" in e.stderr:
            print("Device or resource busy")
            # try again later
            return None
        if "Network is down" in e.stderr:
            print("WiFi is disabled")
            sys.exit(0)
        else:
            print(f"Error executing '{e.cmd}': {e.stderr}")
            sys.exit(3)

    lines = result.stdout.splitlines()

    cc_list: list[str] = []
    # Iterate over the lines and extract the country name
    for line in lines:
        if line.lstrip().startswith("Country: "):
            cc = line.lstrip().split()[1]
            cc_list.append(cc)

    return cc_list


if_name = None

for _ in range(MAX_TRIES):
    if_name = get_wifi_if_name()
    if if_name is not None and os.path.isfile(f"/sys/class/net/{if_name}/address"):
        # Interface detected
        break

    # Wait before next attempt
    time.sleep(1)

# Check if an interface was detected
if if_name is None:
    print("No WiFi interface detected")
    sys.exit(4)

if not reg_domain_is_unset():
    print("Regulatory domain already set")
    sys.exit(0)

# Listen to other access points to get country codes
cc_list = None

for _ in range(MAX_TRIES):
    cc_list = get_wifi_reg_domain(if_name)
    if cc_list is not None:
        break

    # Wait before next attempt
    time.sleep(1)

# Exit if no access point was found with country code
if cc_list is None or len(cc_list) == 0:
    print("No access point with country code found")
    sys.exit(0)

# Sort and count list
counter_list = Counter(cc_list).most_common()

print("Detected domains:")
for cc, count in counter_list:
    print(f"  {cc}: {count} device(s)")

# Get country code for setting regulatory domain
selected_cc = ""
if len(counter_list) == 1:
    selected_cc = counter_list[0][0]

else:
    print(
        "Warning: found multiple regulatory domains.\n"
        "Maybe someone has a badly configured access point."
    )

    # Check if most common country code is leading at least by 2
    if counter_list[0][1] >= counter_list[1][1] + 2:
        selected_cc = counter_list[0][0]
    else:
        print(
            "Can not decide which country code is correct.\n"
            f"Candidates are {counter_list[0][0]} and {counter_list[1][0]}"
        )
        sys.exit(0)

print(f"Set country code {selected_cc}")
try:
    result = run_iw(["reg", "set", selected_cc])
except subprocess.CalledProcessError as e:
    print(f"Error: {e.stderr}")
    sys.exit(4)

# Success
sys.exit(0)
