pi-pm3/app/backend/managers/wifi_manager.py

"""WiFi manager for network configuration and mode switching."""
import asyncio
import re
import subprocess
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Dict
from pathlib import Path

from .. import config


class WiFiMode(str, Enum):
    """WiFi operation modes."""
    AP = "ap"  # Access Point only
    CLIENT = "client"  # Client mode only
    DUAL = "dual"  # AP + Client (requires USB WiFi)
    AUTO = "auto"  # Automatically choose best mode
    OFF = "off"  # WiFi disabled


@dataclass
class WiFiInterface:
    """WiFi interface information."""
    name: str
    mac: str
    is_usb: bool
    is_up: bool
    connected: bool
    ssid: Optional[str] = None
    ip_address: Optional[str] = None


@dataclass
class WiFiNetwork:
    """Available WiFi network."""
    ssid: str
    bssid: str
    signal_strength: int
    frequency: int
    encrypted: bool
    in_use: bool = False


@dataclass
class WiFiStatus:
    """Current WiFi status."""
    mode: WiFiMode
    interfaces: List[WiFiInterface]
    current_ssid: Optional[str]
    current_ip: Optional[str]
    ap_ssid: Optional[str]
    ap_ip: Optional[str]
    supports_dual: bool


class WiFiManager:
    """Manages WiFi configuration and mode switching."""

    def __init__(self):
        """Initialize WiFi manager."""
        self._current_mode = WiFiMode.AUTO
        self._interfaces: List[WiFiInterface] = []
        self._supports_dual = False

    async def detect_interfaces(self) -> List[WiFiInterface]:
        """Detect available WiFi interfaces.

        Returns:
            List of WiFi interfaces found
        """
        interfaces = []

        try:
            # Get all wireless interfaces using iw
            result = await self._run_command("iw dev")

            if not result:
                return interfaces

            # Parse iw dev output
            current_interface = None
            for line in result.split('\n'):
                line = line.strip()

                if line.startswith('Interface '):
                    if current_interface:
                        interfaces.append(current_interface)

                    iface_name = line.split()[1]
                    current_interface = {
                        'name': iface_name,
                        'mac': '',
                        'is_usb': self._is_usb_interface(iface_name),
                        'is_up': False,
                        'connected': False,
                        'ssid': None,
                        'ip_address': None
                    }

                elif current_interface and line.startswith('addr '):
                    current_interface['mac'] = line.split()[1]

                elif current_interface and line.startswith('ssid '):
                    current_interface['ssid'] = ' '.join(line.split()[1:])
                    current_interface['connected'] = True

            if current_interface:
                interfaces.append(current_interface)

            # Get interface status and IP addresses
            for iface_dict in interfaces:
                # Check if interface is up
                iface_dict['is_up'] = await self._is_interface_up(iface_dict['name'])

                # Get IP address if interface is up
                if iface_dict['is_up']:
                    iface_dict['ip_address'] = await self._get_interface_ip(iface_dict['name'])

                # Create WiFiInterface object
                wifi_iface = WiFiInterface(**iface_dict)

            # Convert dicts to WiFiInterface objects
            self._interfaces = [WiFiInterface(**iface) for iface in interfaces]

            # Check if dual mode is supported (requires at least 2 interfaces)
            self._supports_dual = len(self._interfaces) >= 2

            return self._interfaces

        except Exception as e:
            print(f"Error detecting WiFi interfaces: {e}")
            return []

    def _is_usb_interface(self, iface_name: str) -> bool:
        """Check if interface is USB-based.

        Args:
            iface_name: Interface name (e.g., wlan0)

        Returns:
            True if USB interface, False otherwise
        """
        try:
            # Check if device is USB by looking at sysfs
            device_path = Path(f"/sys/class/net/{iface_name}/device")
            if not device_path.exists():
                return False

            # Read uevent to check for USB
            uevent_path = device_path / "uevent"
            if uevent_path.exists():
                content = uevent_path.read_text()
                return "usb" in content.lower()

            return False
        except Exception:
            return False

    async def _is_interface_up(self, iface_name: str) -> bool:
        """Check if interface is up.

        Args:
            iface_name: Interface name

        Returns:
            True if interface is up
        """
        try:
            result = await self._run_command(f"ip link show {iface_name}")
            return "UP" in result
        except Exception:
            return False

    async def _get_interface_ip(self, iface_name: str) -> Optional[str]:
        """Get IP address of interface.

        Args:
            iface_name: Interface name

        Returns:
            IP address or None
        """
        try:
            result = await self._run_command(f"ip -4 addr show {iface_name}")
            match = re.search(r'inet (\d+\.\d+\.\d+\.\d+)', result)
            return match.group(1) if match else None
        except Exception:
            return None

    async def get_status(self) -> WiFiStatus:
        """Get current WiFi status.

        Returns:
            WiFiStatus object with current state
        """
        await self.detect_interfaces()

        # Determine current mode
        current_mode = await self._detect_current_mode()

        # Find client and AP interfaces
        client_ssid = None
        client_ip = None
        ap_ssid = None
        ap_ip = None

        for iface in self._interfaces:
            if iface.connected and iface.ssid:
                client_ssid = iface.ssid
                client_ip = iface.ip_address

            # Check if running as AP (typically 10.3.141.1)
            if iface.ip_address and iface.ip_address.startswith("10.3.141"):
                ap_ip = iface.ip_address
                # Try to get AP SSID from hostapd
                ap_ssid = await self._get_ap_ssid()

        return WiFiStatus(
            mode=current_mode,
            interfaces=self._interfaces,
            current_ssid=client_ssid,
            current_ip=client_ip,
            ap_ssid=ap_ssid or "raspi-webgui",  # Default from existing setup
            ap_ip=ap_ip or "10.3.141.1",
            supports_dual=self._supports_dual
        )

    async def _detect_current_mode(self) -> WiFiMode:
        """Detect current WiFi mode.

        Returns:
            Current WiFi mode
        """
        if not self._interfaces:
            return WiFiMode.OFF

        # Check if any interface is in AP mode
        has_ap = any(
            iface.ip_address and iface.ip_address.startswith("10.3.141")
            for iface in self._interfaces
        )

        # Check if any interface is connected as client
        has_client = any(iface.connected for iface in self._interfaces)

        if has_ap and has_client:
            return WiFiMode.DUAL
        elif has_ap:
            return WiFiMode.AP
        elif has_client:
            return WiFiMode.CLIENT
        else:
            return WiFiMode.AUTO

    async def _get_ap_ssid(self) -> Optional[str]:
        """Get AP SSID from hostapd config.

        Returns:
            AP SSID or None
        """
        try:
            config_path = Path("/etc/hostapd/hostapd.conf")
            if config_path.exists():
                content = config_path.read_text()
                match = re.search(r'ssid=(.+)', content)
                return match.group(1) if match else None
        except Exception:
            return None

    async def scan_networks(self, interface: Optional[str] = None) -> List[WiFiNetwork]:
        """Scan for available WiFi networks.

        Args:
            interface: Interface to scan with (default: first available)

        Returns:
            List of available networks
        """
        if not interface:
            # Use first non-USB interface, or first available
            for iface in self._interfaces:
                if not iface.is_usb:
                    interface = iface.name
                    break

            if not interface and self._interfaces:
                interface = self._interfaces[0].name

        if not interface:
            return []

        try:
            # Request scan
            await self._run_command(f"sudo iw dev {interface} scan trigger", check=False)

            # Wait for scan to complete
            await asyncio.sleep(2)

            # Get scan results
            result = await self._run_command(f"sudo iw dev {interface} scan")

            networks = []
            current_network = None

            for line in result.split('\n'):
                line = line.strip()

                if line.startswith('BSS '):
                    if current_network:
                        networks.append(current_network)

                    bssid = line.split()[1].rstrip('(')
                    current_network = {
                        'ssid': '',
                        'bssid': bssid,
                        'signal_strength': 0,
                        'frequency': 0,
                        'encrypted': False,
                        'in_use': False
                    }

                elif current_network:
                    if line.startswith('SSID: '):
                        current_network['ssid'] = line[6:]
                    elif line.startswith('freq: '):
                        current_network['frequency'] = int(line[6:])
                    elif line.startswith('signal: '):
                        # Parse signal strength (e.g., "-50.00 dBm")
                        signal = line[8:].split()[0]
                        current_network['signal_strength'] = int(float(signal))
                    elif 'RSN:' in line or 'WPA:' in line:
                        current_network['encrypted'] = True

            if current_network:
                networks.append(current_network)

            # Convert to WiFiNetwork objects and filter out empty SSIDs
            return [
                WiFiNetwork(**net)
                for net in networks
                if net['ssid']
            ]

        except Exception as e:
            print(f"Error scanning networks: {e}")
            return []

    async def set_mode(self, mode: WiFiMode) -> bool:
        """Set WiFi mode.

        Args:
            mode: Desired WiFi mode

        Returns:
            True if mode was set successfully
        """
        if mode == WiFiMode.DUAL and not self._supports_dual:
            raise ValueError("Dual mode requires USB WiFi adapter")

        try:
            if mode == WiFiMode.AP:
                await self._enable_ap_mode()
            elif mode == WiFiMode.CLIENT:
                await self._enable_client_mode()
            elif mode == WiFiMode.DUAL:
                await self._enable_dual_mode()
            elif mode == WiFiMode.OFF:
                await self._disable_wifi()
            elif mode == WiFiMode.AUTO:
                await self._enable_auto_mode()

            self._current_mode = mode
            return True
        except Exception as e:
            print(f"Error setting WiFi mode: {e}")
            return False

    async def _enable_ap_mode(self):
        """Enable Access Point mode."""
        # Start hostapd and dnsmasq for AP
        await self._run_command("sudo systemctl start hostapd")
        await self._run_command("sudo systemctl start dnsmasq")
        print("AP mode enabled")

    async def _enable_client_mode(self):
        """Enable Client mode."""
        # Stop AP services
        await self._run_command("sudo systemctl stop hostapd", check=False)
        await self._run_command("sudo systemctl stop dnsmasq", check=False)
        # Start wpa_supplicant
        await self._run_command("sudo systemctl start wpa_supplicant")
        print("Client mode enabled")

    async def _enable_dual_mode(self):
        """Enable Dual mode (AP + Client)."""
        # Enable both AP and client
        await self._run_command("sudo systemctl start hostapd")
        await self._run_command("sudo systemctl start dnsmasq")
        await self._run_command("sudo systemctl start wpa_supplicant")
        print("Dual mode enabled")

    async def _disable_wifi(self):
        """Disable all WiFi."""
        await self._run_command("sudo systemctl stop hostapd", check=False)
        await self._run_command("sudo systemctl stop dnsmasq", check=False)
        await self._run_command("sudo systemctl stop wpa_supplicant", check=False)
        print("WiFi disabled")

    async def _enable_auto_mode(self):
        """Enable Auto mode."""
        # Auto-detect best mode based on saved networks and hardware
        if self._supports_dual:
            await self._enable_dual_mode()
        else:
            # Check if we have saved networks
            saved_networks = await self.get_saved_networks()
            if saved_networks:
                await self._enable_client_mode()
            else:
                await self._enable_ap_mode()
        print("Auto mode enabled")

    async def connect_to_network(
        self,
        ssid: str,
        password: Optional[str] = None,
        interface: Optional[str] = None,
        hidden: bool = False
    ) -> bool:
        """Connect to a WiFi network.

        Args:
            ssid: Network SSID
            password: Network password (if encrypted)
            interface: Interface to use (default: first available)
            hidden: Whether network is hidden

        Returns:
            True if connection successful
        """
        if not interface:
            # Use first non-USB interface, or first available
            for iface in self._interfaces:
                if not iface.is_usb:
                    interface = iface.name
                    break
            if not interface and self._interfaces:
                interface = self._interfaces[0].name

        if not interface:
            return False

        try:
            # Add network to wpa_supplicant configuration
            config_path = Path("/etc/wpa_supplicant/wpa_supplicant.conf")

            # Create network block
            if password:
                # Encrypted network
                network_block = f"""
network={{
    ssid="{ssid}"
    psk="{password}"
    scan_ssid={1 if hidden else 0}
    priority=1
}}
"""
            else:
                # Open network
                network_block = f"""
network={{
    ssid="{ssid}"
    key_mgmt=NONE
    scan_ssid={1 if hidden else 0}
    priority=1
}}
"""

            # Append to config (or use wpa_cli)
            # Using wpa_cli for immediate connection
            await self._add_network_via_wpa_cli(ssid, password, hidden)

            # Reconnect wpa_supplicant
            await self._run_command(f"sudo wpa_cli -i {interface} reconfigure")

            # Wait for connection
            await asyncio.sleep(3)

            # Verify connection
            result = await self._run_command(f"sudo wpa_cli -i {interface} status")
            if f'ssid={ssid}' in result and 'wpa_state=COMPLETED' in result:
                print(f"Successfully connected to {ssid}")
                return True
            else:
                print(f"Connection to {ssid} failed")
                return False

        except Exception as e:
            print(f"Error connecting to network: {e}")
            return False

    async def _add_network_via_wpa_cli(self, ssid: str, password: Optional[str], hidden: bool):
        """Add network using wpa_cli commands."""
        try:
            # Add network
            result = await self._run_command("sudo wpa_cli add_network")
            network_id = result.strip().split('\n')[-1]

            # Set SSID
            await self._run_command(f'sudo wpa_cli set_network {network_id} ssid \\"{ssid}\\"')

            # Set password or open network
            if password:
                await self._run_command(f'sudo wpa_cli set_network {network_id} psk \\"{password}\\"')
            else:
                await self._run_command(f'sudo wpa_cli set_network {network_id} key_mgmt NONE')

            # Set scan_ssid for hidden networks
            if hidden:
                await self._run_command(f'sudo wpa_cli set_network {network_id} scan_ssid 1')

            # Enable network
            await self._run_command(f'sudo wpa_cli enable_network {network_id}')

            # Save configuration
            await self._run_command('sudo wpa_cli save_config')

            return True
        except Exception as e:
            print(f"Error adding network via wpa_cli: {e}")
            return False

    async def disconnect_from_network(self, interface: Optional[str] = None) -> bool:
        """Disconnect from current network.

        Args:
            interface: Interface to disconnect (default: first connected)

        Returns:
            True if disconnection successful
        """
        if not interface:
            for iface in self._interfaces:
                if iface.connected:
                    interface = iface.name
                    break

        if not interface:
            return False

        try:
            await self._run_command(f"sudo wpa_cli -i {interface} disconnect")
            return True
        except Exception as e:
            print(f"Error disconnecting: {e}")
            return False

    async def get_saved_networks(self) -> List[Dict[str, str]]:
        """Get list of saved networks from wpa_supplicant.

        Returns:
            List of saved networks with SSID and other info
        """
        try:
            result = await self._run_command("sudo wpa_cli list_networks")
            networks = []

            for line in result.split('\n')[1:]:  # Skip header
                if line.strip():
                    parts = line.split('\t')
                    if len(parts) >= 2:
                        networks.append({
                            'id': parts[0].strip(),
                            'ssid': parts[1].strip(),
                            'enabled': 'CURRENT' in line or 'ENABLED' in line
                        })

            return networks
        except Exception as e:
            print(f"Error getting saved networks: {e}")
            return []

    async def forget_network(self, ssid: str) -> bool:
        """Forget a saved network.

        Args:
            ssid: SSID of network to forget

        Returns:
            True if network was forgotten
        """
        try:
            # Get network ID
            networks = await self.get_saved_networks()
            network_id = None

            for net in networks:
                if net['ssid'] == ssid:
                    network_id = net['id']
                    break

            if network_id is None:
                return False

            # Remove network
            await self._run_command(f"sudo wpa_cli remove_network {network_id}")
            await self._run_command("sudo wpa_cli save_config")

            return True
        except Exception as e:
            print(f"Error forgetting network: {e}")
            return False

    async def set_static_ip(
        self,
        interface: str,
        ip_address: str,
        netmask: str = "255.255.255.0",
        gateway: Optional[str] = None,
        dns: Optional[List[str]] = None
    ) -> bool:
        """Set static IP for interface.

        Args:
            interface: Interface name
            ip_address: Static IP address
            netmask: Network mask
            gateway: Gateway IP (optional)
            dns: DNS servers (optional)

        Returns:
            True if configuration successful
        """
        try:
            # Configure static IP using dhcpcd
            config_path = Path("/etc/dhcpcd.conf")

            # Read existing config
            if config_path.exists():
                content = config_path.read_text()
            else:
                content = ""

            # Remove existing config for this interface
            lines = content.split('\n')
            new_lines = []
            skip_interface = False

            for line in lines:
                if line.startswith(f'interface {interface}'):
                    skip_interface = True
                    continue
                if skip_interface and (line.startswith('interface ') or line.strip() == ''):
                    skip_interface = False
                if not skip_interface:
                    new_lines.append(line)

            # Add new configuration
            new_config = f"\ninterface {interface}\n"
            new_config += f"static ip_address={ip_address}/{self._netmask_to_cidr(netmask)}\n"

            if gateway:
                new_config += f"static routers={gateway}\n"

            if dns:
                new_config += f"static domain_name_servers={' '.join(dns)}\n"

            new_content = '\n'.join(new_lines) + new_config

            # Write configuration (would need sudo)
            # In production, this should use a proper mechanism
            print(f"Would write to {config_path}:\n{new_config}")

            # Restart interface
            await self._run_command(f"sudo ip link set {interface} down")
            await self._run_command(f"sudo ip link set {interface} up")
            await self._run_command("sudo systemctl restart dhcpcd")

            return True
        except Exception as e:
            print(f"Error setting static IP: {e}")
            return False

    def _netmask_to_cidr(self, netmask: str) -> int:
        """Convert netmask to CIDR notation.

        Args:
            netmask: Netmask (e.g., 255.255.255.0)

        Returns:
            CIDR prefix length (e.g., 24)
        """
        return sum([bin(int(x)).count('1') for x in netmask.split('.')])

    async def enable_dhcp(self, interface: str) -> bool:
        """Enable DHCP for interface.

        Args:
            interface: Interface name

        Returns:
            True if DHCP enabled
        """
        try:
            # Remove static IP configuration
            config_path = Path("/etc/dhcpcd.conf")

            if config_path.exists():
                content = config_path.read_text()
                lines = content.split('\n')
                new_lines = []
                skip_interface = False

                for line in lines:
                    if line.startswith(f'interface {interface}'):
                        skip_interface = True
                        continue
                    if skip_interface and (line.startswith('interface ') or line.strip() == ''):
                        skip_interface = False
                    if not skip_interface:
                        new_lines.append(line)

                # Write back
                print(f"Would update {config_path} to enable DHCP")

            # Restart interface
            await self._run_command(f"sudo ip link set {interface} down")
            await self._run_command(f"sudo ip link set {interface} up")
            await self._run_command("sudo systemctl restart dhcpcd")

            return True
        except Exception as e:
            print(f"Error enabling DHCP: {e}")
            return False

    async def _run_command(self, command: str, check: bool = True) -> str:
        """Run a shell command asynchronously.

        Args:
            command: Command to run
            check: Raise exception on non-zero exit code

        Returns:
            Command output

        Raises:
            subprocess.CalledProcessError: If command fails and check=True
        """
        loop = asyncio.get_event_loop()

        def _execute():
            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                check=check
            )
            return result.stdout

        return await loop.run_in_executor(None, _execute)


# Global instance
wifi_manager = WiFiManager()