#!/usr/bin/env python3
"""WiFi client module — managed wpa_supplicant for WiFi connectivity.

Supports WPA2-PSK and WPA2-Enterprise (PEAP/EAP-TLS).  Generates
wpa_supplicant.conf from config parameters, manages the wpa_supplicant
process, and monitors connection quality.
"""

import logging
import os
import subprocess
import tempfile
import time
from typing import Optional

from modules.base import BaseModule

logger = logging.getLogger(__name__)

WPA_SUPPLICANT_BIN = "/sbin/wpa_supplicant"
WPA_CLI_BIN = "/sbin/wpa_cli"
DEFAULT_INTERFACE = "wlan0"
CONNECT_TIMEOUT = 30  # seconds


class WiFiClient(BaseModule):
    """WiFi client — wpa_supplicant managed WPA2-PSK/Enterprise."""

    name = "wifi_client"
    module_type = "connectivity"
    priority = -100
    requires_root = True

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._wpa_proc: Optional[subprocess.Popen] = None
        self._interface = DEFAULT_INTERFACE
        self._config_path: Optional[str] = None
        self._connected = False
        self._ssid: Optional[str] = None

    # ------------------------------------------------------------------
    # BaseModule interface
    # ------------------------------------------------------------------

    def start(self) -> None:
        if self._running:
            return

        wifi_cfg = self.config.get("connectivity", {}).get("wifi", {})
        if not wifi_cfg:
            logger.error("No wifi config in bigbrother.yaml")
            self.state.set_module_status(self.name, "error")
            return

        self._interface = wifi_cfg.get("interface", DEFAULT_INTERFACE)
        ssid = wifi_cfg.get("ssid")
        psk = wifi_cfg.get("psk")

        if not ssid:
            logger.error("WiFi SSID not configured")
            self.state.set_module_status(self.name, "error")
            return

        # Determine auth type
        eap_cfg = wifi_cfg.get("eap")
        if eap_cfg:
            self._config_path = self._write_enterprise_config(
                ssid=ssid,
                eap_method=eap_cfg.get("method", "PEAP"),
                identity=eap_cfg.get("identity", ""),
                password=eap_cfg.get("password", ""),
                ca_cert=eap_cfg.get("ca_cert"),
                client_cert=eap_cfg.get("client_cert"),
                private_key=eap_cfg.get("private_key"),
                phase2=eap_cfg.get("phase2", "auth=MSCHAPV2"),
            )
        elif psk:
            self._config_path = self._write_psk_config(ssid, psk)
        else:
            logger.error("WiFi config has no PSK and no EAP section")
            self.state.set_module_status(self.name, "error")
            return

        if not self.connect(ssid, psk):
            logger.error("Failed to connect to WiFi network '%s'", ssid)
            self.state.set_module_status(self.name, "error")
            return

        self._running = True
        self._pid = os.getpid()
        self._start_time = time.time()
        self._ssid = ssid
        self._connected = True

        self.state.set_module_status(self.name, "running", pid=self._pid)
        self.bus.emit("CONNECTIVITY_CHANGED",
                      {"module": self.name, "state": "up", "ssid": ssid},
                      source_module=self.name)
        logger.info("WiFi connected to '%s' on %s", ssid, self._interface)

    def stop(self) -> None:
        if not self._running:
            return

        self.disconnect()
        self._running = False
        self._connected = False

        # Cleanup config file
        if self._config_path and os.path.isfile(self._config_path):
            try:
                os.unlink(self._config_path)
            except OSError:
                pass

        self.state.set_module_status(self.name, "stopped")
        self.bus.emit("CONNECTIVITY_CHANGED",
                      {"module": self.name, "state": "down"},
                      source_module=self.name)
        logger.info("WiFi client stopped")

    def status(self) -> dict:
        signal = self.get_signal() if self._running else None
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "connected": self._connected,
            "ssid": self._ssid,
            "interface": self._interface,
            "signal_dbm": signal,
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)

    def health_check(self) -> bool:
        if not self._running:
            return False
        return self.is_connected()

    # ------------------------------------------------------------------
    # WiFi operations
    # ------------------------------------------------------------------

    def connect(self, ssid: str, psk: str = None) -> bool:
        """Start wpa_supplicant and connect to the configured network."""
        # Kill any existing wpa_supplicant on this interface
        self._kill_existing_wpa()

        # Bring interface up
        try:
            subprocess.run(
                ["ip", "link", "set", "up", "dev", self._interface],
                check=True, capture_output=True, timeout=10,
            )
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
            logger.error("Failed to bring up %s: %s", self._interface, exc)
            return False

        # Start wpa_supplicant
        cmd = [
            WPA_SUPPLICANT_BIN,
            "-i", self._interface,
            "-c", self._config_path,
            "-B",  # background
            "-D", "nl80211,wext",
            "-P", f"/var/run/wpa_supplicant_{self._interface}.pid",
        ]

        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
            if result.returncode != 0:
                logger.error("wpa_supplicant failed to start: %s", result.stderr.strip())
                return False
        except (subprocess.TimeoutExpired, FileNotFoundError) as exc:
            logger.error("wpa_supplicant error: %s", exc)
            return False

        # Wait for connection
        return self._wait_for_connection(timeout=CONNECT_TIMEOUT)

    def disconnect(self) -> None:
        """Stop wpa_supplicant and release the interface."""
        self._kill_existing_wpa()
        self._connected = False

        # Release DHCP lease
        try:
            subprocess.run(
                ["dhclient", "-r", self._interface],
                capture_output=True, timeout=10,
            )
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError):
            pass

    def scan(self) -> list:
        """Scan for available WiFi networks. Returns list of dicts."""
        networks = []
        try:
            # Trigger scan
            subprocess.run(
                [WPA_CLI_BIN, "-i", self._interface, "scan"],
                capture_output=True, timeout=10,
            )
            time.sleep(3.0)  # Allow scan to complete

            # Get results
            result = subprocess.run(
                [WPA_CLI_BIN, "-i", self._interface, "scan_results"],
                capture_output=True, text=True, timeout=10,
            )

            if result.returncode == 0:
                for line in result.stdout.strip().splitlines()[1:]:  # Skip header
                    parts = line.split("\t")
                    if len(parts) >= 5:
                        networks.append({
                            "bssid": parts[0],
                            "frequency": int(parts[1]),
                            "signal": int(parts[2]),
                            "flags": parts[3],
                            "ssid": parts[4] if len(parts) > 4 else "",
                        })

        except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError) as exc:
            logger.warning("WiFi scan failed: %s", exc)

        # Fallback to iw if wpa_cli not available
        if not networks:
            networks = self._scan_iw()

        return networks

    def is_connected(self) -> bool:
        """Check if WiFi is associated and has an IP address."""
        try:
            result = subprocess.run(
                [WPA_CLI_BIN, "-i", self._interface, "status"],
                capture_output=True, text=True, timeout=5,
            )
            if result.returncode != 0:
                return False

            status_dict = {}
            for line in result.stdout.strip().splitlines():
                if "=" in line:
                    k, v = line.split("=", 1)
                    status_dict[k.strip()] = v.strip()

            wpa_state = status_dict.get("wpa_state", "")
            self._connected = wpa_state == "COMPLETED"
            return self._connected

        except (subprocess.TimeoutExpired, FileNotFoundError):
            return False

    def get_signal(self) -> Optional[int]:
        """Get current signal strength in dBm."""
        try:
            result = subprocess.run(
                [WPA_CLI_BIN, "-i", self._interface, "signal_poll"],
                capture_output=True, text=True, timeout=5,
            )
            if result.returncode == 0:
                for line in result.stdout.strip().splitlines():
                    if line.startswith("RSSI="):
                        return int(line.split("=", 1)[1])
        except (subprocess.TimeoutExpired, FileNotFoundError, ValueError):
            pass

        # Fallback to /proc/net/wireless
        try:
            with open("/proc/net/wireless", "r") as f:
                for line in f:
                    if self._interface in line:
                        parts = line.split()
                        if len(parts) >= 4:
                            return int(float(parts[3]))
        except (IOError, ValueError, IndexError):
            pass

        return None

    # ------------------------------------------------------------------
    # Config generation
    # ------------------------------------------------------------------

    def _write_psk_config(self, ssid: str, psk: str) -> str:
        """Write WPA2-PSK wpa_supplicant.conf."""
        config = (
            "ctrl_interface=/var/run/wpa_supplicant\n"
            "ctrl_interface_group=0\n"
            "update_config=0\n"
            "ap_scan=1\n"
            "\n"
            "network={\n"
            f'    ssid="{ssid}"\n'
            f'    psk="{psk}"\n'
            "    key_mgmt=WPA-PSK\n"
            "    proto=RSN\n"
            "    pairwise=CCMP\n"
            "    group=CCMP\n"
            "}\n"
        )
        return self._write_config_file(config)

    def _write_enterprise_config(self, ssid: str, eap_method: str,
                                 identity: str, password: str,
                                 ca_cert: str = None,
                                 client_cert: str = None,
                                 private_key: str = None,
                                 phase2: str = "auth=MSCHAPV2") -> str:
        """Write WPA2-Enterprise wpa_supplicant.conf (PEAP or EAP-TLS)."""
        config = (
            "ctrl_interface=/var/run/wpa_supplicant\n"
            "ctrl_interface_group=0\n"
            "update_config=0\n"
            "ap_scan=1\n"
            "\n"
            "network={\n"
            f'    ssid="{ssid}"\n'
            "    key_mgmt=WPA-EAP\n"
            f'    eap={eap_method}\n'
            f'    identity="{identity}"\n'
        )

        if eap_method.upper() == "PEAP":
            config += f'    password="{password}"\n'
            config += f'    phase2="{phase2}"\n'
            if ca_cert:
                config += f'    ca_cert="{ca_cert}"\n'
        elif eap_method.upper() == "TLS":
            if ca_cert:
                config += f'    ca_cert="{ca_cert}"\n'
            if client_cert:
                config += f'    client_cert="{client_cert}"\n'
            if private_key:
                config += f'    private_key="{private_key}"\n'

        config += "}\n"
        return self._write_config_file(config)

    def _write_config_file(self, content: str) -> str:
        """Write config to tmpfs and return path."""
        tmpdir = "/dev/shm" if os.path.isdir("/dev/shm") else tempfile.gettempdir()
        config_path = os.path.join(tmpdir, f".wpa_{self._interface}.conf")

        fd = os.open(config_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        try:
            os.write(fd, content.encode())
        finally:
            os.close(fd)

        return config_path

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _wait_for_connection(self, timeout: int = CONNECT_TIMEOUT) -> bool:
        """Wait for wpa_supplicant to complete association."""
        deadline = time.time() + timeout
        while time.time() < deadline:
            if self.is_connected():
                # Request DHCP lease
                self._request_dhcp()
                return True
            time.sleep(1.0)

        logger.warning("WiFi connection timed out after %ds", timeout)
        return False

    def _request_dhcp(self) -> None:
        """Request a DHCP lease on the WiFi interface."""
        try:
            subprocess.run(
                ["dhclient", "-1", "-nw", self._interface],
                capture_output=True, timeout=15,
            )
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError):
            logger.debug("dhclient not available, trying dhcpcd")
            try:
                subprocess.run(
                    ["dhcpcd", "-1", self._interface],
                    capture_output=True, timeout=15,
                )
            except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError):
                logger.warning("No DHCP client available for %s", self._interface)

    def _kill_existing_wpa(self) -> None:
        """Kill any existing wpa_supplicant on our interface."""
        pid_file = f"/var/run/wpa_supplicant_{self._interface}.pid"
        if os.path.isfile(pid_file):
            try:
                pid = int(open(pid_file).read().strip())
                os.kill(pid, 15)  # SIGTERM
                time.sleep(0.5)
            except (ValueError, ProcessLookupError, PermissionError, IOError):
                pass

        # Also kill by interface match
        try:
            subprocess.run(
                ["pkill", "-f", f"wpa_supplicant.*{self._interface}"],
                capture_output=True, timeout=5,
            )
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError):
            pass

        time.sleep(0.5)

    def _scan_iw(self) -> list:
        """Fallback scan using iw command."""
        networks = []
        try:
            result = subprocess.run(
                ["iw", "dev", self._interface, "scan", "-u"],
                capture_output=True, text=True, timeout=15,
            )
            if result.returncode != 0:
                return networks

            current = {}
            for line in result.stdout.splitlines():
                line = line.strip()
                if line.startswith("BSS "):
                    if current:
                        networks.append(current)
                    bssid = line.split()[1].split("(")[0]
                    current = {"bssid": bssid, "ssid": "", "signal": 0, "frequency": 0}
                elif line.startswith("SSID:"):
                    current["ssid"] = line.split(":", 1)[1].strip()
                elif line.startswith("signal:"):
                    try:
                        current["signal"] = int(float(line.split(":")[1].strip().split()[0]))
                    except (ValueError, IndexError):
                        pass
                elif line.startswith("freq:"):
                    try:
                        current["frequency"] = int(line.split(":")[1].strip())
                    except (ValueError, IndexError):
                        pass

            if current:
                networks.append(current)

        except (subprocess.TimeoutExpired, FileNotFoundError):
            pass

        return networks
