#!/usr/bin/env python3
"""sd_wifi.py — Update WiFi config on a SystemMonitor SD card.

Detects the SD card by looking for the armbi_root label, mounts it to a
temp directory, reads/writes /etc/netplan/20-wifi.yaml, and unmounts on exit.

Usage:
  ./sd_wifi.py                         # interactive menu
  ./sd_wifi.py --list                  # show current configured networks
  ./sd_wifi.py --add SSID PASSWORD     # non-interactive add/update
"""

import json
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Dict, Optional, Tuple

import click
import yaml
from rich.console import Console
from rich.panel import Panel
from rich.prompt import Prompt
from rich.table import Table
from rich.text import Text

console = Console()

NETPLAN_PATH = "etc/netplan/20-wifi.yaml"
WPA_SUPPLICANT_PATH = "etc/wpa_supplicant/wpa_supplicant-wlan0.conf"
SD_LABEL = "armbi_root"


# ---------------------------------------------------------------------------
# SD card detection and mount
# ---------------------------------------------------------------------------

def find_sd_device() -> Optional[str]:
    """Return the block device path for the partition labelled armbi_root."""
    try:
        result = subprocess.run(
            ["lsblk", "-o", "NAME,LABEL", "--json"],
            capture_output=True, text=True, check=True,
        )
        data = json.loads(result.stdout)
    except (subprocess.CalledProcessError, json.JSONDecodeError, FileNotFoundError) as exc:
        console.print(f"[red]Error running lsblk:[/red] {exc}")
        return None

    def _search(devices):
        for dev in devices:
            label = dev.get("label") or ""
            if label.strip() == SD_LABEL:
                return f"/dev/{dev['name']}"
            children = dev.get("children") or []
            found = _search(children)
            if found:
                return found
        return None

    return _search(data.get("blockdevices", []))


def mount_sd(device: str, mountpoint: str) -> bool:
    """Mount device to mountpoint using sudo. Returns True on success."""
    result = subprocess.run(
        ["sudo", "mount", device, mountpoint],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        console.print(f"[red]Mount failed:[/red] {result.stderr.strip()}")
        return False
    return True


def unmount_sd(mountpoint: str) -> None:
    """Unmount mountpoint using sudo (best-effort, ignores errors)."""
    subprocess.run(
        ["sudo", "umount", mountpoint],
        capture_output=True, text=True,
    )


def detect_wifi_config_type(mountpoint: str) -> str:
    """Detect the WiFi config format used on the mounted SD card.

    Returns 'wpa_supplicant' if wpa_supplicant config is found,
    'netplan' if netplan dir exists, otherwise 'netplan' (default).
    """
    wpa_path = Path(mountpoint) / WPA_SUPPLICANT_PATH
    if wpa_path.exists():
        return "wpa_supplicant"

    netplan_dir = Path(mountpoint) / "etc/netplan"
    if netplan_dir.exists():
        return "netplan"

    return "netplan"


# ---------------------------------------------------------------------------
# Netplan read / write
# ---------------------------------------------------------------------------

NETPLAN_TEMPLATE: Dict = {
    "network": {
        "version": 2,
        "renderer": "networkd",
        "wifis": {
            "wlan0": {
                "dhcp4": True,
                "dhcp6": True,
                "regulatory-domain": "US",
                "access-points": {},
            }
        },
    }
}


def read_netplan(mountpoint: str) -> Dict:
    """Read and parse the netplan WiFi config from the mounted SD card.

    Returns the parsed config dict, or a fresh template if the file does not
    exist or cannot be parsed.
    """
    netplan_file = Path(mountpoint) / NETPLAN_PATH
    if not netplan_file.exists():
        console.print(f"[yellow]No netplan file found at {netplan_file}. Starting fresh.[/yellow]")
        return _deep_copy(NETPLAN_TEMPLATE)

    try:
        with open(netplan_file, "r") as fh:
            data = yaml.safe_load(fh)
        if not data or "network" not in data:
            console.print("[yellow]Netplan file appears empty or malformed. Starting fresh.[/yellow]")
            return _deep_copy(NETPLAN_TEMPLATE)
        return data
    except yaml.YAMLError as exc:
        console.print(f"[yellow]YAML parse error ({exc}). Starting fresh.[/yellow]")
        return _deep_copy(NETPLAN_TEMPLATE)


def write_netplan(mountpoint: str, config: Dict) -> None:
    """Write the netplan config to the SD card and set permissions to 600."""
    netplan_file = Path(mountpoint) / NETPLAN_PATH

    # Ensure the directory exists on the card
    netplan_file.parent.mkdir(parents=True, exist_ok=True)

    yaml_text = yaml.dump(config, default_flow_style=False, sort_keys=False,
                          allow_unicode=True)

    # Write via sudo tee to handle root-owned filesystem
    result = subprocess.run(
        ["sudo", "tee", str(netplan_file)],
        input=yaml_text, capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise RuntimeError(f"Failed to write netplan file: {result.stderr.strip()}")

    # Set permissions 600
    chmod_result = subprocess.run(
        ["sudo", "chmod", "600", str(netplan_file)],
        capture_output=True, text=True,
    )
    if chmod_result.returncode != 0:
        console.print(f"[yellow]Warning: chmod 600 failed: {chmod_result.stderr.strip()}[/yellow]")


def _deep_copy(obj):
    """Deep-copy a plain dict/list/scalar structure via JSON round-trip."""
    return json.loads(json.dumps(obj))


# ---------------------------------------------------------------------------
# WPA Supplicant read / write
# ---------------------------------------------------------------------------

def read_wpa_supplicant(mountpoint: str) -> Tuple[Dict[str, Dict], str]:
    """Read and parse wpa_supplicant config.

    Returns (aps_dict, country_code) where aps_dict is {ssid: {"password": "..."}}
    or {ssid: {}} for open networks. Country code extracted from 'country=XX' line.
    """
    wpa_path = Path(mountpoint) / WPA_SUPPLICANT_PATH
    country = "US"
    aps = {}

    if not wpa_path.exists():
        return aps, country

    try:
        with open(wpa_path, "r") as fh:
            content = fh.read()
    except Exception:
        return aps, country

    # Parse country line
    for line in content.split("\n"):
        line = line.strip()
        if line.startswith("country="):
            country = line.split("=", 1)[1].strip()
            break

    # Parse network blocks
    in_network = False
    current_ssid = None
    current_psk = None

    for line in content.split("\n"):
        line = line.strip()

        if line == "network={":
            in_network = True
            current_ssid = None
            current_psk = None
            continue

        if in_network:
            if line == "}":
                if current_ssid:
                    aps[current_ssid] = {"password": current_psk} if current_psk else {}
                in_network = False
                current_ssid = None
                current_psk = None
                continue

            if line.startswith("ssid="):
                # Strip quotes
                ssid_val = line.split("=", 1)[1].strip()
                if ssid_val.startswith('"') and ssid_val.endswith('"'):
                    ssid_val = ssid_val[1:-1]
                current_ssid = ssid_val
            elif line.startswith("psk="):
                # Strip quotes
                psk_val = line.split("=", 1)[1].strip()
                if psk_val.startswith('"') and psk_val.endswith('"'):
                    psk_val = psk_val[1:-1]
                current_psk = psk_val

    return aps, country


def write_wpa_supplicant(mountpoint: str, aps: Dict[str, Dict], country: str) -> None:
    """Write wpa_supplicant config file.

    Args:
        mountpoint: Root mountpoint of the SD card
        aps: {ssid: {"password": "..."}} or {ssid: {}} for open networks
        country: 2-letter country code
    """
    wpa_path = Path(mountpoint) / WPA_SUPPLICANT_PATH

    # Ensure the directory exists
    wpa_path.parent.mkdir(parents=True, exist_ok=True)

    # Build the config content
    lines = [
        "ctrl_interface=DIR=/var/run/wpa_supplicant GROUP=netdev",
        "update_config=1",
        f"country={country.upper()}",
        "",
    ]

    for ssid, details in aps.items():
        password = None
        if isinstance(details, dict):
            password = details.get("password")

        lines.append("network={")
        lines.append(f'    ssid="{ssid}"')

        if password:
            lines.append(f'    psk="{password}"')
            lines.append("    key_mgmt=WPA-PSK")
        else:
            lines.append("    key_mgmt=NONE")

        lines.append("    scan_ssid=1")
        lines.append("}")
        lines.append("")

    config_text = "\n".join(lines)

    # Write via sudo tee
    result = subprocess.run(
        ["sudo", "tee", str(wpa_path)],
        input=config_text, capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise RuntimeError(f"Failed to write wpa_supplicant file: {result.stderr.strip()}")

    # Set permissions 600
    chmod_result = subprocess.run(
        ["sudo", "chmod", "600", str(wpa_path)],
        capture_output=True, text=True,
    )
    if chmod_result.returncode != 0:
        console.print(f"[yellow]Warning: chmod 600 failed: {chmod_result.stderr.strip()}[/yellow]")


# ---------------------------------------------------------------------------
# Config accessors
# ---------------------------------------------------------------------------

def get_access_points(config: Dict) -> Dict[str, str]:
    """Return {ssid: password} mapping from the config."""
    try:
        return dict(
            config["network"]["wifis"]["wlan0"].get("access-points", {}) or {}
        )
    except (KeyError, TypeError):
        return {}


def get_country(config: Dict) -> str:
    """Return the regulatory-domain value."""
    try:
        return config["network"]["wifis"]["wlan0"].get("regulatory-domain", "US")
    except (KeyError, TypeError):
        return "US"


def set_access_points(config: Dict, aps: Dict[str, str]) -> None:
    """Overwrite the access-points section."""
    config["network"]["wifis"]["wlan0"]["access-points"] = aps


def set_country(config: Dict, country: str) -> None:
    """Set the regulatory-domain value."""
    config["network"]["wifis"]["wlan0"]["regulatory-domain"] = country.upper()


# ---------------------------------------------------------------------------
# Display helpers
# ---------------------------------------------------------------------------

def print_header() -> None:
    console.print()
    console.print(Panel(
        "[bold cyan]SystemMonitor[/bold cyan] — SD Card WiFi Configurator\n"
        "Target label: [yellow]armbi_root[/yellow]  |  "
        "Netplan: [dim]" + NETPLAN_PATH + "[/dim]",
        title="sd_wifi",
        border_style="cyan",
    ))
    console.print()


def print_networks(config: Dict) -> None:
    """Print a Rich table of configured access-points."""
    aps = get_access_points(config)
    country = get_country(config)

    console.print()
    if not aps:
        console.print("  [dim]No WiFi networks configured.[/dim]")
    else:
        table = Table(title=f"Configured Networks  (country: {country})",
                      show_lines=True, border_style="cyan")
        table.add_column("#", justify="right", min_width=3, style="dim")
        table.add_column("SSID", min_width=24, style="bold")
        table.add_column("Password", min_width=20)

        for idx, (ssid, details) in enumerate(aps.items(), start=1):
            if isinstance(details, dict):
                pwd = details.get("password", "")
            else:
                pwd = str(details) if details else ""
            table.add_row(str(idx), ssid, pwd or "[dim](none)[/dim]")

        console.print(table)
    console.print()


# ---------------------------------------------------------------------------
# Interactive menu actions
# ---------------------------------------------------------------------------

def action_add_network(config: Dict) -> bool:
    """Prompt for SSID + password and add (or update) the entry. Returns True if changed."""
    console.print()
    ssid = Prompt.ask("  [cyan]SSID[/cyan]").strip()
    if not ssid:
        console.print("  [yellow]Cancelled — empty SSID.[/yellow]")
        return False

    password = Prompt.ask("  [cyan]Password[/cyan] (leave blank for open network)").strip()

    aps = get_access_points(config)
    existed = ssid in aps
    aps[ssid] = {"password": password} if password else {}
    set_access_points(config, aps)

    verb = "Updated" if existed else "Added"
    console.print(f"  [green]{verb}[/green] [bold]{ssid}[/bold]")
    return True


def action_remove_network(config: Dict) -> bool:
    """Prompt the user to choose a network to remove. Returns True if changed."""
    aps = get_access_points(config)
    if not aps:
        console.print("  [yellow]No networks configured.[/yellow]")
        return False

    ssid_list = list(aps.keys())
    console.print()
    for idx, ssid in enumerate(ssid_list, start=1):
        console.print(f"    [bold]{idx}.[/bold] {ssid}")
    console.print()

    choice = Prompt.ask("  Remove # (or [dim]Enter[/dim] to cancel)").strip()
    if not choice:
        return False

    try:
        n = int(choice)
        if not (1 <= n <= len(ssid_list)):
            raise ValueError
    except ValueError:
        console.print("  [yellow]Invalid selection.[/yellow]")
        return False

    target = ssid_list[n - 1]
    del aps[target]
    set_access_points(config, aps)
    console.print(f"  [green]Removed[/green] [bold]{target}[/bold]")
    return True


def action_update_password(config: Dict) -> bool:
    """Update the password for an existing network. Returns True if changed."""
    aps = get_access_points(config)
    if not aps:
        console.print("  [yellow]No networks configured.[/yellow]")
        return False

    ssid_list = list(aps.keys())
    console.print()
    for idx, ssid in enumerate(ssid_list, start=1):
        console.print(f"    [bold]{idx}.[/bold] {ssid}")
    console.print()

    choice = Prompt.ask("  Update password for # (or [dim]Enter[/dim] to cancel)").strip()
    if not choice:
        return False

    try:
        n = int(choice)
        if not (1 <= n <= len(ssid_list)):
            raise ValueError
    except ValueError:
        console.print("  [yellow]Invalid selection.[/yellow]")
        return False

    target = ssid_list[n - 1]
    password = Prompt.ask(f"  New password for [bold]{target}[/bold] "
                          "(leave blank for open network)").strip()

    existing = aps[target]
    if isinstance(existing, dict):
        entry = dict(existing)
    else:
        entry = {}

    if password:
        entry["password"] = password
    else:
        entry.pop("password", None)

    aps[target] = entry if entry else {}
    set_access_points(config, aps)
    console.print(f"  [green]Password updated[/green] for [bold]{target}[/bold]")
    return True


def action_change_country(config: Dict) -> bool:
    """Prompt for a new regulatory-domain country code. Returns True if changed."""
    current = get_country(config)
    console.print()
    new_country = Prompt.ask(
        f"  Country code (current: [yellow]{current}[/yellow])"
    ).strip().upper()

    if not new_country:
        console.print("  [yellow]Cancelled.[/yellow]")
        return False

    if len(new_country) != 2 or not new_country.isalpha():
        console.print("  [yellow]Invalid country code — must be 2 letters (e.g. US, GB, DE).[/yellow]")
        return False

    set_country(config, new_country)
    console.print(f"  [green]Country set to[/green] [bold]{new_country}[/bold]")
    return True


# ---------------------------------------------------------------------------
# Core flow
# ---------------------------------------------------------------------------

def run_interactive(device: str, mountpoint: str) -> None:
    """Mount SD card, run interactive menu, write changes, unmount."""
    console.print(f"  Mounting [cyan]{device}[/cyan] → [dim]{mountpoint}[/dim] ...", end=" ")

    if not mount_sd(device, mountpoint):
        return

    console.print("[green]OK[/green]")

    try:
        config_type = detect_wifi_config_type(mountpoint)

        if config_type == "wpa_supplicant":
            # WPA Supplicant mode: work with aps dict and country string directly
            aps, country = read_wpa_supplicant(mountpoint)
            dirty = False

            while True:
                # Print networks table for wpa_supplicant
                console.print()
                if not aps:
                    console.print("  [dim]No WiFi networks configured.[/dim]")
                else:
                    table = Table(title=f"Configured Networks  (country: {country})",
                                  show_lines=True, border_style="cyan")
                    table.add_column("#", justify="right", min_width=3, style="dim")
                    table.add_column("SSID", min_width=24, style="bold")
                    table.add_column("Password", min_width=20)

                    for idx, (ssid, details) in enumerate(aps.items(), start=1):
                        pwd = details.get("password", "") if isinstance(details, dict) else ""
                        table.add_row(str(idx), ssid, pwd or "[dim](none)[/dim]")

                    console.print(table)
                console.print()
                console.print(f"  Country: [yellow]{country}[/yellow]")
                console.print()
                console.print("  [bold]1.[/bold] Add / update network")
                console.print("  [bold]2.[/bold] Remove network")
                console.print("  [bold]3.[/bold] Update password")
                console.print("  [bold]4.[/bold] Change country code")
                console.print("  [bold]5.[/bold] Done (write & exit)")
                console.print("  [bold]0.[/bold] Quit without saving")
                console.print()

                choice = Prompt.ask("  [cyan]Select[/cyan]").strip()

                if choice == "1":
                    # Add network
                    console.print()
                    ssid = Prompt.ask("  [cyan]SSID[/cyan]").strip()
                    if ssid:
                        password = Prompt.ask("  [cyan]Password[/cyan] (leave blank for open network)").strip()
                        aps[ssid] = {"password": password} if password else {}
                        verb = "Updated" if ssid in aps else "Added"
                        console.print(f"  [green]Added[/green] [bold]{ssid}[/bold]")
                        dirty = True
                    else:
                        console.print("  [yellow]Cancelled — empty SSID.[/yellow]")
                elif choice == "2":
                    # Remove network
                    if not aps:
                        console.print("  [yellow]No networks configured.[/yellow]")
                    else:
                        ssid_list = list(aps.keys())
                        console.print()
                        for idx, ssid in enumerate(ssid_list, start=1):
                            console.print(f"    [bold]{idx}.[/bold] {ssid}")
                        console.print()
                        choice_rem = Prompt.ask("  Remove # (or [dim]Enter[/dim] to cancel)").strip()
                        if choice_rem:
                            try:
                                n = int(choice_rem)
                                if 1 <= n <= len(ssid_list):
                                    target = ssid_list[n - 1]
                                    del aps[target]
                                    console.print(f"  [green]Removed[/green] [bold]{target}[/bold]")
                                    dirty = True
                                else:
                                    console.print("  [yellow]Invalid selection.[/yellow]")
                            except ValueError:
                                console.print("  [yellow]Invalid selection.[/yellow]")
                elif choice == "3":
                    # Update password
                    if not aps:
                        console.print("  [yellow]No networks configured.[/yellow]")
                    else:
                        ssid_list = list(aps.keys())
                        console.print()
                        for idx, ssid in enumerate(ssid_list, start=1):
                            console.print(f"    [bold]{idx}.[/bold] {ssid}")
                        console.print()
                        choice_upd = Prompt.ask("  Update password for # (or [dim]Enter[/dim] to cancel)").strip()
                        if choice_upd:
                            try:
                                n = int(choice_upd)
                                if 1 <= n <= len(ssid_list):
                                    target = ssid_list[n - 1]
                                    password = Prompt.ask(f"  New password for [bold]{target}[/bold] "
                                                          "(leave blank for open network)").strip()
                                    aps[target] = {"password": password} if password else {}
                                    console.print(f"  [green]Password updated[/green] for [bold]{target}[/bold]")
                                    dirty = True
                                else:
                                    console.print("  [yellow]Invalid selection.[/yellow]")
                            except ValueError:
                                console.print("  [yellow]Invalid selection.[/yellow]")
                elif choice == "4":
                    # Change country
                    console.print()
                    new_country = Prompt.ask(
                        f"  Country code (current: [yellow]{country}[/yellow])"
                    ).strip().upper()
                    if new_country:
                        if len(new_country) != 2 or not new_country.isalpha():
                            console.print("  [yellow]Invalid country code — must be 2 letters (e.g. US, GB, DE).[/yellow]")
                        else:
                            country = new_country
                            console.print(f"  [green]Country set to[/green] [bold]{new_country}[/bold]")
                            dirty = True
                    else:
                        console.print("  [yellow]Cancelled.[/yellow]")
                elif choice == "5":
                    break
                elif choice == "0":
                    console.print("\n  [yellow]Discarding changes.[/yellow]")
                    dirty = False
                    return
                else:
                    console.print("  [yellow]Invalid option.[/yellow]")

            if dirty:
                console.print()
                console.print("  Writing config ... ", end="")
                try:
                    write_wpa_supplicant(mountpoint, aps, country)
                    console.print("[green]OK[/green]")
                    console.print(f"  [green]Saved[/green] → [dim]{WPA_SUPPLICANT_PATH}[/dim] (chmod 600)")
                except RuntimeError as exc:
                    console.print(f"[red]FAILED[/red]\n  {exc}")
            else:
                console.print("  [dim]No changes to write.[/dim]")

        else:
            # Netplan mode: use existing config dict-based approach
            config = read_netplan(mountpoint)
            dirty = False

            while True:
                print_networks(config)
                country = get_country(config)
                console.print(f"  Country: [yellow]{country}[/yellow]")
                console.print()
                console.print("  [bold]1.[/bold] Add / update network")
                console.print("  [bold]2.[/bold] Remove network")
                console.print("  [bold]3.[/bold] Update password")
                console.print("  [bold]4.[/bold] Change country code")
                console.print("  [bold]5.[/bold] Done (write & exit)")
                console.print("  [bold]0.[/bold] Quit without saving")
                console.print()

                choice = Prompt.ask("  [cyan]Select[/cyan]").strip()

                if choice == "1":
                    if action_add_network(config):
                        dirty = True
                elif choice == "2":
                    if action_remove_network(config):
                        dirty = True
                elif choice == "3":
                    if action_update_password(config):
                        dirty = True
                elif choice == "4":
                    if action_change_country(config):
                        dirty = True
                elif choice == "5":
                    break
                elif choice == "0":
                    console.print("\n  [yellow]Discarding changes.[/yellow]")
                    dirty = False
                    return
                else:
                    console.print("  [yellow]Invalid option.[/yellow]")

            if dirty:
                console.print()
                console.print("  Writing config ... ", end="")
                try:
                    write_netplan(mountpoint, config)
                    console.print("[green]OK[/green]")
                    console.print(f"  [green]Saved[/green] → [dim]{NETPLAN_PATH}[/dim] (chmod 600)")
                except RuntimeError as exc:
                    console.print(f"[red]FAILED[/red]\n  {exc}")
            else:
                console.print("  [dim]No changes to write.[/dim]")

    finally:
        console.print()
        console.print(f"  Unmounting [dim]{mountpoint}[/dim] ... ", end="")
        unmount_sd(mountpoint)
        console.print("[green]OK[/green]")


def run_list(device: str, mountpoint: str) -> None:
    """Mount SD card, list networks, unmount."""
    if not mount_sd(device, mountpoint):
        return

    try:
        config_type = detect_wifi_config_type(mountpoint)

        if config_type == "wpa_supplicant":
            aps, country = read_wpa_supplicant(mountpoint)
            console.print()
            if not aps:
                console.print("  [dim]No WiFi networks configured.[/dim]")
            else:
                table = Table(title=f"Configured Networks  (country: {country})",
                              show_lines=True, border_style="cyan")
                table.add_column("#", justify="right", min_width=3, style="dim")
                table.add_column("SSID", min_width=24, style="bold")
                table.add_column("Password", min_width=20)

                for idx, (ssid, details) in enumerate(aps.items(), start=1):
                    pwd = details.get("password", "") if isinstance(details, dict) else ""
                    table.add_row(str(idx), ssid, pwd or "[dim](none)[/dim]")

                console.print(table)
            console.print()
            console.print(f"  Regulatory domain: [yellow]{country}[/yellow]")
        else:
            config = read_netplan(mountpoint)
            print_networks(config)
            country = get_country(config)
            console.print(f"  Regulatory domain: [yellow]{country}[/yellow]")
    finally:
        unmount_sd(mountpoint)


def run_add(device: str, mountpoint: str, ssid: str, password: str) -> None:
    """Mount SD card, add/update a network non-interactively, unmount."""
    if not mount_sd(device, mountpoint):
        return

    try:
        config_type = detect_wifi_config_type(mountpoint)

        if config_type == "wpa_supplicant":
            aps, country = read_wpa_supplicant(mountpoint)
            existed = ssid in aps
            aps[ssid] = {"password": password} if password else {}
            write_wpa_supplicant(mountpoint, aps, country)
            verb = "Updated" if existed else "Added"
            console.print(f"  [green]{verb}[/green] [bold]{ssid}[/bold]  (chmod 600 applied)")
        else:
            config = read_netplan(mountpoint)
            aps = get_access_points(config)
            existed = ssid in aps
            aps[ssid] = {"password": password} if password else {}
            set_access_points(config, aps)

            write_netplan(mountpoint, config)

            verb = "Updated" if existed else "Added"
            console.print(f"  [green]{verb}[/green] [bold]{ssid}[/bold]  (chmod 600 applied)")
    except RuntimeError as exc:
        console.print(f"  [red]Error:[/red] {exc}")
    finally:
        unmount_sd(mountpoint)


# ---------------------------------------------------------------------------
# Click CLI
# ---------------------------------------------------------------------------

@click.command(context_settings={"help_option_names": ["-h", "--help"]})
@click.option("--list", "do_list", is_flag=True, default=False,
              help="Show current configured networks and exit.")
@click.option("--add", nargs=2, metavar="SSID PASSWORD",
              default=(None, None),
              help="Non-interactively add or update a network.")
def main(do_list: bool, add: Tuple[Optional[str], Optional[str]]) -> None:
    """Update WiFi config on a SystemMonitor SD card (armbi_root label)."""
    print_header()

    # Locate the SD card
    console.print(f"  Scanning for SD card with label [yellow]{SD_LABEL}[/yellow] ...", end=" ")
    device = find_sd_device()

    if not device:
        console.print("[red]NOT FOUND[/red]")
        console.print()
        console.print(f"  [red]No block device with label '{SD_LABEL}' detected.[/red]")
        console.print("  Insert the SystemMonitor SD card and try again.")
        sys.exit(1)

    console.print(f"[green]{device}[/green]")

    # Create temp mountpoint
    mountpoint = tempfile.mkdtemp(prefix="bb_sd_")

    try:
        ssid, password = add
        if ssid is not None:
            # --add mode
            run_add(device, mountpoint, ssid, password or "")
        elif do_list:
            # --list mode
            run_list(device, mountpoint)
        else:
            # Interactive mode
            run_interactive(device, mountpoint)
    finally:
        # Always clean up the temp directory
        try:
            os.rmdir(mountpoint)
        except OSError:
            pass


if __name__ == "__main__":
    main()
