"""Network interface probing for distributed mesh system.

Probes NICs and network performance to feed into a JEPA-based system optimizer.
Uses psutil, subprocess, and /proc /sys filesystem reads.
"""
from __future__ import annotations

import json
import logging
import re
import subprocess
import socket
import struct
import threading
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

import psutil

logger = logging.getLogger(__name__)


class NetworkProbe:
    """Probes local and remote network characteristics for mesh optimization."""

    def detect_nics(self) -> List[Dict]:
        """Detect network interfaces via psutil and Linux sysfs/ethtool.

        Returns:
            List of dicts with keys: name, mac, ipv4, speed_mbps, mtu,
            numa_node, driver, is_up.
        """
        interfaces: List[Dict] = []

        try:
            addrs = psutil.net_if_addrs()
            stats = psutil.net_if_stats()
        except Exception as exc:
            logger.warning("psutil NIC detection failed: %s", exc)
            return interfaces

        for name in sorted(addrs.keys()):
            if name == "lo":
                continue

            iface: Dict[str, Any] = {"name": name}

            # Extract MAC and IPv4 from psutil addresses
            mac = None
            ipv4 = None
            for addr in addrs.get(name, []):
                if addr.family == socket.AF_INET:
                    ipv4 = addr.address
                elif addr.family == psutil.AF_LINK:
                    mac = addr.address
            iface["mac"] = mac or "unknown"
            iface["ipv4"] = ipv4

            # Stats from psutil
            st = stats.get(name)
            if st:
                iface["mtu"] = st.mtu
                iface["is_up"] = st.isup
            else:
                iface["mtu"] = -1
                iface["is_up"] = False

            # Speed from sysfs (psutil speed can be unreliable)
            speed_path = Path(f"/sys/class/net/{name}/speed")
            try:
                iface["speed_mbps"] = int(speed_path.read_text().strip())
            except Exception:
                # Fall back to psutil speed if sysfs unavailable
                iface["speed_mbps"] = st.speed if st and st.speed > 0 else -1

            # NUMA node from sysfs
            numa_path = Path(f"/sys/class/net/{name}/device/numa_node")
            try:
                if numa_path.exists():
                    iface["numa_node"] = int(numa_path.read_text().strip())
                else:
                    iface["numa_node"] = -1
            except Exception:
                iface["numa_node"] = -1

            # Driver from ethtool -i
            iface["driver"] = "unknown"
            try:
                out = subprocess.check_output(
                    ["ethtool", "-i", name],
                    stderr=subprocess.DEVNULL,
                    timeout=5,
                ).decode()
                for line in out.splitlines():
                    if line.startswith("driver:"):
                        iface["driver"] = line.split(":", 1)[1].strip()
                        break
            except Exception:
                # Fallback: resolve sysfs driver symlink
                try:
                    driver_link = Path(f"/sys/class/net/{name}/device/driver")
                    if driver_link.is_symlink() or driver_link.exists():
                        import os
                        iface["driver"] = os.path.basename(
                            os.readlink(str(driver_link))
                        )
                except Exception:
                    pass

            interfaces.append(iface)

        logger.info("Detected %d network interfaces", len(interfaces))
        return interfaces

    def probe_tcp_throughput(
        self,
        target_host: str,
        port: int = 5201,
        duration: int = 5,
    ) -> Dict:
        """Measure TCP throughput to target_host.

        Tries iperf3 first (subprocess). If iperf3 is not available, falls
        back to a simple socket-based throughput test.

        Args:
            target_host: Remote host to test against.
            port: Port for iperf3 or socket test.
            duration: Test duration in seconds.

        Returns:
            Dict with bandwidth_mbps, retransmits (iperf3 only), method used.
        """
        result: Dict[str, Any] = {
            "target": target_host,
            "port": port,
            "bandwidth_mbps": 0.0,
            "retransmits": -1,
            "method": "none",
        }

        # Try iperf3 first
        if self._try_iperf3(target_host, port, duration, result):
            return result

        # Fallback: socket-based throughput test
        self._socket_throughput(target_host, port, duration, result)
        return result

    def _try_iperf3(
        self, host: str, port: int, duration: int, result: Dict
    ) -> bool:
        """Attempt iperf3 throughput test. Returns True if successful."""
        try:
            cmd = [
                "iperf3", "-c", host, "-p", str(port),
                "-t", str(duration), "-J",
            ]
            out = subprocess.check_output(
                cmd, stderr=subprocess.DEVNULL, timeout=duration + 30
            ).decode()
            data = json.loads(out)

            end = data.get("end", {})
            sum_sent = end.get("sum_sent", {})
            bw = sum_sent.get("bits_per_second", 0)
            result["bandwidth_mbps"] = round(bw / 1e6, 2)
            result["retransmits"] = sum_sent.get("retransmits", -1)
            result["method"] = "iperf3"

            logger.info(
                "iperf3 to %s: %.2f Mbps, %d retransmits",
                host, result["bandwidth_mbps"], result["retransmits"],
            )
            return True
        except FileNotFoundError:
            logger.debug("iperf3 not found, falling back to socket test")
            return False
        except subprocess.TimeoutExpired:
            logger.warning("iperf3 to %s timed out", host)
            result["error"] = "iperf3 timed out"
            return False
        except Exception as exc:
            logger.debug("iperf3 failed (%s), falling back to socket test", exc)
            return False

    def _socket_throughput(
        self, host: str, port: int, duration: int, result: Dict
    ) -> None:
        """Simple socket-based throughput test between mesh nodes."""
        chunk = b"\x00" * (64 * 1024)
        bytes_received = [0]
        server_ready = threading.Event()
        server_error: List[Optional[str]] = [None]
        is_local = host in ("localhost", "127.0.0.1", "::1")

        def _receiver():
            try:
                srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                srv.bind(("0.0.0.0", port))
                srv.listen(1)
                srv.settimeout(10)
                server_ready.set()
                conn, _ = srv.accept()
                conn.settimeout(duration + 5)
                try:
                    while True:
                        data = conn.recv(65536)
                        if not data:
                            break
                        bytes_received[0] += len(data)
                except socket.timeout:
                    pass
                finally:
                    conn.close()
                    srv.close()
            except Exception as exc:
                server_error[0] = str(exc)
                server_ready.set()

        if is_local:
            recv_thread = threading.Thread(target=_receiver, daemon=True)
            recv_thread.start()
            server_ready.wait(timeout=5)
            if server_error[0]:
                result["error"] = f"Receiver failed: {server_error[0]}"
                return

        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(10)
            sock.connect((host, port))
            try:
                sock.setsockopt(
                    socket.SOL_SOCKET, socket.SO_SNDBUF, 4 * 1024 * 1024
                )
            except Exception:
                pass

            bytes_sent = 0
            t0 = time.perf_counter()
            deadline = t0 + duration
            while time.perf_counter() < deadline:
                try:
                    n = sock.send(chunk)
                    bytes_sent += n
                except (socket.timeout, BrokenPipeError):
                    break
            elapsed = time.perf_counter() - t0
            sock.close()

            if is_local:
                recv_thread.join(timeout=5)

            if elapsed > 0:
                # Convert bytes/sec to megabits/sec
                result["bandwidth_mbps"] = round(
                    (bytes_sent * 8) / (1e6 * elapsed), 2
                )
            result["method"] = "socket"
            result["bytes_sent"] = bytes_sent

            logger.info(
                "Socket throughput to %s: %.2f Mbps (%d bytes in %.1f s)",
                host, result["bandwidth_mbps"], bytes_sent, elapsed,
            )
        except Exception as exc:
            result["error"] = str(exc)
            result["method"] = "socket"
            logger.warning("Socket throughput test to %s failed: %s", host, exc)

    def probe_latency(self, target_host: str, count: int = 100) -> Dict:
        """Measure ICMP ping latency to target_host.

        Args:
            target_host: Host to ping.
            count: Number of ping packets to send.

        Returns:
            Dict with min_ms, avg_ms, max_ms, stddev_ms, loss_pct.
        """
        result: Dict[str, Any] = {
            "target": target_host,
            "count": count,
            "min_ms": 0.0,
            "avg_ms": 0.0,
            "max_ms": 0.0,
            "stddev_ms": 0.0,
            "loss_pct": 100.0,
        }

        try:
            cmd = [
                "ping", "-c", str(count), "-i", "0.05", "-W", "2",
                target_host,
            ]
            out = subprocess.check_output(
                cmd, stderr=subprocess.STDOUT, timeout=count * 2 + 10
            ).decode()
            self._parse_ping_output(out, result)
        except subprocess.CalledProcessError as exc:
            # ping returns non-zero on partial packet loss
            out = exc.output.decode() if exc.output else ""
            self._parse_ping_output(out, result)
            logger.warning("Ping to %s had packet loss", target_host)
        except subprocess.TimeoutExpired:
            result["error"] = "ping timed out"
            logger.warning("Ping to %s timed out", target_host)
        except FileNotFoundError:
            result["error"] = "ping command not found"
            logger.warning("ping binary not found")
        except Exception as exc:
            result["error"] = str(exc)
            logger.warning("Latency probe to %s failed: %s", target_host, exc)

        logger.info(
            "Latency to %s: avg %.2f ms, loss %.1f%%",
            target_host, result["avg_ms"], result["loss_pct"],
        )
        return result

    @staticmethod
    def _parse_ping_output(out: str, result: Dict) -> None:
        """Parse ping command output into result dict."""
        loss_match = re.search(r"(\d+(?:\.\d+)?)% packet loss", out)
        if loss_match:
            result["loss_pct"] = float(loss_match.group(1))

        rtt_match = re.search(
            r"rtt min/avg/max/mdev = "
            r"([\d.]+)/([\d.]+)/([\d.]+)/([\d.]+)",
            out,
        )
        if rtt_match:
            result["min_ms"] = float(rtt_match.group(1))
            result["avg_ms"] = float(rtt_match.group(2))
            result["max_ms"] = float(rtt_match.group(3))
            result["stddev_ms"] = float(rtt_match.group(4))

    def probe_buffer_sizes(self) -> Dict:
        """Read TCP/socket buffer size settings from /proc/sys.

        Returns:
            Dict with rmem_max, wmem_max, tcp_rmem (min/default/max),
            tcp_wmem (min/default/max).
        """
        result: Dict[str, Any] = {}

        def _read_single(path: str) -> int:
            try:
                return int(Path(path).read_text().strip())
            except Exception as exc:
                logger.debug("Could not read %s: %s", path, exc)
                return 0

        def _read_triple(path: str) -> Dict[str, int]:
            try:
                parts = Path(path).read_text().strip().split()
                return {
                    "min": int(parts[0]),
                    "default": int(parts[1]),
                    "max": int(parts[2]),
                }
            except Exception as exc:
                logger.debug("Could not read %s: %s", path, exc)
                return {"min": 0, "default": 0, "max": 0}

        result["rmem_max"] = _read_single("/proc/sys/net/core/rmem_max")
        result["wmem_max"] = _read_single("/proc/sys/net/core/wmem_max")
        result["tcp_rmem"] = _read_triple("/proc/sys/net/ipv4/tcp_rmem")
        result["tcp_wmem"] = _read_triple("/proc/sys/net/ipv4/tcp_wmem")

        logger.info(
            "Buffer sizes: rmem_max=%d, wmem_max=%d",
            result["rmem_max"], result["wmem_max"],
        )
        return result

    def probe_irq_distribution(self) -> Dict:
        """Parse /proc/interrupts for network-related IRQs.

        Returns:
            Dict with:
              - irqs: {irq_label: [count_per_cpu, ...]}
              - cpu_totals: {cpu_index: total_nic_irq_count}
              - handling_cpus: list of CPU indices that handle NIC IRQs
        """
        result: Dict[str, Any] = {
            "irqs": {},
            "cpu_totals": {},
            "handling_cpus": [],
        }

        nic_patterns = re.compile(
            r"(eth|eno|enp|ens|enx|wl|ib|mlx|i40e|ixgbe|igb|e1000|"
            r"r8169|virtio|xen-net|be2net|bnxt|ice|ena)",
            re.IGNORECASE,
        )

        try:
            with open("/proc/interrupts") as f:
                header = f.readline()
                cpu_cols = [c for c in header.split() if c.startswith("CPU")]
                num_cpus = len(cpu_cols)

                cpu_totals: Dict[int, int] = {i: 0 for i in range(num_cpus)}

                for line in f:
                    if not nic_patterns.search(line):
                        continue

                    parts = line.split()
                    if len(parts) < num_cpus + 2:
                        continue

                    irq_id = parts[0].rstrip(":")
                    counts: List[int] = []
                    for i in range(1, num_cpus + 1):
                        try:
                            c = int(parts[i])
                            counts.append(c)
                        except (ValueError, IndexError):
                            break

                    if len(counts) != num_cpus:
                        continue

                    # Build label from description columns
                    desc_parts = parts[num_cpus + 1:]
                    label = f"IRQ{irq_id}_{'_'.join(desc_parts)}"
                    result["irqs"][label] = counts

                    for cpu_idx, cnt in enumerate(counts):
                        cpu_totals[cpu_idx] += cnt

                result["cpu_totals"] = cpu_totals
                result["handling_cpus"] = [
                    cpu for cpu, total in cpu_totals.items() if total > 0
                ]

        except Exception as exc:
            logger.warning("Could not read /proc/interrupts: %s", exc)
            result["error"] = str(exc)

        logger.info(
            "Found %d NIC-related IRQ entries across %d CPUs",
            len(result["irqs"]), len(result["handling_cpus"]),
        )
        return result

    def probe_numa_affinity(self) -> Dict:
        """Report NUMA affinity for each NIC and IRQ pinning status.

        For each NIC, reports its NUMA node and whether its IRQs are
        pinned to CPUs on that same NUMA node.

        Returns:
            Dict of {nic_name: {numa_node, irq_cpus, properly_pinned}}.
        """
        result: Dict[str, Any] = {}
        net_dir = Path("/sys/class/net")

        if not net_dir.exists():
            logger.warning("/sys/class/net not found")
            return result

        # Build NUMA-to-CPU mapping
        numa_cpus: Dict[int, set] = {}
        try:
            import os
            numa_dir = Path("/sys/devices/system/node")
            if numa_dir.exists():
                for node_entry in numa_dir.iterdir():
                    if not node_entry.name.startswith("node"):
                        continue
                    try:
                        node_id = int(node_entry.name[4:])
                    except ValueError:
                        continue
                    cpulist_path = node_entry / "cpulist"
                    if cpulist_path.exists():
                        cpulist = cpulist_path.read_text().strip()
                        numa_cpus[node_id] = self._parse_cpulist(cpulist)
        except Exception as exc:
            logger.debug("Could not build NUMA-CPU map: %s", exc)

        # Get IRQ distribution for cross-reference
        irq_dist = self.probe_irq_distribution()

        for entry in sorted(net_dir.iterdir()):
            name = entry.name
            if name == "lo":
                continue

            nic_info: Dict[str, Any] = {"numa_node": -1, "irq_cpus": [], "properly_pinned": False}

            # Get NIC NUMA node
            numa_path = entry / "device" / "numa_node"
            try:
                if numa_path.exists():
                    nic_info["numa_node"] = int(numa_path.read_text().strip())
            except Exception:
                pass

            # Find which CPUs handle this NIC's IRQs
            irq_cpus: set = set()
            for label, counts in irq_dist.get("irqs", {}).items():
                if name in label:
                    for cpu_idx, cnt in enumerate(counts):
                        if cnt > 0:
                            irq_cpus.add(cpu_idx)
            nic_info["irq_cpus"] = sorted(irq_cpus)

            # Check if IRQs are pinned to the NIC's NUMA node CPUs
            node = nic_info["numa_node"]
            if node >= 0 and node in numa_cpus and irq_cpus:
                nic_info["properly_pinned"] = irq_cpus.issubset(numa_cpus[node])

            result[name] = nic_info

        logger.info("NUMA affinity for %d NICs", len(result))
        return result

    @staticmethod
    def _parse_cpulist(cpulist: str) -> set:
        """Parse a CPU list string like '0-3,8-11' into a set of ints."""
        cpus: set = set()
        for part in cpulist.split(","):
            part = part.strip()
            if "-" in part:
                lo, hi = part.split("-", 1)
                cpus.update(range(int(lo), int(hi) + 1))
            else:
                cpus.add(int(part))
        return cpus

    def run_all(self, target_hosts: Optional[List[str]] = None) -> Dict:
        """Run all network probes and return combined results.

        Runs all local probes (NIC detection, buffers, IRQs, NUMA affinity),
        plus latency and throughput tests to each target host.

        Args:
            target_hosts: List of remote hosts to probe. If None, only
                          local probes are run.

        Returns:
            Dict with all probe results keyed by probe name.
        """
        logger.info("=== NetworkProbe.run_all starting ===")
        t0 = time.perf_counter()

        results: Dict[str, Any] = {
            "probe_type": "network",
            "timestamp": time.time(),
        }

        # --- Local probes ---
        for probe_name, probe_fn in [
            ("nics", self.detect_nics),
            ("buffer_sizes", self.probe_buffer_sizes),
            ("irq_distribution", self.probe_irq_distribution),
            ("numa_affinity", self.probe_numa_affinity),
        ]:
            try:
                results[probe_name] = probe_fn()
            except Exception as exc:
                logger.error("%s probe failed: %s", probe_name, exc)
                results[probe_name] = {"error": str(exc)}

        # --- Remote probes ---
        peer_results: List[Dict] = []
        if target_hosts:
            for host in target_hosts:
                peer: Dict[str, Any] = {"host": host}

                # Latency
                try:
                    peer["latency"] = self.probe_latency(host)
                except Exception as exc:
                    logger.warning("Latency probe to %s failed: %s", host, exc)
                    peer["latency"] = {"error": str(exc)}

                # Throughput (only if host is reachable)
                loss = peer.get("latency", {}).get("loss_pct", 100.0)
                if loss < 100.0:
                    try:
                        peer["throughput"] = self.probe_tcp_throughput(host)
                    except Exception as exc:
                        logger.warning(
                            "Throughput probe to %s failed: %s", host, exc
                        )
                        peer["throughput"] = {"error": str(exc)}
                else:
                    peer["throughput"] = {
                        "error": "host unreachable (100% packet loss)"
                    }

                peer_results.append(peer)

        results["peers"] = peer_results

        elapsed = time.perf_counter() - t0
        results["duration_s"] = round(elapsed, 2)
        logger.info(
            "=== NetworkProbe.run_all complete in %.2f s ===", elapsed
        )
        return results
