netbox-diode-project/collectors/network_collector.py

#!/usr/bin/env python3
"""Network device collector for NetBox via Diode SDK.

Discovers Cisco and Brocade network devices via NAPALM (and optionally pyATS
for CDP/OSPF/IS-IS), then ingests devices, interfaces, IPs, cables, VLANs,
VRFs, prefixes, configs, and inventory into NetBox through the Diode pipeline.

BGP sessions are pushed to the netbox-bgp plugin API (not via Diode).

Usage:
    python collectors/network_collector.py --inventory inventory.yaml
    python collectors/network_collector.py --inventory inventory.yaml --dry-run
"""

import argparse
import json
import logging
import os
import re
import sys
from typing import Any

import yaml
from napalm import get_network_driver

from netboxlabs.diode.sdk import DiodeClient, DiodeDryRunClient
from netboxlabs.diode.sdk.ingester import (
    ASN,
    Cable,
    Device,
    DeviceConfig,
    DeviceRole,
    DeviceType,
    Entity,
    GenericObject,
    Interface,
    InventoryItem,
    InventoryItemRole,
    IPAddress,
    Manufacturer,
    Platform,
    Prefix,
    Site,
    VRF,
    VLAN,
    VLANGroup,
)

# Optional: pyATS/Genie for Cisco-specific parsing (CDP, OSPF, IS-IS)
try:
    from genie.conf.base import Device as GenieDevice
    from genie.libs.parser.utils import common as genie_common

    HAS_PYATS = True
except ImportError:
    HAS_PYATS = False

# Optional: requests for pushing data to NetBox plugin APIs
try:
    import requests

    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False

# Netmiko for devices without NAPALM drivers (e.g. Brocade VDX/NOS)
from netmiko import ConnectHandler

log = logging.getLogger("network-collector")

# Drivers that should use Netmiko SSH fallback instead of NAPALM.
# napalm-ruckus-fastiron installs but fails on open(); NOS has no NAPALM driver.
# Keys = inventory driver names, values = Netmiko device_type strings.
NETMIKO_ONLY_DRIVERS = {
    "ruckus_fastiron": "ruckus_fastiron",
    "nos": "extreme_nos",
    "brocade_nos": "extreme_nos",
    "brocade_vdx": "extreme_nos",
    "extreme_nos": "extreme_nos",
}

# ---------------------------------------------------------------------------
# Interface type mapping
# ---------------------------------------------------------------------------

SPEED_TO_TYPE = {
    10: "10base-t",
    100: "100base-tx",
    1000: "1000base-t",
    2500: "2.5gbase-t",
    5000: "5gbase-t",
    10000: "10gbase-x-sfpp",
    25000: "25gbase-x-sfp28",
    40000: "40gbase-x-qsfpp",
    50000: "50gbase-x-sfp56",
    100000: "100gbase-x-qsfp28",
}

NAME_TO_TYPE = {
    r"^(Gi|GigabitEthernet)": "1000base-t",
    r"^(Te|TenGigabitEthernet|TenGigE)": "10gbase-x-sfpp",
    r"^(Tw|TwentyFiveGig)": "25gbase-x-sfp28",
    r"^(Fo|FortyGig|FortyGigE)": "40gbase-x-qsfpp",
    r"^(Hu|HundredGig|HundredGigE)": "100gbase-x-qsfp28",
    r"^(Fa|FastEthernet)": "100base-tx",
    r"^(Et|Ethernet)\d": "1000base-t",
    r"^(Lo|Loopback)": "virtual",
    r"^(Vl|Vlan)": "virtual",
    r"^(Tu|Tunnel)": "virtual",
    r"^(Mg|MgmtEth|Management)": "1000base-t",
    r"^(Nu|Null)": "virtual",
    r"^(Po|Port-channel|port-channel)": "lag",
    r"^(BV|BVI)": "bridge",
    r"^(Se|Serial)": "other",
    # Brocade ICX / FastIron patterns
    r"^e\s?\d": "1000base-t",                          # e 1/1/1
    r"^ethernet\s?\d": "1000base-t",                    # ethernet 1/1/1
    r"^(ve|VirtualEthernet)": "virtual",                # ve 10
    r"^(lb|loopback)\s?\d": "virtual",                  # lb 1
    r"^(lag|trunk)\s?\d": "lag",                         # lag 1
    r"^management\s?\d": "1000base-t",                  # management 1
}

DRIVER_TO_PLATFORM = {
    "ios": "Cisco IOS",
    "iosxr": "Cisco IOS-XR",
    "eos": "Arista EOS",
    "junos": "Juniper Junos",
    "nxos": "Cisco NX-OS",
    "nxos_ssh": "Cisco NX-OS",
    "ruckus_fastiron": "Ruckus FastIron",
    "nos": "Brocade NOS",
    "brocade_nos": "Brocade NOS",
    "brocade_vdx": "Brocade NOS",
    "extreme_nos": "Brocade NOS",
}

DRIVER_TO_MANUFACTURER = {
    "ios": "Cisco",
    "iosxr": "Cisco",
    "eos": "Arista",
    "junos": "Juniper",
    "nxos": "Cisco",
    "nxos_ssh": "Cisco",
    "ruckus_fastiron": "Brocade",
    "nos": "Brocade",
    "brocade_nos": "Brocade",
    "brocade_vdx": "Brocade",
    "extreme_nos": "Brocade",
}


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------


def load_dotenv(path: str = ".env") -> None:
    """Minimal .env loader — no extra dependency."""
    if not os.path.isfile(path):
        return
    with open(path) as fh:
        for line in fh:
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, _, val = line.partition("=")
            key = key.strip()
            val = val.strip().strip("\"'")
            os.environ.setdefault(key, val)


def load_inventory(path: str) -> dict:
    """Load the device inventory from a YAML file.

    Expected format:
        defaults:
          site: main
          role: Network Device
          username: admin
          password: cisco
          secret: cisco          # enable secret
          driver: ios            # NAPALM driver
          timeout: 60

        devices:
          - host: 10.10.20.1
            driver: ios
            role: Router
          - host: 10.10.20.55
            driver: ios
            role: Switch
          - host: 10.10.20.100
            driver: iosxr
            role: Router
    """
    with open(path) as fh:
        data = yaml.safe_load(fh)
    if not data or "devices" not in data:
        raise ValueError(f"Inventory file {path} must have a 'devices' list")
    return data


def merge_device_config(device_entry: dict, defaults: dict) -> dict:
    """Merge a single device entry with defaults (device overrides defaults)."""
    merged = dict(defaults)
    merged.update({k: v for k, v in device_entry.items() if v is not None})
    return merged


# ---------------------------------------------------------------------------
# Interface type helpers
# ---------------------------------------------------------------------------


def map_interface_type(name: str, speed: int = 0) -> str:
    """Map an interface name and/or speed to a NetBox interface type."""
    for pattern, iface_type in NAME_TO_TYPE.items():
        if re.match(pattern, name, re.IGNORECASE):
            return iface_type
    if speed and speed in SPEED_TO_TYPE:
        return SPEED_TO_TYPE[speed]
    return "other"


def normalize_interface_name(name: str) -> str:
    """Normalize short interface names from LLDP/CDP to long form for matching."""
    abbrevs = {
        "Gi": "GigabitEthernet",
        "Te": "TenGigabitEthernet",
        "Fa": "FastEthernet",
        "Et": "Ethernet",
        "Fo": "FortyGigabitEthernet",
        "Hu": "HundredGigE",
        "Lo": "Loopback",
        "Vl": "Vlan",
        "Po": "Port-channel",
        "Mg": "MgmtEth",
        "Tu": "Tunnel",
        "Se": "Serial",
    }
    for short, long in abbrevs.items():
        if name.startswith(short) and not name.startswith(long):
            rest = name[len(short):]
            # Handle both "Fo1/0/1" and "Fo 1/0/1" (NOS uses spaces)
            rest = rest.lstrip()
            if rest and (rest[0].isdigit() or rest[0] == "/"):
                return long + rest
    return name


# ---------------------------------------------------------------------------
# NAPALM data collection
# ---------------------------------------------------------------------------


def connect_device(host: str, driver: str, username: str, password: str,
                   secret: str = "", timeout: int = 60,
                   optional_args: dict | None = None) -> Any:
    """Connect to a device via NAPALM and return the open driver."""
    driver_cls = get_network_driver(driver)
    opts = optional_args or {}
    if secret:
        opts.setdefault("secret", secret)
    opts.setdefault("transport", "ssh")
    dev = driver_cls(host, username, password, timeout=timeout, optional_args=opts)
    dev.open()
    return dev


def collect_napalm_data(dev: Any) -> dict:
    """Collect all available data from a NAPALM device."""
    data = {}

    # Facts (always available)
    try:
        data["facts"] = dev.get_facts()
        log.info("  Facts: %s (%s)", data["facts"].get("hostname"), data["facts"].get("model"))
    except Exception as exc:
        log.error("  get_facts() failed: %s", exc)
        return data

    # Interfaces
    try:
        data["interfaces"] = dev.get_interfaces()
        log.debug("  Interfaces: %d", len(data["interfaces"]))
    except Exception as exc:
        log.warning("  get_interfaces() failed: %s", exc)

    # Interface IPs (IPv4 + IPv6)
    try:
        data["interfaces_ip"] = dev.get_interfaces_ip()
        log.debug("  Interface IPs: %d interfaces with IPs", len(data["interfaces_ip"]))
    except Exception as exc:
        log.warning("  get_interfaces_ip() failed: %s", exc)

    # LLDP neighbors
    try:
        data["lldp_neighbors"] = dev.get_lldp_neighbors_detail()
        total = sum(len(v) for v in data["lldp_neighbors"].values())
        log.debug("  LLDP neighbors: %d", total)
    except Exception as exc:
        log.debug("  get_lldp_neighbors_detail() failed: %s", exc)

    # Config
    try:
        data["config"] = dev.get_config()
        log.debug("  Config: running=%d bytes", len(data["config"].get("running", "")))
    except Exception as exc:
        log.warning("  get_config() failed: %s", exc)

    # VLANs
    try:
        data["vlans"] = dev.get_vlans()
        log.debug("  VLANs: %d", len(data["vlans"]))
    except Exception as exc:
        log.debug("  get_vlans() unavailable: %s", exc)

    # Network instances (VRFs)
    try:
        data["network_instances"] = dev.get_network_instances()
        log.debug("  VRFs: %d", len(data["network_instances"]))
    except Exception as exc:
        log.debug("  get_network_instances() unavailable: %s", exc)

    # BGP neighbors
    try:
        data["bgp_neighbors"] = dev.get_bgp_neighbors_detail()
        total = sum(len(peers) for vrf in data["bgp_neighbors"].values() for peers in vrf.values())
        log.debug("  BGP peers: %d", total)
    except Exception as exc:
        log.debug("  get_bgp_neighbors_detail() unavailable: %s", exc)

    return data


# ---------------------------------------------------------------------------
# pyATS/Genie data collection (optional)
# ---------------------------------------------------------------------------


def collect_pyats_data(host: str, driver: str, username: str, password: str,
                       secret: str = "") -> dict:
    """Collect CDP, OSPF, IS-IS data via pyATS/Genie parsers."""
    if not HAS_PYATS:
        return {}

    os_map = {"ios": "iosxe", "iosxr": "iosxr", "nxos": "nxos", "nxos_ssh": "nxos"}
    genie_os = os_map.get(driver)
    if not genie_os:
        log.debug("  pyATS: no mapping for driver %s, skipping", driver)
        return {}

    data = {}
    try:
        dev = GenieDevice(name=host, os=genie_os, credentials={
            "default": {"username": username, "password": password}
        })
        dev.connect(ip=host, init_exec_commands=[], init_config_commands=[], log_stdout=False)
    except Exception as exc:
        log.warning("  pyATS connect failed: %s", exc)
        return data

    # CDP neighbors
    try:
        data["cdp_neighbors"] = dev.parse("show cdp neighbors detail")
        log.debug("  CDP neighbors: %d", len(data["cdp_neighbors"].get("index", {})))
    except Exception as exc:
        log.debug("  CDP parse failed: %s", exc)

    # OSPF neighbors
    try:
        data["ospf_neighbors"] = dev.parse("show ip ospf neighbor")
        log.debug("  OSPF neighbors parsed")
    except Exception as exc:
        log.debug("  OSPF parse unavailable: %s", exc)

    # IS-IS adjacencies
    try:
        data["isis_adjacencies"] = dev.parse("show isis adjacency")
        log.debug("  IS-IS adjacencies parsed")
    except Exception as exc:
        log.debug("  IS-IS parse unavailable: %s", exc)

    # Inventory (modules, transceivers)
    try:
        data["inventory"] = dev.parse("show inventory")
        log.debug("  Inventory parsed")
    except Exception as exc:
        log.debug("  Inventory parse unavailable: %s", exc)

    try:
        dev.disconnect()
    except Exception:
        pass

    return data


# ---------------------------------------------------------------------------
# Netmiko fallback collection (devices without NAPALM drivers)
# ---------------------------------------------------------------------------


def connect_netmiko(host: str, driver: str, username: str, password: str,
                    secret: str = "", timeout: int = 60) -> ConnectHandler:
    """Connect to a device via Netmiko and return the connection handler."""
    device_type = NETMIKO_ONLY_DRIVERS.get(driver, driver)
    params = {
        "device_type": device_type,
        "host": host,
        "username": username,
        "password": password,
        "timeout": timeout,
        "conn_timeout": timeout,
    }
    if secret:
        params["secret"] = secret
    conn = ConnectHandler(**params)
    return conn


def collect_netmiko_data(conn: ConnectHandler, driver: str) -> dict:
    """Collect device data via Netmiko CLI commands and parse output.

    Returns a dict shaped like collect_napalm_data() output so the same
    entity builders work for both paths.
    """
    data = {}

    # --- Facts ---
    facts = _netmiko_parse_facts(conn, driver)
    if not facts:
        return data
    data["facts"] = facts
    log.info("  Facts: %s (%s)", facts.get("hostname"), facts.get("model"))

    # --- Interfaces ---
    try:
        data["interfaces"] = _netmiko_parse_interfaces(conn, driver)
        log.debug("  Interfaces: %d", len(data["interfaces"]))
    except Exception as exc:
        log.warning("  Interface collection failed: %s", exc)

    # --- Interface IPs ---
    try:
        data["interfaces_ip"] = _netmiko_parse_interfaces_ip(conn, driver)
        log.debug("  Interface IPs: %d interfaces with IPs", len(data["interfaces_ip"]))
    except Exception as exc:
        log.warning("  IP collection failed: %s", exc)

    # --- LLDP neighbors ---
    try:
        data["lldp_neighbors"] = _netmiko_parse_lldp(conn, driver)
        total = sum(len(v) for v in data["lldp_neighbors"].values())
        log.debug("  LLDP neighbors: %d", total)
    except Exception as exc:
        log.debug("  LLDP collection failed: %s", exc)

    # --- Running config ---
    try:
        running = conn.send_command("show running-config", read_timeout=120)
        data["config"] = {
            "running": running,
            "startup": "",
            "candidate": "",
        }
        log.debug("  Config: running=%d bytes", len(running))
    except Exception as exc:
        log.warning("  Config collection failed: %s", exc)

    return data


def _netmiko_parse_facts(conn: ConnectHandler, driver: str) -> dict:
    """Parse device facts from Netmiko CLI output.

    Handles both ICX/FastIron and NOS/VDX output formats.
    """
    facts = {
        "hostname": "",
        "model": "Unknown",
        "vendor": "Brocade",
        "serial_number": "",
        "os_version": "",
        "fqdn": "",
        "uptime": -1,
        "interface_list": [],
    }

    try:
        # Hostname from prompt (works for both ICX and NOS)
        # ICX: "SSH@Brocade40G-01#" → "Brocade40G-01"
        # NOS: "Brocade-VDX-6940-01#" → "Brocade-VDX-6940-01"
        prompt = conn.find_prompt().strip().rstrip("#>")
        # ICX prepends "SSH@" to the prompt
        if prompt.startswith("SSH@"):
            prompt = prompt[4:]
        if prompt:
            facts["hostname"] = prompt

        # 'show version' — works on both ICX and NOS (different formats)
        ver_out = conn.send_command("show version")
        for line in ver_out.splitlines():
            stripped = line.strip()
            low = stripped.lower()

            # ICX: "  HW: Stackable ICX6610-24"
            if low.startswith("hw:"):
                model = stripped.split(":", 1)[1].strip()
                # Remove "Stackable " prefix
                model = re.sub(r"^Stackable\s+", "", model)
                facts["model"] = model

            # ICX: "  Serial  #: 2ax5o2jk68e"
            elif "serial" in low and "#:" in low:
                sn = stripped.split("#:", 1)[1].strip()
                if sn:
                    facts["serial_number"] = sn

            # ICX: "  SW: Version 08.0.30uT7f3"
            elif low.startswith("sw:") and "version" in low:
                ver = stripped.split("Version", 1)[-1].strip() if "Version" in stripped else stripped.split(":", 1)[1].strip()
                facts["os_version"] = ver

            # NOS: "Firmware name:  ..." or "Network Operating System ..."
            elif "firmware name:" in low or "network operating system" in low:
                facts["os_version"] = stripped.split(":", 1)[-1].strip() if ":" in stripped else stripped

            # NOS: "Firmware version:  ..."
            elif "firmware version:" in low:
                facts["os_version"] = stripped.split(":", 1)[1].strip()

        # Try 'show chassis' for NOS (ICX may not have this or it returns different data)
        try:
            chassis_out = conn.send_command("show chassis")
            for line in chassis_out.splitlines():
                low = line.strip().lower()
                # NOS: "Chassis Name:  VDX6940-36Q" or "Switch Type: ..."
                if ("chassis name:" in low or "switchtype:" in low
                        or "switch type:" in low):
                    val = line.split(":", 1)[1].strip()
                    if val and val.lower() not in ("", "unknown"):
                        facts["model"] = val
                # NOS: "Chassis Serial Number: ..." or "Serial Number: ..."
                elif "serial number:" in low and not facts["serial_number"]:
                    sn = line.split(":", 1)[1].strip()
                    if sn:
                        facts["serial_number"] = sn
        except Exception:
            pass  # ICX may error on 'show chassis'

        # Fallback: get model from running config's chassis-name
        if facts["model"] in ("Unknown", "") or facts["model"].isdigit():
            try:
                cfg_out = conn.send_command("show running-config", read_timeout=120)
                for line in cfg_out.splitlines():
                    stripped = line.strip()
                    if stripped.lower().startswith("chassis-name "):
                        val = stripped.split(None, 1)[1].strip()
                        if val:
                            facts["model"] = val
                            break
            except Exception:
                pass

        # Try to get hostname from running config (ICX uses 'hostname', NOS uses 'switch-name')
        try:
            host_out = conn.send_command("show running-config | include hostname|switch-name|host-name")
            for line in host_out.splitlines():
                stripped = line.strip()
                low = stripped.lower()
                if (low.startswith("hostname ") or "switch-name" in low
                        or "host-name" in low):
                    parts = stripped.split()
                    if len(parts) >= 2:
                        facts["hostname"] = parts[-1].strip('"')
        except Exception:
            pass

        # Build interface list from 'show interface brief'
        try:
            iface_out = conn.send_command("show interface brief")
            if "syntax error" not in iface_out.lower():
                for line in iface_out.splitlines():
                    parts = line.split()
                    if not parts:
                        continue
                    # ICX: ports like "1/1/1", "1/2/1", "mgmt1"
                    # NOS: ports like "Te 0/0", "Gi 0/1", "Ve 10"
                    if re.match(r"^\d+/\d+/\d+", parts[0]):
                        facts["interface_list"].append(parts[0])
                    elif parts[0].startswith(("Te", "Gi", "Fo", "Hu", "Eth", "mgmt",
                                              "Ve", "Lo", "Po", "te", "gi", "fo")):
                        facts["interface_list"].append(parts[0])
        except Exception:
            pass

    except Exception as exc:
        log.error("  Facts parsing failed: %s", exc)
        return {}

    return facts


def _netmiko_parse_interfaces(conn: ConnectHandler, driver: str) -> dict:
    """Parse interface data from CLI output.

    Handles both ICX/FastIron and NOS/VDX formats.
    Returns dict matching NAPALM get_interfaces() format:
    {iface_name: {is_up, is_enabled, description, mac_address, speed, mtu, last_flapped}}
    """
    interfaces = {}
    is_icx = driver in ("ruckus_fastiron",)

    # --- Try brief output first (ICX only; NOS errors on this) ---
    brief_output = ""
    try:
        brief_output = conn.send_command("show interfaces brief")
        if "syntax error" in brief_output.lower() or "invalid" in brief_output.lower():
            brief_output = ""
    except Exception:
        pass
    if not brief_output:
        try:
            brief_output = conn.send_command("show interface brief")
            if "syntax error" in brief_output.lower() or "invalid" in brief_output.lower():
                brief_output = ""
        except Exception:
            pass

    if brief_output and is_icx:
        for line in brief_output.splitlines():
            parts = line.split()
            if not parts:
                continue
            # ICX format:
            # Port       Link    State   Dupl Speed Trunk Tag Pvid Pri MAC             Name
            # 1/1/1      Up      Forward Full 1G    None  No  20   0   cc4e.2413.2562
            if not re.match(r"^\d+/\d+/\d+|^mgmt\d*|^lb\d*|^ve\d*", parts[0]):
                continue
            name = parts[0]
            is_up = len(parts) > 1 and parts[1].lower() == "up"
            speed = 0
            mac = ""
            for part in parts:
                # Speed: "1G", "10G", "40G", "100G", "100M"
                speed_m = re.match(r"^(\d+)([GM])$", part)
                if speed_m:
                    val = int(speed_m.group(1))
                    speed = val * 1000 if speed_m.group(2) == "G" else val  # → Mbps
                # MAC: cc4e.2413.2562
                if re.match(r"^[0-9a-f]{4}\.[0-9a-f]{4}\.[0-9a-f]{4}$", part, re.IGNORECASE):
                    mac = part
            interfaces[name] = {
                "is_up": is_up,
                "is_enabled": True,
                "description": "",
                "mac_address": mac,
                "speed": speed,
                "mtu": 0,
                "last_flapped": -1.0,
            }

    # --- Parse full 'show interface' detail output ---
    # For NOS/VDX this is the primary source (brief doesn't exist).
    # For ICX this enriches the brief data with descriptions, MACs, MTUs.
    try:
        detail_out = conn.send_command("show interface", read_timeout=120)
        current_iface = None
        for line in detail_out.splitlines():
            # Header line pattern (both ICX and NOS):
            # NOS: "FortyGigabitEthernet 1/0/1 is up, line protocol is up (connected)"
            # NOS: "TenGigabitEthernet 1/0/34:1 is up, line protocol is up"
            # NOS: "Management 1/0 is up, line protocol is up"
            # ICX: "GigabitEthernet1/1/1 is up, line protocol is up"
            iface_hdr = re.match(
                r"^((?:TenGigabit|Gigabit|FortyGigabit|HundredGigabit)?Ethernet|"
                r"Management|Loopback|Port-channel|Ve)\s*(\S+)\s+is\s+(up|down|administratively down)",
                line,
            )
            if iface_hdr:
                iface_type_str = iface_hdr.group(1)
                iface_id = iface_hdr.group(2)
                link_state = iface_hdr.group(3)
                full_name = f"{iface_type_str} {iface_id}"
                # Abbreviate NOS long names for consistency with LLDP
                abbrev_map = {
                    "FortyGigabitEthernet": "Fo",
                    "TenGigabitEthernet": "Te",
                    "GigabitEthernet": "Gi",
                    "HundredGigabitEthernet": "Hu",
                    "Management": "Management",
                    "Loopback": "Lo",
                    "Port-channel": "Po",
                    "Ve": "Ve",
                }
                abbrev = abbrev_map.get(iface_type_str, iface_type_str)
                short_name = f"{abbrev} {iface_id}"

                if is_icx:
                    # For ICX, match to existing brief entry
                    current_iface = None
                    for iname in interfaces:
                        if full_name.endswith(iname) or iname == full_name:
                            current_iface = iname
                            break
                else:
                    # NOS: create new interface entry using abbreviated name
                    current_iface = short_name
                    is_up = link_state == "up"
                    # Check "line protocol is up/down" portion
                    proto_up = "line protocol is up" in line.lower()
                    if current_iface not in interfaces:
                        interfaces[current_iface] = {
                            "is_up": proto_up,
                            "is_enabled": link_state != "administratively down",
                            "description": "",
                            "mac_address": "",
                            "speed": 0,
                            "mtu": 0,
                            "last_flapped": -1.0,
                        }
                continue

            if current_iface and current_iface in interfaces:
                stripped = line.strip()
                low = stripped.lower()
                if "port name is" in low or low.startswith("description:"):
                    if low.startswith("description:"):
                        desc = stripped.split(":", 1)[-1].strip()
                    else:
                        desc = stripped.split("is", 1)[-1].strip()
                    interfaces[current_iface]["description"] = desc[:200]
                elif "hardware is" in low and "address is" in low:
                    mac_match = re.search(r"address is\s+([0-9a-fA-F.:]+)", stripped)
                    if mac_match:
                        interfaces[current_iface]["mac_address"] = mac_match.group(1)
                elif low.startswith("mtu") or "mtu " in low:
                    mtu_match = re.search(r"mtu\s+(\d+)", stripped, re.IGNORECASE)
                    if mtu_match:
                        interfaces[current_iface]["mtu"] = int(mtu_match.group(1))
                elif "linespeed actual" in low:
                    # NOS: "LineSpeed Actual     : 40000 Mbit"
                    speed_match = re.search(r":\s*(\d+)\s*Mbit", stripped, re.IGNORECASE)
                    if speed_match:
                        interfaces[current_iface]["speed"] = int(speed_match.group(1))
    except Exception as exc:
        log.debug("  Detailed interface parsing failed: %s", exc)

    return interfaces


def _netmiko_parse_interfaces_ip(conn: ConnectHandler, driver: str) -> dict:
    """Parse IP addresses from CLI output.

    Handles both ICX and NOS formats.
    Returns dict matching NAPALM get_interfaces_ip() format:
    {iface_name: {ipv4: {addr: {prefix_length: N}}, ipv6: {addr: {prefix_length: N}}}}
    """
    interfaces_ip: dict[str, dict] = {}

    # 'show ip interface' — works on both ICX and NOS
    # ICX: "Ve 1                192.168.1.250   YES  NVRAM     up    up    default-vrf"
    # NOS: "Interface  IP-Address  Status  Protocol"
    try:
        output = conn.send_command("show ip interface")
        for line in output.splitlines():
            # Match lines with an IP address
            ip_match = re.search(r"(\d+\.\d+\.\d+\.\d+)(?:/(\d+))?", line)
            if not ip_match:
                continue
            # Skip header lines and lines that look like headers
            if "IP-Address" in line or "Interface" in line.split()[0:1]:
                continue
            ip = ip_match.group(1)
            if ip.startswith("0.") or ip == "127.0.0.1":
                continue
            prefix = int(ip_match.group(2)) if ip_match.group(2) else None

            # Extract interface name — everything before the IP
            before_ip = line[:ip_match.start()].strip()
            # Interface names may have spaces: "Ve 1", "Ve 20", "mgmt 1"
            iface_name = before_ip.strip()
            if not iface_name:
                continue

            # If prefix not in output, try to get it from running config later
            if prefix is None:
                prefix = 24  # default fallback, will try to refine below

            interfaces_ip.setdefault(iface_name, {"ipv4": {}, "ipv6": {}})
            interfaces_ip[iface_name]["ipv4"][ip] = {"prefix_length": prefix}
    except Exception as exc:
        log.debug("  'show ip interface' failed: %s", exc)

    # Also try running-config to get accurate prefix lengths.
    # NOS has nested interfaces under "rbridge-id 1" so pipe-filtered
    # includes won't catch them. Parse the full running config instead.
    try:
        output = conn.send_command("show running-config", read_timeout=120)
        current_iface = None
        for line in output.splitlines():
            stripped = line.strip()
            # Match "interface Ve 40", "interface Management 1/0", etc.
            iface_match = re.match(r"^interface\s+(\S+(?:\s+\S+)?)", stripped, re.IGNORECASE)
            if iface_match:
                current_iface = iface_match.group(1).strip()
                continue
            # Reset on non-indented lines that aren't "interface"
            if not line.startswith(" ") and not line.startswith("\t") and stripped and not stripped.startswith("!"):
                if not stripped.lower().startswith("interface"):
                    current_iface = None
                continue
            if current_iface and "ip address" in stripped.lower():
                ip_match = re.search(
                    r"ip address\s+(\d+\.\d+\.\d+\.\d+)[/\s]+(\d+\.[\d.]+|\d+)", stripped,
                    re.IGNORECASE,
                )
                if ip_match:
                    ip = ip_match.group(1)
                    mask_or_prefix = ip_match.group(2)
                    if "." in mask_or_prefix:
                        prefix = sum(bin(int(x)).count("1") for x in mask_or_prefix.split("."))
                    else:
                        prefix = int(mask_or_prefix)
                    # Normalize interface name: "Ve 40" vs "ve 40"
                    normalized = current_iface
                    for existing_name in interfaces_ip:
                        if existing_name.lower().replace(" ", "") == normalized.lower().replace(" ", ""):
                            normalized = existing_name
                            break
                    interfaces_ip.setdefault(normalized, {"ipv4": {}, "ipv6": {}})
                    interfaces_ip[normalized]["ipv4"][ip] = {"prefix_length": prefix}
    except Exception as exc:
        log.debug("  Running-config IP parsing failed: %s", exc)

    # IPv6
    try:
        output = conn.send_command("show ipv6 interface brief")
        current_iface = None
        for line in output.splitlines():
            parts = line.split()
            if not parts:
                continue
            # Interface line (not indented, not an address)
            if not line.startswith(" ") and not re.match(r"^[0-9a-f]{4}:", parts[0], re.IGNORECASE):
                iface_match = re.match(r"^(\S+(?:\s+\d+)?)", line)
                if iface_match:
                    current_iface = iface_match.group(1).strip()
            # IPv6 address line
            ipv6_match = re.search(r"([0-9a-fA-F:]+(?:::[0-9a-fA-F:]*)?)/(\d+)", line)
            if ipv6_match and current_iface:
                ipv6 = ipv6_match.group(1)
                prefix = int(ipv6_match.group(2))
                if not ipv6.lower().startswith("fe80"):
                    interfaces_ip.setdefault(current_iface, {"ipv4": {}, "ipv6": {}})
                    interfaces_ip[current_iface]["ipv6"][ipv6] = {"prefix_length": prefix}
    except Exception as exc:
        log.debug("  IPv6 interface parsing failed: %s", exc)

    return interfaces_ip


def _netmiko_parse_lldp(conn: ConnectHandler, driver: str) -> dict:
    """Parse LLDP neighbor data from CLI output.

    Handles both ICX and NOS LLDP output formats.
    Returns dict matching NAPALM get_lldp_neighbors_detail() format:
    {local_iface: [{"parent_interface": "", "remote_system_name": "...", ...}]}
    """
    lldp: dict[str, list] = {}

    try:
        output = conn.send_command("show lldp neighbors detail")
    except Exception:
        try:
            output = conn.send_command("show lldp neighbors")
        except Exception:
            return lldp

    if not output:
        return lldp

    # Both ICX and NOS use similar LLDP detail format:
    # ICX: "Local port: 1/1/1" / NOS: "Local Interface: TenGigabitEthernet 0/0"
    # Both: "+ Chassis ID (MAC address): xxxx.xxxx.xxxx"
    # Both: "+ Port ID (interface name): ..."
    # Both: "+ System name         : ..."
    # Entries separated by "Local port:" / "Local Interface:" headers
    current_local = ""
    entry: dict[str, str] = {}

    def _flush_entry():
        nonlocal entry
        if entry and current_local:
            lldp.setdefault(current_local, []).append({
                "parent_interface": current_local,
                "remote_system_name": entry.get("system_name", ""),
                "remote_port": entry.get("port_id", "").strip('"'),
                "remote_port_description": entry.get("port_description", "").strip('"'),
                "remote_chassis_id": entry.get("chassis_id", ""),
                "remote_system_description": entry.get("system_description", ""),
                "remote_system_capab": [],
                "remote_system_enable_capab": [],
            })
        entry = {}

    for line in output.splitlines():
        stripped = line.strip()
        low = stripped.lower()

        # New neighbor block — ICX: "Local port: 1/1/1"
        #                       NOS: "Local Interface: Fo 1/0/1  (Local Interface MAC: ...)"
        #                       NOS: "Neighbors for Interface Fo 1/0/1" (section header)
        if low.startswith("neighbors for interface"):
            _flush_entry()
            current_local = stripped.split("Interface", 1)[1].strip()
            continue

        if "local port:" in low or "local interface:" in low:
            _flush_entry()
            raw = stripped.split(":", 1)[1].strip()
            # NOS appends "(Local Interface MAC: xxxx.xxxx.xxxx)" — strip it
            paren_idx = raw.find("(")
            if paren_idx > 0:
                raw = raw[:paren_idx].strip()
            current_local = raw
            continue

        if not stripped:
            continue

        # Skip section headers
        if stripped in ("MANDATORY TLVs", "OPTIONAL TLVs", "DCBX TLVs") or stripped.startswith("==="):
            continue

        # ICX uses "+ Field : value" format; NOS uses "Field: value"
        # Normalize: strip leading "+" and whitespace
        clean = stripped.lstrip("+ ").strip()
        clean_low = clean.lower()

        if clean_low.startswith("chassis id"):
            # "+ Chassis ID (MAC address): 0013.c602.6961"
            val = clean.split(":", 1)[-1].strip() if ":" in clean else ""
            if val:
                entry["chassis_id"] = val
        elif clean_low.startswith("port id"):
            val = clean.split(":", 1)[-1].strip() if ":" in clean else ""
            if val:
                entry["port_id"] = val
        elif clean_low.startswith("remote interface"):
            # NOS: "Remote Interface: 748e.f8e9.d41b (Remote Interface MAC: ...)"
            # The value is a MAC or port name
            raw = clean.split(":", 1)[-1].strip() if ":" in clean else ""
            paren_idx = raw.find("(")
            if paren_idx > 0:
                raw = raw[:paren_idx].strip()
            if raw and "remote_port_id" not in entry:
                entry["remote_port_id"] = raw
        elif clean_low.startswith("port interface description"):
            # NOS: "Port Interface Description: 40GigabitEthernet1/2/1"
            val = clean.split(":", 1)[-1].strip() if ":" in clean else ""
            if val:
                entry["port_description"] = val.strip('"')
                # On NOS, this is often the best port identifier
                if not entry.get("port_id"):
                    entry["port_id"] = val.strip('"')
        elif clean_low.startswith("system name"):
            val = clean.split(":", 1)[-1].strip() if ":" in clean else ""
            if val:
                entry["system_name"] = val.strip('"')
        elif clean_low.startswith("system description"):
            val = clean.split(":", 1)[-1].strip() if ":" in clean else ""
            if val:
                entry["system_description"] = val.strip('"')
        elif clean_low.startswith("port description"):
            val = clean.split(":", 1)[-1].strip() if ":" in clean else ""
            if val:
                entry["port_description"] = val.strip('"')

    # Flush final entry
    _flush_entry()

    return lldp


# ---------------------------------------------------------------------------
# Device reference helper (rich references for Diode reconciler)
# ---------------------------------------------------------------------------


def _device_ref(hostname: str, model: str, manufacturer: str, role: str,
                site_name: str) -> Device:
    """Build a rich Device reference with enough context for the reconciler."""
    return Device(
        name=hostname,
        device_type=DeviceType(
            model=model or "Unknown",
            manufacturer=Manufacturer(name=manufacturer or "Unknown"),
        ),
        role=DeviceRole(name=role),
        site=Site(name=site_name),
    )


# ---------------------------------------------------------------------------
# Entity builders
# ---------------------------------------------------------------------------


def build_device_entity(facts: dict, driver: str, role: str, site_name: str,
                        host: str) -> Entity:
    """Build a Device entity from NAPALM facts."""
    hostname = facts.get("hostname") or host
    model = facts.get("model") or "Unknown"
    vendor = facts.get("vendor") or DRIVER_TO_MANUFACTURER.get(driver, "Unknown")
    serial = facts.get("serial_number") or ""
    os_version = facts.get("os_version") or ""
    platform_name = DRIVER_TO_PLATFORM.get(driver, driver)

    return Entity(device=Device(
        name=hostname,
        device_type=DeviceType(
            model=model,
            manufacturer=Manufacturer(name=vendor),
        ),
        role=DeviceRole(name=role),
        platform=Platform(name=platform_name),
        site=Site(name=site_name),
        serial=serial,
        status="active",
        comments=f"OS: {os_version}" if os_version else "",
        tags=["network-collector"],
    ))


def build_interface_entities(interfaces: dict, hostname: str, model: str,
                             manufacturer: str, role: str,
                             site_name: str) -> list[Entity]:
    """Build Interface entities from NAPALM get_interfaces()."""
    entities = []
    dev_ref = _device_ref(hostname, model, manufacturer, role, site_name)

    for name, iface_data in interfaces.items():
        speed = iface_data.get("speed", 0)
        iface_type = map_interface_type(name, speed)

        mac = iface_data.get("mac_address") or ""
        iface_kwargs = dict(
            device=dev_ref,
            name=name,
            type=iface_type,
            enabled=iface_data.get("is_enabled", True),
            mtu=int(iface_data["mtu"]) if iface_data.get("mtu") else None,
            speed=int(speed * 1000) if speed else 0,  # NAPALM Mbps → NetBox Kbps
            description=iface_data.get("description") or "",
            tags=["network-collector"],
        )
        if mac:
            iface_kwargs["primary_mac_address"] = mac
        iface = Interface(**iface_kwargs)
        entities.append(Entity(interface=iface))

    return entities


def build_ip_entities(interfaces_ip: dict, hostname: str, model: str,
                      manufacturer: str, role: str,
                      site_name: str) -> list[Entity]:
    """Build IPAddress entities from NAPALM get_interfaces_ip()."""
    entities = []
    dev_ref = _device_ref(hostname, model, manufacturer, role, site_name)

    for iface_name, af_data in interfaces_ip.items():
        iface_type = map_interface_type(iface_name)

        for af in ("ipv4", "ipv6"):
            addrs = af_data.get(af, {})
            for addr, meta in addrs.items():
                prefix_len = meta.get("prefix_length", 32 if af == "ipv4" else 128)
                ip_str = f"{addr}/{prefix_len}"

                # Skip link-local IPv6
                if af == "ipv6" and addr.lower().startswith("fe80"):
                    continue

                entities.append(Entity(ip_address=IPAddress(
                    address=ip_str,
                    status="active",
                    assigned_object_interface=Interface(
                        device=dev_ref,
                        name=iface_name,
                        type=iface_type,
                    ),
                    tags=["network-collector"],
                )))

    return entities


def build_vlan_entities(vlans: dict, site_name: str) -> list[Entity]:
    """Build VLAN entities from NAPALM get_vlans()."""
    entities = []
    for vid_str, vlan_data in vlans.items():
        vid = int(vid_str)
        if vid in (0, 4095):
            continue
        name = vlan_data.get("name") or f"VLAN{vid}"
        entities.append(Entity(vlan=VLAN(
            vid=vid,
            name=name,
            site=Site(name=site_name),
            status="active",
            tags=["network-collector"],
        )))
    return entities


def build_vrf_entities(network_instances: dict) -> list[Entity]:
    """Build VRF entities from NAPALM get_network_instances()."""
    entities = []
    for vrf_name, vrf_data in network_instances.items():
        if vrf_name in ("default", "global"):
            continue
        rd = ""
        if vrf_data.get("state", {}).get("route_distinguisher"):
            rd = vrf_data["state"]["route_distinguisher"]
        entities.append(Entity(vrf=VRF(
            name=vrf_name,
            rd=rd or None,
            tags=["network-collector"],
        )))
    return entities


def build_prefix_entities(interfaces_ip: dict, site_name: str) -> list[Entity]:
    """Build Prefix entities from discovered interface IPs."""
    import ipaddress
    seen = set()
    entities = []

    for iface_name, af_data in interfaces_ip.items():
        for af in ("ipv4", "ipv6"):
            for addr, meta in af_data.get(af, {}).items():
                prefix_len = meta.get("prefix_length", 32 if af == "ipv4" else 128)
                try:
                    network = ipaddress.ip_network(f"{addr}/{prefix_len}", strict=False)
                except ValueError:
                    continue
                # Skip host routes and link-local
                if af == "ipv4" and prefix_len >= 31:
                    continue
                if af == "ipv6" and (prefix_len >= 127 or addr.lower().startswith("fe80")):
                    continue
                prefix_str = str(network)
                if prefix_str not in seen:
                    seen.add(prefix_str)
                    entities.append(Entity(prefix=Prefix(
                        prefix=prefix_str,
                        scope_site=Site(name=site_name),
                        status="active",
                        tags=["network-collector"],
                    )))

    return entities


def build_config_entity(config: dict, hostname: str, model: str,
                        manufacturer: str, role: str,
                        site_name: str) -> Entity | None:
    """Build a DeviceConfig entity from NAPALM get_config()."""
    running = config.get("running", "")
    startup = config.get("startup", "")
    if not running and not startup:
        return None
    dev_ref = _device_ref(hostname, model, manufacturer, role, site_name)
    return Entity(device=dev_ref, device_config=DeviceConfig(
        running=running.encode("utf-8") if running else None,
        startup=startup.encode("utf-8") if startup else None,
    ))


def build_inventory_entities(inventory: dict, hostname: str, model: str,
                             manufacturer: str, role: str,
                             site_name: str) -> list[Entity]:
    """Build InventoryItem entities from pyATS show inventory."""
    entities = []
    dev_ref = _device_ref(hostname, model, manufacturer, role, site_name)

    # pyATS inventory format varies by platform; handle the common structure
    items = inventory.get("main", {}).get("chassis", {})
    if not items:
        items = inventory
    for item_name, item_data in items.items():
        if isinstance(item_data, dict):
            pid = item_data.get("pid") or item_data.get("name") or ""
            sn = item_data.get("sn") or ""
            desc = item_data.get("descr") or item_data.get("description") or ""
            entities.append(Entity(inventory_item=InventoryItem(
                device=dev_ref,
                name=str(item_name)[:64],
                part_id=str(pid)[:50] if pid else None,
                serial=str(sn)[:50] if sn else None,
                description=str(desc)[:200] if desc else None,
                discovered=True,
                tags=["network-collector"],
            )))

    return entities


# ---------------------------------------------------------------------------
# Cable discovery from LLDP/CDP
# ---------------------------------------------------------------------------


def build_cable_entities_from_lldp(
    lldp_all: dict[str, dict],
    site_name: str,
    device_models: dict[str, tuple[str, str, str]],
) -> list[Entity]:
    """Build Cable entities from LLDP neighbor data collected from all devices.

    lldp_all: {hostname: {local_iface: [{remote_system_name, remote_port, ...}]}}
    device_models: {hostname: (model, manufacturer, role)}

    Deduplication: each link appears on both ends. We sort the pair and keep
    only one Cable per unique (deviceA:portA, deviceB:portB) tuple.
    """
    seen_links: set[tuple] = set()
    entities = []

    for local_host, iface_neighbors in lldp_all.items():
        local_info = device_models.get(local_host, ("Unknown", "Unknown", "Network Device"))

        for local_iface, neighbors in iface_neighbors.items():
            for neighbor in neighbors:
                remote_host = (
                    neighbor.get("remote_system_name")
                    or neighbor.get("remote_system_description", "")
                ).split(".")[0]  # Strip FQDN
                remote_port = neighbor.get("remote_port") or neighbor.get("remote_port_description", "")

                if not remote_host or not remote_port:
                    continue

                # Normalize interface names for matching
                local_norm = normalize_interface_name(local_iface)
                remote_norm = normalize_interface_name(remote_port)

                # Deduplicate: sorted link key
                link_a = f"{local_host}:{local_norm}"
                link_b = f"{remote_host}:{remote_norm}"
                link_key = tuple(sorted([link_a, link_b]))
                if link_key in seen_links:
                    continue
                seen_links.add(link_key)

                # Build device references for both ends
                remote_info = device_models.get(
                    remote_host, ("Unknown", "Unknown", "Network Device")
                )

                a_dev = _device_ref(local_host, *local_info, site_name)
                b_dev = _device_ref(remote_host, *remote_info, site_name)

                a_iface_type = map_interface_type(local_norm)
                b_iface_type = map_interface_type(remote_norm)

                cable = Cable(
                    a_terminations=[GenericObject(object_interface=Interface(
                        device=a_dev,
                        name=local_norm,
                        type=a_iface_type,
                    ))],
                    b_terminations=[GenericObject(object_interface=Interface(
                        device=b_dev,
                        name=remote_norm,
                        type=b_iface_type,
                    ))],
                    status="connected",
                    tags=["lldp-discovered"],
                )
                entities.append(Entity(cable=cable))
                log.info("  Cable: %s:%s <-> %s:%s",
                         local_host, local_norm, remote_host, remote_norm)

    return entities


def build_cable_entities_from_cdp(
    cdp_all: dict[str, dict],
    site_name: str,
    device_models: dict[str, tuple[str, str, str]],
    existing_links: set[tuple] | None = None,
) -> list[Entity]:
    """Build Cable entities from CDP neighbor data (pyATS parsed).

    cdp_all: {hostname: pyats_parsed_cdp_output}
    Only creates cables for links NOT already discovered via LLDP.
    """
    if existing_links is None:
        existing_links = set()
    entities = []

    for local_host, cdp_data in cdp_all.items():
        local_info = device_models.get(local_host, ("Unknown", "Unknown", "Network Device"))

        # pyATS CDP format: {"index": {1: {device_id, local_interface, port_id, ...}}}
        for idx, entry in cdp_data.get("index", {}).items():
            remote_host = (entry.get("device_id") or "").split(".")[0]
            local_iface = entry.get("local_interface") or ""
            remote_port = entry.get("port_id") or ""

            if not remote_host or not local_iface or not remote_port:
                continue

            local_norm = normalize_interface_name(local_iface)
            remote_norm = normalize_interface_name(remote_port)

            link_a = f"{local_host}:{local_norm}"
            link_b = f"{remote_host}:{remote_norm}"
            link_key = tuple(sorted([link_a, link_b]))
            if link_key in existing_links:
                continue
            existing_links.add(link_key)

            remote_info = device_models.get(
                remote_host, ("Unknown", "Unknown", "Network Device")
            )

            a_dev = _device_ref(local_host, *local_info, site_name)
            b_dev = _device_ref(remote_host, *remote_info, site_name)

            cable = Cable(
                a_terminations=[GenericObject(object_interface=Interface(
                    device=a_dev,
                    name=local_norm,
                    type=map_interface_type(local_norm),
                ))],
                b_terminations=[GenericObject(object_interface=Interface(
                    device=b_dev,
                    name=remote_norm,
                    type=map_interface_type(remote_norm),
                ))],
                status="connected",
                tags=["cdp-discovered"],
            )
            entities.append(Entity(cable=cable))
            log.info("  Cable (CDP): %s:%s <-> %s:%s",
                     local_host, local_norm, remote_host, remote_norm)

    return entities


# ---------------------------------------------------------------------------
# NetBox plugin API pushers (BGP, OSPF, IS-IS)
# ---------------------------------------------------------------------------


def push_bgp_to_netbox(bgp_data: dict, hostname: str, netbox_url: str,
                       netbox_token: str, site_name: str) -> None:
    """Push BGP neighbor data to the netbox-bgp plugin API."""
    if not HAS_REQUESTS or not bgp_data:
        return

    headers = {"Authorization": f"Bearer {netbox_token}", "Content-Type": "application/json"}
    base = netbox_url.rstrip("/")

    # First look up the local device in NetBox to get its ID
    resp = requests.get(f"{base}/api/dcim/devices/", params={"name": hostname},
                        headers=headers, timeout=10)
    if resp.status_code != 200:
        log.warning("  BGP push: cannot find device %s in NetBox", hostname)
        return
    devices = resp.json().get("results", [])
    if not devices:
        log.warning("  BGP push: device %s not found in NetBox", hostname)
        return
    local_device_id = devices[0]["id"]

    for vrf_name, peers in bgp_data.items():
        for peer_ip, peer_list in peers.items():
            for peer in peer_list:
                local_as = peer.get("local_as")
                remote_as = peer.get("remote_as")
                remote_id = peer.get("remote_id", "")

                if not local_as or not remote_as:
                    continue

                session_data = {
                    "name": f"{hostname} <-> {peer_ip}",
                    "device": local_device_id,
                    "local_address": None,  # Would need to resolve
                    "remote_address": None,
                    "local_as": {"asn": local_as},
                    "remote_as": {"asn": remote_as},
                    "status": "active" if peer.get("is_enabled") else "offline",
                    "description": f"Auto-discovered by network-collector",
                }

                try:
                    resp = requests.post(
                        f"{base}/api/plugins/bgp/sessions/",
                        headers=headers,
                        json=session_data,
                        timeout=10,
                    )
                    if resp.status_code in (200, 201):
                        log.info("  BGP session created: %s AS%s <-> %s AS%s",
                                 hostname, local_as, peer_ip, remote_as)
                    elif resp.status_code == 400 and "already exists" in resp.text.lower():
                        log.debug("  BGP session already exists: %s <-> %s", hostname, peer_ip)
                    else:
                        log.warning("  BGP push failed (%d): %s", resp.status_code, resp.text[:200])
                except Exception as exc:
                    log.warning("  BGP push error: %s", exc)


# ---------------------------------------------------------------------------
# Orchestration
# ---------------------------------------------------------------------------


def collect_all_entities(inventory: dict, env_file: str = ".env") -> tuple[list[Entity], dict]:
    """Walk all devices in inventory, collect data, build entities."""
    defaults = inventory.get("defaults", {})
    site_name = defaults.get("site", "main")
    entities: list[Entity] = []
    lldp_all: dict[str, dict] = {}
    cdp_all: dict[str, dict] = {}
    device_models: dict[str, tuple[str, str, str]] = {}
    bgp_all: dict[str, dict] = {}

    for dev_entry in inventory["devices"]:
        cfg = merge_device_config(dev_entry, defaults)
        host = cfg["host"]
        driver = cfg.get("driver", "ios")
        role = cfg.get("role", "Network Device")
        username = cfg.get("username", "admin")
        password = cfg.get("password", "")
        secret = cfg.get("secret", "")
        timeout = int(cfg.get("timeout", 60))
        optional_args = cfg.get("optional_args", {})

        log.info("Connecting to %s (driver=%s, role=%s)...", host, driver, role)

        # --- Decide collection path: NAPALM or Netmiko fallback ---
        use_netmiko = driver in NETMIKO_ONLY_DRIVERS

        if use_netmiko:
            log.info("  Using Netmiko fallback (no NAPALM driver for %s)", driver)
            try:
                conn = connect_netmiko(host, driver, username, password, secret, timeout)
            except Exception as exc:
                log.error("Failed to connect to %s via Netmiko: %s", host, exc)
                continue
            try:
                collected_data = collect_netmiko_data(conn, driver)
            finally:
                try:
                    conn.disconnect()
                except Exception:
                    pass
        else:
            # --- NAPALM collection ---
            try:
                dev = connect_device(host, driver, username, password, secret,
                                     timeout, optional_args)
            except Exception as exc:
                log.error("Failed to connect to %s: %s", host, exc)
                continue

            try:
                collected_data = collect_napalm_data(dev)
            finally:
                try:
                    dev.close()
                except Exception:
                    pass

        if not collected_data.get("facts"):
            log.error("No facts for %s, skipping", host)
            continue

        facts = collected_data["facts"]
        hostname = facts.get("hostname") or host
        model = facts.get("model") or "Unknown"
        vendor = facts.get("vendor") or DRIVER_TO_MANUFACTURER.get(driver, "Unknown")

        # Track device info for cable building
        device_models[hostname] = (model, vendor, role)

        # Device entity
        entities.append(build_device_entity(facts, driver, role, site_name, host))

        # Interface entities
        if collected_data.get("interfaces"):
            entities.extend(build_interface_entities(
                collected_data["interfaces"], hostname, model, vendor, role, site_name
            ))

        # IP entities (IPv4 + IPv6)
        if collected_data.get("interfaces_ip"):
            entities.extend(build_ip_entities(
                collected_data["interfaces_ip"], hostname, model, vendor, role, site_name
            ))
            # Prefix entities from discovered IPs
            entities.extend(build_prefix_entities(
                collected_data["interfaces_ip"], site_name
            ))

        # VLAN entities
        if collected_data.get("vlans"):
            entities.extend(build_vlan_entities(collected_data["vlans"], site_name))

        # VRF entities
        if collected_data.get("network_instances"):
            entities.extend(build_vrf_entities(collected_data["network_instances"]))

        # Config entity
        if collected_data.get("config"):
            config_entity = build_config_entity(
                collected_data["config"], hostname, model, vendor, role, site_name
            )
            if config_entity:
                entities.append(config_entity)

        # LLDP neighbors (saved for cable building later)
        if collected_data.get("lldp_neighbors"):
            lldp_all[hostname] = collected_data["lldp_neighbors"]

        # BGP data (saved for plugin API push later)
        if collected_data.get("bgp_neighbors"):
            bgp_all[hostname] = collected_data["bgp_neighbors"]

        # --- pyATS collection (optional, NAPALM devices only) ---
        if not use_netmiko and HAS_PYATS and driver in ("ios", "iosxr", "nxos", "nxos_ssh"):
            log.info("  Running pyATS parsers...")
            pyats_data = collect_pyats_data(host, driver, username, password, secret)

            if pyats_data.get("cdp_neighbors"):
                cdp_all[hostname] = pyats_data["cdp_neighbors"]

            if pyats_data.get("inventory"):
                entities.extend(build_inventory_entities(
                    pyats_data["inventory"], hostname, model, vendor, role, site_name
                ))

    # --- Cable entities from LLDP ---
    if lldp_all:
        log.info("Building cable entities from LLDP data...")
        cable_entities = build_cable_entities_from_lldp(lldp_all, site_name, device_models)
        entities.extend(cable_entities)

        # Extract seen links for CDP dedup
        seen_links = set()
        for local_host, iface_neighbors in lldp_all.items():
            for local_iface, neighbors in iface_neighbors.items():
                for neighbor in neighbors:
                    remote_host = (neighbor.get("remote_system_name") or "").split(".")[0]
                    remote_port = neighbor.get("remote_port") or ""
                    if remote_host and remote_port:
                        local_norm = normalize_interface_name(local_iface)
                        remote_norm = normalize_interface_name(remote_port)
                        link_key = tuple(sorted([
                            f"{local_host}:{local_norm}",
                            f"{remote_host}:{remote_norm}",
                        ]))
                        seen_links.add(link_key)
    else:
        seen_links = set()

    # --- Cable entities from CDP (only new links) ---
    if cdp_all:
        log.info("Building cable entities from CDP data...")
        cdp_cable_entities = build_cable_entities_from_cdp(
            cdp_all, site_name, device_models, seen_links
        )
        entities.extend(cdp_cable_entities)

    return entities, bgp_all


def ingest_entities(entities: list[Entity], dry_run: bool = False) -> None:
    """Send entities to Diode (or dry-run print them)."""
    if not entities:
        log.warning("No entities to ingest")
        return

    target = os.environ.get("DIODE_TARGET", "grpc://localhost:8080/diode")
    client_id = os.environ.get("DIODE_CLIENT_ID",
                               os.environ.get("INGESTER_CLIENT_ID", "diode-ingester"))
    client_secret = os.environ.get("DIODE_CLIENT_SECRET",
                                   os.environ.get("INGESTER_CLIENT_SECRET", ""))

    if dry_run:
        log.info("DRY RUN: %d entities would be ingested", len(entities))
        for i, e in enumerate(entities):
            log.info("  [%d] %s", i, e)
        return

    if not client_secret:
        log.error("DIODE_CLIENT_SECRET not set — cannot ingest")
        sys.exit(1)

    log.info("Ingesting %d entities to %s ...", len(entities), target)

    with DiodeClient(
        target=target,
        client_id=client_id,
        client_secret=client_secret,
        app_name="network-collector",
        app_version="0.1.0",
    ) as client:
        resp = client.ingest(entities=entities)
        if resp.errors:
            log.error("Ingestion errors: %s", resp.errors)
        else:
            log.info("Ingestion successful")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def main():
    parser = argparse.ArgumentParser(description="Network device collector for NetBox")
    parser.add_argument("--inventory", "-i", required=True,
                        help="Path to device inventory YAML file")
    parser.add_argument("--dry-run", action="store_true",
                        help="Collect data but don't ingest")
    parser.add_argument("--log-level", default="INFO",
                        choices=["DEBUG", "INFO", "WARNING", "ERROR"])
    parser.add_argument("--env-file", default=".env",
                        help="Path to .env file (default: .env)")
    parser.add_argument("--no-bgp-push", action="store_true",
                        help="Skip pushing BGP data to netbox-bgp plugin")
    parser.add_argument("--no-pyats", action="store_true",
                        help="Skip pyATS/Genie collection even if available")
    args = parser.parse_args()

    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format="%(asctime)s %(name)s %(levelname)s %(message)s",
    )

    load_dotenv(args.env_file)

    if args.no_pyats:
        global HAS_PYATS
        HAS_PYATS = False

    # Load inventory
    inventory = load_inventory(args.inventory)
    log.info("Loaded %d devices from inventory", len(inventory["devices"]))

    # Collect and build entities
    entities, bgp_all = collect_all_entities(inventory, args.env_file)
    log.info("Total entities: %d", len(entities))

    # Ingest via Diode
    ingest_entities(entities, dry_run=args.dry_run)

    # Push BGP sessions to netbox-bgp plugin API
    if not args.no_bgp_push and not args.dry_run and bgp_all:
        netbox_url = os.environ.get("NETBOX_URL", "http://172.19.77.160:8000")
        netbox_token = os.environ.get("NETBOX_API_TOKEN", "")
        if netbox_token:
            for hostname, bgp_data in bgp_all.items():
                log.info("Pushing BGP data for %s to netbox-bgp...", hostname)
                push_bgp_to_netbox(bgp_data, hostname, netbox_url, netbox_token,
                                   inventory.get("defaults", {}).get("site", "main"))
        else:
            log.warning("NETBOX_API_TOKEN not set, skipping BGP plugin push")

    log.info("Done!")


if __name__ == "__main__":
    main()