from __future__ import annotations

import json
import re
import subprocess
from pathlib import Path
from typing import Final

from hytop.net.models import NetCounter, NetKind

VIRTUAL_PREFIXES: Final[tuple[str, ...]] = (
    "veth",
    "docker",
    "cni",
    "flannel",
    "virbr",
    "br-",
    "tunl",
)

# Remote collector script shipped verbatim to remote hosts via SSH.
# Stored as a standalone file so it gets syntax highlighting and linting.
# Keep _remote_collect.py in sync with collect_local_counters() below.
REMOTE_COLLECTOR_PY: Final[str] = (Path(__file__).parent / "_remote_collect.py").read_text()


def parse_kind_filter(kind: str) -> set[NetKind]:
    if kind == "all":
        return {"eth", "ib"}
    return {kind}  # type: ignore[return-value]


def parse_name_filter(raw: str) -> set[str]:
    return {token.strip() for token in raw.split(",") if token.strip()}


def should_keep_iface(
    name: str,
    include: set[str],
) -> bool:
    if include and name not in include:
        return False
    if name == "lo":
        return False
    return not name.startswith(VIRTUAL_PREFIXES)


def _safe_read_int(path: Path) -> int | None:
    try:
        return int(path.read_text().strip())
    except (OSError, ValueError):
        return None


def _safe_read_text(path: Path) -> str | None:
    try:
        return path.read_text().strip()
    except OSError:
        return None


def _is_infiniband_netdev(iface: Path) -> bool:
    return _safe_read_text(iface / "type") == "32"


def _normalize_ib_port_state(raw: str | None) -> str:
    if not raw:
        return "unknown"
    # Example: "4: ACTIVE"
    if ":" in raw:
        _, _, state = raw.partition(":")
        return state.strip().lower()
    return raw.strip().lower()


def _normalize_link_layer(raw: str | None) -> NetKind:
    if not raw:
        return "ib"
    text = raw.strip().lower()
    return "eth" if text.startswith("ethernet") else "ib"


def _should_keep_ib_iface(
    port_name: str,
    mapped_netdev: str | None,
    include: set[str],
) -> bool:
    names = {port_name}
    if mapped_netdev:
        names.add(mapped_netdev)
    if include and not (names & include):
        return False
    if mapped_netdev == "lo":
        return False
    return not (mapped_netdev and mapped_netdev.startswith(VIRTUAL_PREFIXES))


def _parse_ibdev2netdev_output(stdout: str) -> dict[tuple[str, str], tuple[str, str]]:
    """Parse `ibdev2netdev` output into (dev,port)->(netdev,state)."""

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


def _get_ibdev2netdev_mapping() -> dict[tuple[str, str], tuple[str, str]]:
    """Best-effort ibdev2netdev mapping for classification/state labeling."""

    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 collect_local_counters(
    kind_filter: set[NetKind],
    include: set[str],
) -> dict[str, NetCounter]:
    counters: dict[str, NetCounter] = {}
    ib_mapping = _get_ibdev2netdev_mapping()
    ib_netdevs = {netdev for netdev, _state in ib_mapping.values()}

    if "eth" in kind_filter:
        for iface in Path("/sys/class/net").iterdir():
            if not iface.is_dir():
                continue
            name = iface.name
            if name in ib_netdevs:
                continue
            if _is_infiniband_netdev(iface):
                continue
            if not should_keep_iface(name, include):
                continue
            rx = _safe_read_int(iface / "statistics" / "rx_bytes")
            tx = _safe_read_int(iface / "statistics" / "tx_bytes")
            if rx is None or tx is None:
                continue
            counter = NetCounter(
                kind="eth",
                name=name,
                rx_bytes=rx,
                tx_bytes=tx,
                link_state=(_safe_read_text(iface / "operstate") or "unknown").lower(),
            )
            counters[counter.key] = counter

    # 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:
        ib_root = Path("/sys/class/infiniband")
        if ib_root.exists():
            for hca in ib_root.iterdir():
                ports_dir = hca / "ports"
                if not ports_dir.is_dir():
                    continue
                for port in ports_dir.iterdir():
                    name = f"{hca.name}/p{port.name}"
                    mapped = ib_mapping.get((hca.name, port.name))
                    mapped_netdev = mapped[0] if mapped is not None else None
                    link_layer = _safe_read_text(port / "link_layer")
                    mode = _normalize_link_layer(link_layer)
                    if mode not in kind_filter:
                        continue
                    if not _should_keep_ib_iface(
                        port_name=name,
                        mapped_netdev=mapped_netdev,
                        include=include,
                    ):
                        continue
                    iface_name = mapped_netdev or name
                    rx_words = _safe_read_int(port / "counters" / "port_rcv_data")
                    tx_words = _safe_read_int(port / "counters" / "port_xmit_data")
                    if rx_words is None or tx_words is None:
                        continue
                    counter = NetCounter(
                        kind=mode,
                        name=iface_name,
                        rx_bytes=rx_words * 4,
                        tx_bytes=tx_words * 4,
                        link_state=(
                            mapped[1]
                            if mapped is not None
                            else _normalize_ib_port_state(_safe_read_text(port / "state"))
                        ),
                        device_name=name,
                    )
                    counters[counter.key] = counter

    return counters


def parse_counter_payload(payload: str) -> dict[str, NetCounter]:
    try:
        data = json.loads(payload)
    except json.JSONDecodeError as exc:
        raise ValueError(f"invalid collector payload: {exc}") from exc
    raw_counters = data.get("counters")
    if not isinstance(raw_counters, list):
        raise ValueError("invalid collector payload: missing counters list")

    counters: dict[str, NetCounter] = {}
    for item in raw_counters:
        if not isinstance(item, dict):
            continue
        kind = item.get("kind")
        name = item.get("name")
        rx = item.get("rx_bytes")
        tx = item.get("tx_bytes")
        link_state = item.get("link_state")
        device_name = item.get("device_name")
        if kind not in {"eth", "ib"}:
            continue
        if not isinstance(name, str):
            continue
        try:
            rx_value = int(rx)
            tx_value = int(tx)
        except (TypeError, ValueError):
            continue
        normalized_state: str | None
        if isinstance(link_state, str):
            normalized_state = (
                _normalize_ib_port_state(link_state) if kind == "ib" else link_state.strip().lower()
            )
        else:
            normalized_state = None
        counter = NetCounter(
            kind=kind,
            name=name,
            rx_bytes=rx_value,
            tx_bytes=tx_value,
            link_state=normalized_state,
            device_name=device_name if isinstance(device_name, str) else None,
        )
        counters[counter.key] = counter
    return counters
