# This script is shipped verbatim to remote hosts via SSH (python3 -c).
# It must be fully self-contained: no imports from hytop.* are allowed.
#
# IMPORTANT: The caller prepends the following line before running this script:
#   INPUT_JSON='{"kind_filter": [...], "include": [...]}'
#
# Keep the collection logic in sync with collect_local_counters() in collector.py.

import json
import re
import subprocess
from pathlib import Path

vprefix = ("veth", "docker", "cni", "flannel", "virbr", "br-", "tunl")


def read_int(path):
    try:
        return int(path.read_text().strip())
    except (OSError, ValueError):
        return None


def read_text(path):
    try:
        return path.read_text().strip()
    except OSError:
        return None


def normalize_link_layer(value):
    if not value:
        return "ib"
    text = value.strip().lower()
    return "eth" if text.startswith("ethernet") else "ib"


def parse_ibdev2netdev_output(stdout):
    mapping = {}
    pattern = re.compile(
        r"^(?P<dev>\S+)\s+port\s+(?P<port>\d+)\s+==>\s+(?P<netdev>\S+)\s+\((?P<state>[^)]+)\)$"
    )
    for line in stdout.splitlines():
        m = pattern.match(line.strip())
        if not m:
            continue
        dev = m.group("dev")
        port = m.group("port")
        netdev = m.group("netdev")
        state = m.group("state").strip().lower()
        mapping[(dev, port)] = {"netdev": netdev, "state": state}
    return mapping


def get_ibdev2netdev_mapping():
    try:
        proc = subprocess.run(
            ["ibdev2netdev"],
            check=False,
            capture_output=True,
            text=True,
            timeout=2.0,
        )
    except (OSError, subprocess.TimeoutExpired):
        return {}
    if proc.returncode != 0:
        return {}
    return parse_ibdev2netdev_output(proc.stdout)


def is_ib_netdev(iface_path):
    type_text = read_text(iface_path / "type")
    return type_text == "32"


def want(name, include):
    if include and name not in include:
        return False
    return not (name == "lo" or name.startswith(vprefix))


def collect(kind_filter, include):
    out = []
    ib_map = get_ibdev2netdev_mapping()
    ib_netdevs = {item["netdev"] for item in ib_map.values()}

    if "eth" in kind_filter:
        for p in Path("/sys/class/net").iterdir():
            if not p.is_dir():
                continue
            name = p.name
            if name in ib_netdevs:
                continue
            if is_ib_netdev(p):
                continue
            if not want(name, include):
                continue
            rx = read_int(p / "statistics" / "rx_bytes")
            tx = read_int(p / "statistics" / "tx_bytes")
            if rx is None or tx is None:
                continue
            out.append(
                {
                    "kind": "eth",
                    "name": name,
                    "rx_bytes": rx,
                    "tx_bytes": tx,
                    "link_state": (read_text(p / "operstate") or "unknown").lower(),
                }
            )

    # Always scan /sys/class/infiniband even for eth-only filters:
    # IB ports operating in RoCE mode report link_layer=Ethernet and are
    # classified as "eth" by normalize_link_layer(). The mode check below
    # filters them out if they don't match kind_filter.
    if "ib" in kind_filter or "eth" in kind_filter:
        base = Path("/sys/class/infiniband")
        if base.exists():
            for dev in base.iterdir():
                ports = dev / "ports"
                if not ports.is_dir():
                    continue
                for port in ports.iterdir():
                    pname = f"{dev.name}/p{port.name}"
                    mapped = ib_map.get((dev.name, port.name))
                    mapped_netdev = mapped["netdev"] if mapped is not None else None
                    mode = normalize_link_layer(read_text(port / "link_layer"))
                    if mode not in kind_filter:
                        continue
                    candidate_name = mapped_netdev or pname
                    if not want(candidate_name, include):
                        continue
                    if include and pname in include:
                        pass
                    elif include and candidate_name not in include:
                        continue
                    rxw = read_int(port / "counters" / "port_rcv_data")
                    txw = read_int(port / "counters" / "port_xmit_data")
                    if rxw is None or txw is None:
                        continue
                    out.append(
                        {
                            "kind": mode,
                            "name": candidate_name,
                            "rx_bytes": rxw * 4,
                            "tx_bytes": txw * 4,
                            "link_state": (
                                mapped["state"]
                                if mapped is not None
                                else (read_text(port / "state") or "unknown")
                            ),
                            "device_name": pname,
                        }
                    )
    return out


payload = json.loads(INPUT_JSON)  # noqa: F821  # injected by caller
res = {
    "counters": collect(
        kind_filter=set(payload.get("kind_filter", ["eth", "ib"])),
        include=set(payload.get("include", [])),
    )
}
print(json.dumps(res, separators=(",", ":")))
