#!/usr/bin/env python3
"""
NEnG Device Scanner — Fast network scanner for SCPI instruments.

Scans local networks for devices with SCPI ports open (5025, 5026, 1394),
then verifies they are SCPI instruments by querying *IDN?. Both NEnG
instruments and any other standard SCPI device (oscilloscopes, bench
instruments, etc.) are discovered. Supports concurrent scanning for maximum
speed (~3-5 seconds for a /24 network).

Usage (installed via pip):
    # Scan default network (auto-detected)
    neng-device-scan

    # Scan specific network
    neng-device-scan -n 192.168.1.0/24

    # JSON output
    neng-device-scan -n 192.168.1.0/24 --json

    # Save results to file
    neng-device-scan -n 192.168.1.0/24 -o devices.txt

    # Verbose mode (show all port checks)
    neng-device-scan -v

(c) 2026 Prof. Flavio ABREU ARAUJO. All rights reserved.
"""

from __future__ import annotations

import argparse
import concurrent.futures
import json
import platform
import socket
import sys
import threading
import time
from ipaddress import IPv4Network, ip_network

# Conditional import of colorama (for Windows ANSI color support)
try:
    import colorama
    COLORAMA_AVAILABLE = True
except ImportError:
    COLORAMA_AVAILABLE = False

# ---------------------------------------------------------------------------
# ANSI Color Codes
# ---------------------------------------------------------------------------
CYAN = "\033[96m"
YELLOW = "\033[93m"
GREEN = "\033[92m"
MAGENTA = "\033[95m"
RED = "\033[91m"
BLUE = "\033[94m"
RESET = "\033[0m"
BOLD = "\033[1m"
DIM = "\033[2m"

DEFAULT_SCPI_PORT = 5025
# Additional SCPI ports probed when scanning for non-NEnG devices.
# 5025  — NEnG instruments (primary)
# 5026  — alternative SCPI-RAW (some Keysight / Rohde & Schwarz instruments)
# 1394  — LXI / IEEE-1394 SCPI-RAW (LeCroy WaveRunner, etc.)
ALL_SCPI_PORTS = [5025, 5026, 1394]
DEFAULT_PORT_TIMEOUT = 0.5  # seconds
DEFAULT_SCPI_TIMEOUT = 1.5  # seconds
DEFAULT_MAX_WORKERS = 100


# ---------------------------------------------------------------------------
# Network Utilities
# ---------------------------------------------------------------------------
def get_local_network() -> str | None:
    """Auto-detect local network CIDR (e.g., 192.168.1.0/24)."""
    try:
        # Create a socket to find the local IP
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.settimeout(0.1)
        # Connect to a public DNS server (doesn't actually send data)
        s.connect(("8.8.8.8", 80))
        local_ip = s.getsockname()[0]
        s.close()

        # Convert to network address with /24 subnet
        network = ip_network(f"{local_ip}/24", strict=False)
        return str(network)
    except Exception:
        return None


def get_all_local_networks() -> list[str]:
    """
    Get all local network CIDRs for all network interfaces.

    Supports Linux/macOS (via ``ifconfig``), Windows (via ``ipconfig``),
    and falls back to a socket-based single-interface method.

    Returns:
        List of network CIDR strings (e.g., ['192.168.1.0/24', '10.0.0.0/24'])
    """
    import subprocess

    networks: list[str] = []

    # ------------------------------------------------------------------
    # Helper: add a (ip, mask) pair to the list, skipping loopback/link-local
    # ------------------------------------------------------------------
    def _add(ip_addr: str, prefix_len: int) -> None:
        if ip_addr.startswith("127.") or ip_addr.startswith("169.254."):
            return
        try:
            net_str = str(ip_network(f"{ip_addr}/{prefix_len}", strict=False))
            if net_str not in networks:
                networks.append(net_str)
        except ValueError:
            pass

    # ------------------------------------------------------------------
    # Linux / macOS — ifconfig
    # ------------------------------------------------------------------
    try:
        result = subprocess.run(["ifconfig"], capture_output=True, text=True, timeout=2)
        if result.returncode == 0:
            current_iface = None
            for line in result.stdout.split("\n"):
                line = line.strip()
                if line and not line.startswith(" ") and ":" in line:
                    current_iface = line.split(":")[0]
                if line.startswith("inet ") and current_iface:
                    parts = line.split()
                    if len(parts) >= 4:
                        ip_addr = parts[1]
                        netmask = parts[3]
                        try:
                            if netmask.startswith("0x"):
                                prefix_len = bin(int(netmask, 16)).count("1")
                            else:
                                nm_parts = [int(x) for x in netmask.split(".")]
                                nm_int = (
                                    (nm_parts[0] << 24)
                                    | (nm_parts[1] << 16)
                                    | (nm_parts[2] << 8)
                                    | nm_parts[3]
                                )
                                prefix_len = bin(nm_int).count("1")
                            _add(ip_addr, prefix_len)
                        except (ValueError, IndexError):
                            pass
    except (subprocess.TimeoutExpired, FileNotFoundError, OSError):
        pass

    if networks:
        return networks

    # ------------------------------------------------------------------
    # Windows — ipconfig
    # ------------------------------------------------------------------
    if platform.system() == "Windows":
        try:
            result = subprocess.run(
                ["ipconfig"], capture_output=True, text=True, timeout=2
            )
            if result.returncode == 0:
                current_ip: str | None = None
                for line in result.stdout.split("\n"):
                    line = line.strip()
                    # "IPv4 Address. . . : 192.168.x.x (Preferred)"
                    if ("IPv4 Address" in line or "IP Address" in line) and ":" in line:
                        ip = line.split(":")[-1].strip().replace("(Preferred)", "").strip()
                        if ip and not ip.startswith("127.") and not ip.startswith("169.254."):
                            current_ip = ip
                    elif "Subnet Mask" in line and ":" in line and current_ip:
                        mask = line.split(":")[-1].strip()
                        try:
                            nm_parts = [int(x) for x in mask.split(".")]
                            nm_int = (
                                (nm_parts[0] << 24)
                                | (nm_parts[1] << 16)
                                | (nm_parts[2] << 8)
                                | nm_parts[3]
                            )
                            _add(current_ip, bin(nm_int).count("1"))
                        except (ValueError, IndexError):
                            pass
                        current_ip = None
        except (subprocess.TimeoutExpired, FileNotFoundError, OSError):
            pass

    if networks:
        return networks

    # ------------------------------------------------------------------
    # Last-resort: single-interface socket trick (any platform)
    # ------------------------------------------------------------------
    fallback = get_local_network()
    if fallback:
        networks.append(fallback)

    return networks


# ---------------------------------------------------------------------------
# Port Scanning
# ---------------------------------------------------------------------------
def check_port(
    ip: str, port: int = DEFAULT_SCPI_PORT, timeout: float = DEFAULT_PORT_TIMEOUT
) -> bool:
    """
    Check if TCP port is open on given IP address.

    Args:
        ip: IP address to check
        port: TCP port number (default: 5025)
        timeout: Connection timeout in seconds (default: 0.5)

    Returns:
        True if port is open, False otherwise
    """
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.settimeout(timeout)
            result = sock.connect_ex((ip, port))
            return result == 0
    except (socket.timeout, OSError):
        return False


def verify_neng_device(
    ip: str, port: int = DEFAULT_SCPI_PORT, timeout: float = DEFAULT_SCPI_TIMEOUT
) -> str | None:
    """
    Verify if a device responds to *IDN? on the given port.

    Accepts any standard SCPI instrument — both NEnG devices and third-party
    bench instruments (oscilloscopes, source-measure units, etc.).

    Args:
        ip: IP address of device
        port: SCPI port (default: 5025)
        timeout: SCPI query timeout in seconds (default: 1.5)

    Returns:
        Device identification string if the device responds to *IDN?,
        None otherwise.
    """
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            sock.settimeout(timeout)
            sock.connect((ip, port))

            # Send *IDN? query
            sock.sendall(b"*IDN?\n")

            # Read response
            response = sock.recv(1024).decode().strip()

            # Reject error/rejection messages from the SCPI server
            # (e.g. '-1,"Max TCP clients reached, connection refused"')
            if response and not response.startswith("-"):
                return response

    except (socket.timeout, OSError, UnicodeDecodeError):
        pass

    return None


# ---------------------------------------------------------------------------
# Network Scanner
# ---------------------------------------------------------------------------
def scan_network(
    network: str,
    port: int = DEFAULT_SCPI_PORT,
    ports: list[int] | None = None,
    max_workers: int = DEFAULT_MAX_WORKERS,
    verbose: bool = False,
    progress_callback=None,
    stop_event: threading.Event | None = None,
) -> list[tuple[str, int, str]]:
    """
    Scan a network for SCPI instruments on one or more TCP ports.

    Probes every host in the network for open SCPI ports, then verifies
    each candidate with ``*IDN?``. Both NEnG instruments and any other
    standard SCPI device (oscilloscopes, bench instruments, etc.) are
    returned.

    When ``ports`` is provided it overrides the single ``port`` argument.
    Pass ``ports=ALL_SCPI_PORTS`` (``[5025, 5026, 1394]``) to discover
    non-NEnG devices that use alternative SCPI ports.

    Args:
        network: Network in CIDR notation (e.g., "192.168.1.0/24")
        port: Single TCP port to scan (default: 5025). Ignored when
              ``ports`` is supplied.
        ports: List of TCP ports to scan. Supersedes ``port`` when given.
        max_workers: Number of concurrent scanning threads (default: 100)
        verbose: Print detailed progress (default: False)
        progress_callback: Optional callback function(current, total)
        stop_event: Optional threading.Event; scanning stops when set.

    Returns:
        List of tuples ``(ip_address, scpi_port, device_id_string)``,
        deduplicated by IP so that instruments reachable on multiple ports
        appear only once. The ``scpi_port`` is the actual TCP port on which
        the device responded.
    """
    scan_ports: list[int] = ports if ports is not None else [port]

    try:
        net = IPv4Network(network, strict=False)
    except ValueError as e:
        print(f"{RED}✗ Invalid network: {network}{RESET}")
        print(f"  Error: {e}")
        return []

    try:
        hosts = [str(ip) for ip in net.hosts()]
    except KeyboardInterrupt:
        if stop_event:
            stop_event.set()
        return []
    total_hosts = len(hosts)

    if verbose:
        ports_str = ", ".join(str(p) for p in scan_ports)
        print(f"{CYAN}Scanning {network} ({total_hosts} hosts) on port(s) {ports_str}...{RESET}")

    # Stage 1: Fast port scan — collect (ip, port) pairs with open ports.
    # We submit one check per (host, port) combination and collect the winners.
    open_endpoints: list[tuple[str, int]] = []
    checked_count = 0
    total_checks = total_hosts * len(scan_ports)

    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_endpoint: dict = {
            executor.submit(check_port, ip, p): (ip, p)
            for ip in hosts
            for p in scan_ports
        }

        try:
            for future in concurrent.futures.as_completed(future_to_endpoint):
                if stop_event and stop_event.is_set():
                    for f in future_to_endpoint:
                        f.cancel()
                    break

                endpoint = future_to_endpoint[future]
                checked_count += 1

                if progress_callback:
                    progress_callback(checked_count, total_checks)

                try:
                    is_open = future.result()
                    if is_open:
                        open_endpoints.append(endpoint)
                        if verbose:
                            ip, p = endpoint
                            print(f"{YELLOW}  Port {p} open: {ip}{RESET}")
                except Exception as e:
                    if verbose:
                        ip, p = endpoint
                        print(f"{RED}  Error checking {ip}:{p}: {e}{RESET}")

        except KeyboardInterrupt:
            if stop_event:
                stop_event.set()
            for f in future_to_endpoint:
                f.cancel()

    # Stage 2: Verify devices — query *IDN? on each open endpoint.
    # Deduplicate by IP: once we have a valid IDN string for an IP we skip
    # any remaining ports for that same host.
    if stop_event and stop_event.is_set():
        return []

    if verbose and open_endpoints:
        print(f"{CYAN}Verifying {len(open_endpoints)} endpoint(s)...{RESET}")

    devices_found: list[tuple[str, int, str]] = []
    seen_ips: set[str] = set()

    for ip, p in open_endpoints:
        if stop_event and stop_event.is_set():
            break
        if ip in seen_ips:
            continue
        idn = verify_neng_device(ip, p)
        if idn:
            devices_found.append((ip, p, idn))
            seen_ips.add(ip)
            if verbose:
                print(f"{GREEN}  ✓ SCPI device found: {ip}:{p} → {idn}{RESET}")
        else:
            if verbose:
                print(f"{DIM}  ✗ No SCPI response: {ip}:{p}{RESET}")

    return devices_found


# ---------------------------------------------------------------------------
# Output Formatters
# ---------------------------------------------------------------------------
def format_table(devices: list[tuple[str, int, str]]) -> str:
    """Format device list as a table."""
    if not devices:
        return f"{YELLOW}No SCPI devices found.{RESET}"

    lines = [
        "",
        f"{BOLD}{'=' * 70}{RESET}",
        f"{BOLD}Found {len(devices)} SCPI device(s):{RESET}",
        f"{BOLD}{'=' * 70}{RESET}",
    ]

    for ip, p, idn in devices:
        port_hint = f":{p}" if p != DEFAULT_SCPI_PORT else ""
        lines.append(f"{GREEN}  • {ip:15s}{port_hint:<6}{RESET} - {idn}")

    lines.append(f"{BOLD}{'=' * 70}{RESET}")
    return "\n".join(lines)


def format_json(devices: list[tuple[str, int, str]]) -> str:
    """Format device list as JSON."""
    device_list = [{"ip": ip, "port": p, "id": idn} for ip, p, idn in devices]
    return json.dumps({"devices": device_list, "count": len(device_list)}, indent=2)


def format_simple(devices: list[tuple[str, int, str]]) -> str:
    """Format device list as simple IP list."""
    return "\n".join(ip for ip, _, _ in devices)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> int:
    """Main entry point for neng-device-scan command."""
    # Enable ANSI colors on Windows terminals
    if COLORAMA_AVAILABLE and platform.system() == "Windows":
        colorama.just_fix_windows_console()

    parser = argparse.ArgumentParser(
        description="Scan network for SCPI instruments (NEnG and third-party) on SCPI ports",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s                          # Scan auto-detected local network
  %(prog)s -n 192.168.1.0/24        # Scan specific network
  %(prog)s -n 10.0.0.0/16 -w 200    # Scan large network with more workers
  %(prog)s --json                   # JSON output
  %(prog)s -o devices.txt           # Save to file
  %(prog)s -v                       # Verbose output

(c) 2026 Prof. Flavio ABREU ARAUJO. All rights reserved.
        """,
    )

    parser.add_argument(
        "-n",
        "--network",
        help="Network to scan in CIDR notation (e.g., 192.168.1.0/24). "
        "Auto-detects local network if not specified.",
    )

    parser.add_argument(
        "-p",
        "--port",
        type=int,
        default=None,
        help=(
            f"Single SCPI port to scan (default: scan all standard SCPI ports "
            f"{ALL_SCPI_PORTS}). Use this to restrict scanning to one port."
        ),
    )

    parser.add_argument(
        "-w",
        "--workers",
        type=int,
        default=DEFAULT_MAX_WORKERS,
        help=f"Number of parallel scanning threads (default: {DEFAULT_MAX_WORKERS})",
    )

    parser.add_argument(
        "-o",
        "--output",
        help="Save results to file (default: print to stdout)",
    )

    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Show detailed scanning progress",
    )

    parser.add_argument(
        "--large-networks",
        action="store_true",
        help="Also scan auto-detected networks larger than /24 (>254 hosts). "
        "Skipped by default to keep scans fast.",
    )

    format_group = parser.add_mutually_exclusive_group()
    format_group.add_argument(
        "--json",
        action="store_true",
        help="Output results as JSON",
    )
    format_group.add_argument(
        "--simple",
        action="store_true",
        help="Output only IP addresses (one per line)",
    )

    args = parser.parse_args()

    # Resolve ports to scan
    scan_ports = [args.port] if args.port is not None else ALL_SCPI_PORTS
    ports_str = ", ".join(str(p) for p in scan_ports)

    # Determine network(s) to scan
    if args.network:
        # Single network specified by user
        networks = [args.network]
        if args.verbose:
            print(f"{CYAN}Scanning specified network: {args.network} on port(s) {ports_str}{RESET}")
        else:
            print(f"{CYAN}Scanning {args.network} for SCPI devices...{RESET}")
    else:
        # Auto-detect all local networks
        networks = get_all_local_networks()
        if not networks:
            print(f"{RED}✗ Could not auto-detect local networks.{RESET}")
            print(f"{YELLOW}  Please specify network with -n/--network{RESET}")
            return 1

        # Drop networks larger than /24 unless the user opts in
        if not args.large_networks:
            small = [n for n in networks if IPv4Network(n, strict=False).prefixlen >= 24]
            skipped = len(networks) - len(small)
            if skipped:
                print(
                    f"{DIM}  Skipping {skipped} large network(s) with prefix < /24 "
                    f"(use --large-networks to include){RESET}"
                )
            networks = small
            if not networks:
                print(f"{RED}✗ Only large networks detected (prefix < /24).{RESET}")
                print(f"{YELLOW}  Use --large-networks to scan them.{RESET}")
                return 1

        if args.verbose:
            print(f"{CYAN}Auto-detected {len(networks)} network(s), ports {ports_str}:{RESET}")
            for net in networks:
                print(f"{DIM}  • {net}{RESET}")
        else:
            print(f"{CYAN}Scanning {len(networks)} network(s) for SCPI devices...{RESET}")

    # Sort networks: smallest first so fast results appear early
    networks.sort(key=lambda n: IPv4Network(n, strict=False).num_addresses)

    # Start scan
    start_time = time.time()
    stop_event = threading.Event()

    # Scan all networks
    devices = []
    try:
        for network in networks:
            network_devices = scan_network(
                network=network,
                ports=scan_ports,
                max_workers=args.workers,
                verbose=args.verbose,
                stop_event=stop_event,
            )
            devices.extend(network_devices)
            if stop_event.is_set():
                break
    except KeyboardInterrupt:
        stop_event.set()

    elapsed = time.time() - start_time
    interrupted = stop_event.is_set()

    if interrupted:
        print(f"\n{YELLOW}Scan interrupted after {elapsed:.2f}s.{RESET}")
        if not devices:
            return 130

    # Format output
    if args.json:
        output = format_json(devices)
    elif args.simple:
        output = format_simple(devices)
    else:
        output = format_table(devices)

    # Print or save
    if args.output:
        try:
            with open(args.output, "w") as f:
                # Write plain text (no ANSI codes) to file
                if args.json or args.simple:
                    f.write(output)
                else:
                    # Strip ANSI codes for file output
                    import re

                    clean_output = re.sub(r"\033\[[0-9;]+m", "", output)
                    f.write(clean_output)
            print(f"{GREEN}✓ Results saved to {args.output}{RESET}")
            print(f"{DIM}  Found {len(devices)} device(s) in {elapsed:.2f}s{RESET}")
        except OSError as e:
            print(f"{RED}✗ Error writing to {args.output}: {e}{RESET}")
            return 1
    else:
        print(output)
        if not args.simple and not args.json:
            if interrupted:
                print(f"{DIM}Partial results — scan interrupted after {elapsed:.2f}s{RESET}")
            else:
                print(f"{DIM}Scan completed in {elapsed:.2f}s{RESET}")

    return 130 if interrupted else 0


if __name__ == "__main__":
    sys.exit(main())
