#!/usr/bin/env python3
"""sd_wifi.py — Update WiFi config on a BigBrother 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"
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,
    )


# ---------------------------------------------------------------------------
# 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))


# ---------------------------------------------------------------------------
# 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]BigBrother[/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 = 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 = 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 = 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 BigBrother 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 BigBrother 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()