pi-multifx-pedal/src/system/network.py

"""WiFi network management module for Pi Multi-FX Pedal.

Uses nmcli (NetworkManager) by default — works on RPi OS and most
desktop Linux distros. Falls back to iwlist + manual wpa_supplicant
for minimized DietPi installs without NetworkManager.

All public methods return dicts suitable for JSON serialization.
"""

from __future__ import annotations

import json
import logging
import os
import re
import subprocess
import tempfile
import time
from pathlib import Path
from typing import Any, Optional

logger = logging.getLogger(__name__)

# ── Constants ─────────────────────────────────────────────────────────────────

HOTSPOT_SCRIPT = Path(__file__).resolve().parent.parent.parent / "scripts" / "setup-wifi-ap.sh"
WPA_SUPPLICANT_CONF = Path("/etc/wpa_supplicant/wpa_supplicant.conf")

# ── Capability detection ──────────────────────────────────────────────────────


def _has_nmcli() -> bool:
    """Check if NetworkManager (nmcli) is available on this system."""
    try:
        subprocess.run(
            ["which", "nmcli"],
            capture_output=True, timeout=5, check=True,
        )
        return True
    except (subprocess.CalledProcessError, FileNotFoundError, OSError):
        return False


def _run(cmd: list[str], timeout: int = 15, check: bool = True) -> subprocess.CompletedProcess:
    """Run a subprocess and return the result."""
    logger.debug("Running: %s", " ".join(cmd))
    return subprocess.run(
        cmd, capture_output=True, text=True, timeout=timeout, check=check,
    )


# ── nmcli helpers ─────────────────────────────────────────────────────────────


def _nmcli_scan() -> list[dict[str, Any]]:
    """Scan WiFi networks using nmcli.

    Returns a list of networks with SSID, signal strength (%), and encryption type.
    """
    result = _run(
        ["nmcli", "-t", "-f", "SSID,SIGNAL,SECURITY", "device", "wifi", "list"],
        timeout=20,
        check=False,
    )
    networks: list[dict[str, Any]] = []
    seen: set[str] = set()
    for line in result.stdout.strip().split("\n"):
        if not line.strip() or ":" not in line:
            continue
        parts = line.split(":", 2)
        if len(parts) < 3:
            continue
        ssid, signal_str, security = parts
        ssid = ssid.strip()
        if not ssid or ssid in seen:
            continue
        seen.add(ssid)
        try:
            signal = int(signal_str)
        except ValueError:
            signal = 0
        networks.append({
            "ssid": ssid,
            "signal": signal,
            "encryption": security.strip() if security.strip() and security.strip() != "" else "Open",
        })
    return sorted(networks, key=lambda n: n["signal"], reverse=True)


def _nmcli_status() -> dict[str, Any]:
    """Get current WiFi connection status via nmcli.

    Returns SSID, IP address, signal strength, and device name.
    """
    result = _run(
        ["nmcli", "-t", "-f", "DEVICE,TYPE,STATE,CONNECTION", "device", "status"],
        timeout=10, check=False,
    )
    wifi_device = None
    for line in result.stdout.strip().split("\n"):
        parts = line.split(":")
        if len(parts) >= 3 and parts[1] == "wifi":
            wifi_device = parts[0]
            break

    if not wifi_device:
        return {"connected": False, "ssid": None, "ip": None, "signal": 0}

    # Get connection details
    detail = _run(
        ["nmcli", "-t", "-f", "GENERAL.CONNECTION,IP4.ADDRESS", "device", "show", wifi_device],
        timeout=10, check=False,
    )
    ssid = None
    ip = None
    for line in detail.stdout.strip().split("\n"):
        if line.startswith("GENERAL.CONNECTION:"):
            val = line.split(":", 1)[1].strip()
            if val and val != "--":
                ssid = val
        elif line.startswith("IP4.ADDRESS:"):
            val = line.split(":", 1)[1].strip()
            if val and val != "--":
                ip = val.split("/")[0]

    if not ssid:
        return {"connected": False, "ssid": None, "ip": None, "signal": 0}

    # Get signal strength
    signal = 0
    sig_result = _run(
        ["nmcli", "-t", "-f", "SSID,SIGNAL", "device", "wifi", "list"],
        timeout=10, check=False,
    )
    for line in sig_result.stdout.strip().split("\n"):
        if ":" in line:
            s, sig = line.split(":", 1)
            if s.strip() == ssid:
                try:
                    signal = int(sig)
                except ValueError:
                    pass
                break

    return {
        "connected": True,
        "ssid": ssid,
        "ip": ip,
        "signal": signal,
        "device": wifi_device,
    }


def _nmcli_saved() -> list[dict[str, Any]]:
    """List saved/known WiFi connections."""
    result = _run(
        ["nmcli", "-t", "-f", "NAME,UUID,TYPE", "connection", "show"],
        timeout=10, check=False,
    )
    saved: list[dict[str, Any]] = []
    for line in result.stdout.strip().split("\n"):
        parts = line.split(":")
        if len(parts) >= 3 and "wireless" in parts[2]:
            saved.append({
                "ssid": parts[0],
                "uuid": parts[1],
            })
    return saved


def _nmcli_connect(ssid: str, password: str) -> bool:
    """Connect to a WiFi network.

    If the network is already saved, just activate it.
    Otherwise create a new connection profile.
    """
    # Check if already saved
    saved = _nmcli_saved()
    for net in saved:
        if net["ssid"] == ssid:
            _run(
                ["nmcli", "connection", "up", net["uuid"]],
                timeout=30, check=False,
            )
            logger.info("Activated saved network: %s", ssid)
            return True

    # Connect with password
    result = _run(
        ["nmcli", "device", "wifi", "connect", ssid, "password", password],
        timeout=30, check=False,
    )
    ok = result.returncode == 0
    if ok:
        logger.info("Connected to new network: %s", ssid)
    else:
        logger.warning("Failed to connect to %s: %s", ssid, result.stderr.strip())
    return ok


def _nmcli_disconnect() -> bool:
    """Disconnect from the current WiFi network."""
    status = _nmcli_status()
    if not status.get("connected"):
        return True  # already disconnected
    device = status.get("device", "wlan0")
    result = _run(
        ["nmcli", "device", "disconnect", device],
        timeout=15, check=False,
    )
    return result.returncode == 0


def _nmcli_forget(ssid: str) -> bool:
    """Delete a saved WiFi network."""
    saved = _nmcli_saved()
    for net in saved:
        if net["ssid"] == ssid:
            _run(
                ["nmcli", "connection", "delete", net["uuid"]],
                timeout=15, check=False,
            )
            logger.info("Forgot network: %s", ssid)
            return True
    return False


# ── iwlist fallback (minimized DietPi) ────────────────────────────────────────


def _iwlist_scan() -> list[dict[str, Any]]:
    """Fallback WiFi scan using iwlist (no NetworkManager)."""
    result = _run(
        ["sudo", "iwlist", "wlan0", "scan"],
        timeout=30, check=False,
    )
    networks: list[dict[str, Any]] = []
    current: dict[str, Any] = {}
    for line in result.stdout.split("\n"):
        line = line.strip()
        if "ESSID:" in line:
            essid = line.split("ESSID:")[1].strip().strip('"')
            if essid and essid != "":
                current["ssid"] = essid
        elif "Signal level=" in line:
            match = re.search(r"Signal level=(-?\d+)", line)
            if match:
                dbm = int(match.group(1))
                # Convert dBm to percentage (0% = -100dBm, 100% = -30dBm)
                signal = max(0, min(100, int((dbm + 100) * 100 / 70)))
                current["signal"] = signal
        elif "Encryption key:" in line:
            val = line.split(":")[1].strip()
            current["encryption"] = "WPA/WPA2" if val == "on" else "Open"
        elif "IE:" in line and "WPA" in line.upper():
            if current.get("encryption") in (None, "Open"):
                current["encryption"] = "WPA/WPA2"
        if current.get("ssid") and "signal" in current and "encryption" in current:
            networks.append(dict(current))
            current = {}

    seen: set[str] = set()
    unique: list[dict[str, Any]] = []
    for n in networks:
        if n["ssid"] not in seen:
            seen.add(n["ssid"])
            unique.append(n)
    return sorted(unique, key=lambda n: n.get("signal", 0), reverse=True)


def _iwlist_status() -> dict[str, Any]:
    """Fallback WiFi status using iwconfig."""
    result = _run(
        ["iwconfig", "wlan0"],
        timeout=10, check=False,
    )
    stdout = result.stdout
    ssid = None
    signal = 0
    match_ssid = re.search(r'ESSID:"([^"]+)"', stdout)
    if match_ssid:
        ssid = match_ssid.group(1)
    match_signal = re.search(r"Signal level=(-?\d+)", stdout)
    if match_signal:
        dbm = int(match_signal.group(1))
        signal = max(0, min(100, int((dbm + 100) * 100 / 70)))

    # Get IP
    ip = None
    ip_result = _run(
        ["ip", "-4", "addr", "show", "wlan0"],
        timeout=5, check=False,
    )
    ip_match = re.search(r"inet (\d+\.\d+\.\d+\.\d+)", ip_result.stdout)
    if ip_match:
        ip = ip_match.group(1)

    if not ssid or ssid == "off/any":
        return {"connected": False, "ssid": None, "ip": None, "signal": 0}
    return {"connected": True, "ssid": ssid, "ip": ip, "signal": signal, "device": "wlan0"}


def _wpa_connect(ssid: str, password: str) -> bool:
    """Fallback: add network to wpa_supplicant.conf and restart."""
    try:
        # Generate psk
        psk_result = _run(
            ["wpa_passphrase", ssid, password],
            timeout=10, check=True,
        )
        network_block = psk_result.stdout.strip()

        # Append to wpa_supplicant.conf
        if WPA_SUPPLICANT_CONF.exists():
            current = WPA_SUPPLICANT_CONF.read_text()
            if ssid in current:
                # Remove old config for this SSID and re-add
                lines = current.split("\n")
                new_lines = []
                skip = False
                for line in lines:
                    if f'ssid="{ssid}"' in line:
                        skip = True
                        continue
                    if skip and "}" in line:
                        skip = False
                        continue
                    if not skip:
                        new_lines.append(line)
                current = "\n".join(new_lines)

            updated = current.strip() + "\n\n" + network_block + "\n"
        else:
            updated = network_block + "\n"

        WPA_SUPPLICANT_CONF.parent.mkdir(parents=True, exist_ok=True)
        # Write via temp to avoid partial writes
        with tempfile.NamedTemporaryFile(
            mode="w", dir=str(WPA_SUPPLICANT_CONF.parent),
            prefix="wpa_temp_", delete=False,
        ) as f:
            f.write(updated)
            tmp_path = f.name
        _run(["sudo", "cp", tmp_path, str(WPA_SUPPLICANT_CONF)], timeout=10)
        os.unlink(tmp_path)

        # Restart networking
        _run(["sudo", "wpa_cli", "-i", "wlan0", "reconfigure"], timeout=10, check=False)
        time.sleep(3)
        return True
    except Exception as e:
        logger.error("Failed to connect with wpa_supplicant: %s", e)
        return False


# ── Hotspot control ───────────────────────────────────────────────────────────


def _hotspot_status() -> dict[str, Any]:
    """Check if the WiFi hotspot is currently active."""
    # Check if hostapd is running
    result = _run(
        ["systemctl", "is-active", "hostapd"],
        timeout=10, check=False,
    )
    active = result.stdout.strip() == "active"

    if not active:
        return {"active": False, "ssid": None, "clients": 0, "ip": None}

    # Read SSID from config
    ssid = None
    hostapd_conf = Path("/etc/hostapd/hostapd.conf")
    if hostapd_conf.exists():
        for line in hostapd_conf.read_text().split("\n"):
            if line.startswith("ssid="):
                ssid = line.split("=", 1)[1].strip()
                break

    # Count connected clients via iw
    clients = 0
    iw_result = _run(
        ["iw", "dev", "wlan0", "station", "dump"],
        timeout=10, check=False,
    )
    if iw_result.returncode == 0:
        clients = len([l for l in iw_result.stdout.split("\n") if "Station" in l])

    return {
        "active": True,
        "ssid": ssid or "Pi-Pedal",
        "clients": clients,
        "ip": "192.168.4.1",
    }


def _hotspot_enable(ssid: str = "Pi-Pedal", password: str = "pedal1234") -> bool:
    """Enable WiFi hotspot mode by calling the setup script."""
    result = _run(
        ["sudo", "bash", str(HOTSPOT_SCRIPT), "--ssid", ssid, "--psk", password],
        timeout=60, check=False,
    )
    ok = result.returncode == 0
    if ok:
        logger.info("Hotspot enabled: SSID=%s", ssid)
    else:
        logger.warning("Failed to enable hotspot: %s", result.stderr.strip())
    return ok


def _hotspot_disable() -> bool:
    """Disable WiFi hotspot, return to client mode."""
    result = _run(
        ["sudo", "bash", str(HOTSPOT_SCRIPT), "--disable"],
        timeout=60, check=False,
    )
    ok = result.returncode == 0
    if ok:
        logger.info("Hotspot disabled")
    else:
        logger.warning("Failed to disable hotspot: %s", result.stderr.strip())
    return ok


# ── Public API ─────────────────────────────────────────────────────────────────


def wifi_scan() -> list[dict[str, Any]]:
    """Scan for available WiFi networks.

    Returns a list of dicts with keys: ssid, signal (%), encryption.
    """
    if _has_nmcli():
        try:
            return _nmcli_scan()
        except Exception as e:
            logger.warning("nmcli scan failed, falling back to iwlist: %s", e)
    try:
        return _iwlist_scan()
    except Exception as e:
        logger.error("All WiFi scan methods failed: %s", str(e))
        return []


def wifi_status() -> dict[str, Any]:
    """Get current WiFi client connection status.

    Returns dict with: connected (bool), ssid (str|None), ip (str|None),
    signal (int), device (str|None).
    """
    if _has_nmcli():
        try:
            return _nmcli_status()
        except Exception as e:
            logger.warning("nmcli status failed, falling back: %s", e)
    try:
        return _iwlist_status()
    except Exception as e:
        logger.error("All WiFi status methods failed: %s", str(e))
        return {"connected": False, "ssid": None, "ip": None, "signal": 0}


def wifi_saved_networks() -> list[dict[str, Any]]:
    """List saved/known WiFi networks.

    Returns list of dicts with keys: ssid, uuid (nmcli) or just ssid.
    """
    if _has_nmcli():
        try:
            return _nmcli_saved()
        except Exception as e:
            logger.warning("nmcli saved networks failed: %s", e)
    # Fallback: parse wpa_supplicant.conf
    try:
        if WPA_SUPPLICANT_CONF.exists():
            networks: list[dict[str, Any]] = []
            for line in WPA_SUPPLICANT_CONF.read_text().split("\n"):
                match = re.search(r'ssid="([^"]+)"', line)
                if match:
                    networks.append({"ssid": match.group(1)})
            return networks
    except Exception as e:
        logger.warning("Failed to read wpa_supplicant.conf: %s", e)
    return []


def wifi_connect(ssid: str, password: str) -> dict[str, Any]:
    """Connect to a WiFi network.

    Args:
        ssid: Network SSID.
        password: Network password.

    Returns:
        dict with ok (bool) and error (str|None) keys.
    """
    if not ssid:
        return {"ok": False, "error": "SSID is required"}
    if not password:
        return {"ok": False, "error": "Password is required"}

    if _has_nmcli():
        try:
            ok = _nmcli_connect(ssid, password)
            return {"ok": ok, "error": None if ok else "Failed to connect via NetworkManager"}
        except Exception as e:
            logger.warning("nmcli connect failed, falling back: %s", e)

    try:
        ok = _wpa_connect(ssid, password)
        return {"ok": ok, "error": None if ok else "Failed to connect via wpa_supplicant"}
    except Exception as e:
        logger.error("All connect methods failed: %s", str(e))
        return {"ok": False, "error": str(e)}


def wifi_disconnect() -> dict[str, Any]:
    """Disconnect from the current WiFi network.

    Returns dict with ok (bool) key.
    """
    if _has_nmcli():
        try:
            ok = _nmcli_disconnect()
            return {"ok": ok}
        except Exception as e:
            logger.warning("nmcli disconnect failed: %s", e)
    try:
        _run(["sudo", "iwconfig", "wlan0", "essid", "off"], timeout=10, check=False)
        return {"ok": True}
    except Exception as e:
        logger.error("Disconnect failed: %s", str(e))
        return {"ok": False}


def wifi_forget(ssid: str) -> dict[str, Any]:
    """Forget a saved WiFi network.

    Args:
        ssid: Network SSID to forget.

    Returns:
        dict with ok (bool) key.
    """
    if _has_nmcli():
        try:
            ok = _nmcli_forget(ssid)
            return {"ok": ok}
        except Exception as e:
            logger.warning("nmcli forget failed: %s", e)
    # Fallback: remove from wpa_supplicant.conf
    try:
        if WPA_SUPPLICANT_CONF.exists():
            lines = WPA_SUPPLICANT_CONF.read_text().split("\n")
            new_lines: list[str] = []
            skip = False
            removed = False
            for line in lines:
                if f'ssid="{ssid}"' in line:
                    skip = True
                    removed = True
                    continue
                if skip and "}" in line:
                    skip = False
                    continue
                if not skip:
                    new_lines.append(line)
            if removed:
                WPA_SUPPLICANT_CONF.write_text("\n".join(new_lines))
                _run(["sudo", "wpa_cli", "-i", "wlan0", "reconfigure"], timeout=10, check=False)
                return {"ok": True}
        return {"ok": False}
    except Exception as e:
        logger.error("Forget failed: %s", str(e))
        return {"ok": False}


def hotspot_status() -> dict[str, Any]:
    """Get the current hotspot status.

    Returns dict with active (bool), ssid (str|None), clients (int), ip (str|None).
    """
    try:
        return _hotspot_status()
    except Exception as e:
        logger.error("Hotspot status failed: %s", str(e))
        return {"active": False, "ssid": None, "clients": 0, "ip": None}


def hotspot_enable(ssid: str = "Pi-Pedal", password: str = "pedal1234") -> dict[str, Any]:
    """Enable the WiFi hotspot.

    Args:
        ssid: Hotspot SSID.
        password: Hotspot password (min 8 chars).

    Returns:
        dict with ok (bool) and error (str|None) keys.
    """
    if len(password) < 8:
        return {"ok": False, "error": "Password must be at least 8 characters"}
    try:
        ok = _hotspot_enable(ssid, password)
        return {"ok": ok, "error": None if ok else "Failed to enable hotspot"}
    except Exception as e:
        logger.error("Hotspot enable failed: %s", str(e))
        return {"ok": False, "error": str(e)}


def hotspot_disable() -> dict[str, Any]:
    """Disable the WiFi hotspot.

    Returns dict with ok (bool) key.
    """
    try:
        ok = _hotspot_disable()
        return {"ok": ok}
    except Exception as e:
        logger.error("Hotspot disable failed: %s", str(e))
        return {"ok": False}