"""SNMP v1/v2c probe engine: UDP transport and ICMP handling over the codec.

A single unconnected UDP socket fires an SNMP GET at each target and a receiver
callback matches replies by request-id + source address. There is no per-host
engine, socket, or thread, and the in-flight request window is bounded
(max_inflight), so the scan's working set stays flat regardless of the address
space -- though the responders/unreachable maps below do grow with the number
of hosts that answer or are terminally classified.
scan() consumes any iterable of Target and paces sends to a packets-per-second
rate; each result is delivered as a plain dict to an on_result(record) callback,
which decides what to do with it (serialize, store, forward, ...). The callback
runs on the event loop thread, so it should not block. Every result is also
retained on the Prober (responders / unreachable), so a caller that doesn't need
streaming can pass a no-op callback and read those after the scan.

When ProbeConfig.icmp is set (Linux only, no root), the socket also collects
ICMP errors via IP_RECVERR / the socket error queue to classify non-responders
(e.g. port-closed). An optional label_for(ip) callback tags those errors with a
group name, since the ICMP error carries only the IP, not the original target.

The BER codec and response decoding are hand-rolled in the sibling module
snmp_wire (shared with the bulkwalk transport), and the ICMP error-queue
decoding in icmp_errqueue, so there are no external dependencies; this module
builds only one PDU shape -- a GET of a fixed set of scalar OIDs.

Entry points:
    Prober(config, on_result, label_for=None)
    await run(config, on_result, targets, label_for=None) -> (Prober, seconds)
"""
import asyncio
import socket
import sys
import time
import traceback
from collections import namedtuple
from dataclasses import dataclass
from typing import Optional

# Codec from snmp_wire; socket error-queue / ICMP decoding from icmp_errqueue.
# The star import brings snmp_wire's public codec names; the explicit snmp_wire
# line re-exposes the private codec helpers this module's tests reach.
from preseem.snmp_wire import *  # noqa: F401,F403
from preseem.snmp_wire import (_decode_oid, _decode_value, _read_tlv)
from preseem.icmp_errqueue import (CLOSED_STATUSES, IP_RECVERR, MSG_ERRQUEUE,
                                   parse_icmp_error)

# ---------------------------------------------------------------------------
# OIDs we collect (RFC1213 system group scalars)
# ---------------------------------------------------------------------------
SYS_OIDS = [
    ('sysDescr', '1.3.6.1.2.1.1.1.0'),
    ('sysObjectID', '1.3.6.1.2.1.1.2.0'),
    ('sysUpTime', '1.3.6.1.2.1.1.3.0'),
    ('sysName', '1.3.6.1.2.1.1.5.0'),
]
OID_TO_FIELD = {oid: name for name, oid in SYS_OIDS}
SCAN_OIDS = [oid for _, oid in SYS_OIDS]
# The scan OID set is fixed, so its variable-bindings are identical in every
# probe -- encode them once and reuse (see assemble_get in the send path).
_SCAN_VARBINDS = encode_varbinds(SCAN_OIDS)

# A probe target. ``subnet`` is a free-form grouping label echoed unchanged
# into the result record. A host's targets should be produced consecutively so
# the "skip a host once it resolves" short-circuit applies. That skip is only
# best-effort: a host's communities are sent back-to-back within a burst, so a
# later community's probe often goes out before the earlier one's reply lands.
# _record dedupes by host, so at most one result is kept regardless -- the
# short-circuit saves packets when a reply beats the next send; it does not
# guarantee only one community is ever tried.
Target = namedtuple('Target', ('subnet', 'ip', 'community'))
# A request we have sent and are waiting on. ``tries`` is the number of times
# it has been sent so far; the packet is rebuilt per send because the version
# can change between attempts (v2c early, v1 on the final try).
Pending = namedtuple('Pending', ('target', 'deadline', 'tries'))


@dataclass
class ProbeConfig:
    """Transport/protocol knobs for the probe engine.

    Typically only a couple are overridden, e.g. ProbeConfig(rate=500, retries=2).
    """
    port: int = 161
    timeout: float = 2.0
    retries: int = 1
    rate: float = 2000.0
    burst: int = 100
    max_inflight: int = 2500
    icmp: bool = False
    emit_icmp: Optional[bool] = None  # None -> follow icmp
    primary_version: int = SNMP_V2C
    fallback_version: Optional[int] = None  # v1-fallback piggyback if set

    def __post_init__(self):
        if self.emit_icmp is None:
            self.emit_icmp = self.icmp
        if self.rate <= 0 or self.burst <= 0:
            # scan() computes interval = burst / rate; reject up front rather
            # than raising ZeroDivisionError mid-scan.
            raise ValueError('rate and burst must be positive')
        if self.fallback_version is not None and self.retries < 1:
            raise ValueError(
                'fallback_version requires retries >= 1 so the primary version '
                'gets a clean first attempt before the final-retry fallback')


# ---------------------------------------------------------------------------
# Prober -- the probe engine
# ---------------------------------------------------------------------------
class Prober:

    def __init__(self, config, on_result, label_for=None):
        self.config = config
        self.on_result = on_result  # callback invoked once per result dict
        self.label_for = label_for  # optional callable(ip) -> group label
        self.icmp = config.icmp
        self.emit_icmp = config.emit_icmp
        self.icmp_enabled = False
        # Version strategy: probe primary_version on every attempt except, when
        # fallback_version is set, the final one (the v1-fallback / retry
        # piggyback). None fallback => single-version run.
        self.primary_version = config.primary_version
        self.fallback_version = config.fallback_version
        self.sock = None
        self.pending = {}  # request_id -> Pending
        self.responders = {}  # ip -> result dict (dedupe by host)
        self.unreachable = {}  # ip -> status (from ICMP)
        self._rid = 0
        self.sent = 0
        self.retransmits = 0
        self.send_errors = 0
        self.late_replies = 0  # replies matched after their deadline
        self.callback_errors = 0  # exceptions raised by on_result

    # -- socket lifecycle ---------------------------------------------------
    def open(self, loop):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.setblocking(False)
        # Socket buffer sizes are left to the OS defaults (net.core.rmem_default
        # / wmem_default), which the deployment tunes via sysctl.
        if self.icmp:
            try:
                self.sock.setsockopt(socket.IPPROTO_IP, IP_RECVERR, 1)
                self.icmp_enabled = True
            except OSError as e:
                print('warning: ICMP error reporting unavailable (%s); '
                      'continuing without it' % e,
                      file=sys.stderr)
        self.sock.bind(('0.0.0.0', 0))
        loop.add_reader(self.sock.fileno(), self._on_readable)

    def close(self, loop):
        if self.sock is None:  # open() failed before creating the socket
            return
        try:
            loop.remove_reader(self.sock.fileno())
        except (OSError, ValueError):
            pass
        self.sock.close()

    # -- readable callback --------------------------------------------------
    def _on_readable(self):
        # Drain the error queue first so ICMP verdicts land before we would
        # otherwise retransmit to the same host.
        if self.icmp_enabled:
            self._drain_errqueue()
        self._drain_data()

    def _drain_data(self):
        while True:
            try:
                data, addr = self.sock.recvfrom(4096)
            except (BlockingIOError, InterruptedError):
                return
            except ConnectionRefusedError:
                # A queued ICMP error surfaced on the normal path; its detail
                # comes via the error queue, so keep draining.
                continue
            except OSError:
                return
            try:
                resp = decode_response(data)
            except (IndexError, ValueError):
                continue
            if resp is None:
                continue
            # Peek, don't pop: a reply whose source IP doesn't match (stray,
            # spoofed, or a multi-homed device answering from another
            # interface) must not destroy the real in-flight entry, which still
            # needs to retransmit/time out normally.
            entry = self.pending.get(resp.request_id)
            if entry is None or addr[0] != entry.target.ip:
                continue
            del self.pending[resp.request_id]
            if entry.deadline <= time.monotonic():
                # Reply arrived after its own timeout deadline: it was in the
                # socket but not serviced in time. We still record it (the
                # pending entry survived), but count it for visibility.
                self.late_replies += 1
            self._record(entry.target, resp)

    def _drain_errqueue(self):
        while True:
            try:
                _data, ancdata, _flags, address = self.sock.recvmsg(
                    2048, 1024, MSG_ERRQUEUE)
            except (BlockingIOError, InterruptedError):
                return
            except OSError:
                return
            ip, status = parse_icmp_error(ancdata, address)
            if ip and status:
                self._record_icmp(ip, status)

    # -- result recording ---------------------------------------------------
    def _deliver(self, record):
        """Hand a result to on_result, isolating callback faults.

        Results arrive from both the receive path and the reaper, so an
        unguarded raise would kill one or the other and abort the scan.
        Ordinary exceptions are counted (the first logged with its traceback)
        and the scan continues; BaseException -- CancelledError,
        KeyboardInterrupt, SystemExit -- is left to propagate.
        """
        try:
            self.on_result(record)
        except Exception:
            self.callback_errors += 1
            if self.callback_errors == 1:
                print(
                    'warning: on_result callback raised (subsequent errors '
                    'are counted only):',
                    file=sys.stderr)
                traceback.print_exc()

    def _record(self, target, resp):
        if target.ip in self.responders or target.ip in self.unreachable:
            return
        result = {
            'ip': target.ip,
            'subnet': target.subnet,
            'status': 'responder',
            'community': target.community,
            'version': VERSION_LABEL.get(resp.version, str(resp.version)),
            'error_status': resp.error_status,
        }
        for oid, value in resp.varbinds:
            # OCTET STRING / Opaque decode to bytes; the system-group scalars we
            # collect (sysDescr, sysName) are DisplayString, so render them as
            # text for the (JSON-friendly) result dict.
            if isinstance(value, (bytes, bytearray)):
                value = value.decode('latin-1')
            result[OID_TO_FIELD.get(oid, oid)] = value
        self.responders[target.ip] = result
        self._deliver(result)

    def _record_icmp(self, ip, status):
        # Only terminal statuses mean "not an SNMP responder; stop probing this
        # host". Transient ones (ttl-exceeded, unknown icmp types) must not
        # resolve the host -- otherwise a routing blip would cut its community
        # fallback short or suppress a reply still in flight.
        if status not in CLOSED_STATUSES:
            return
        if ip in self.responders or ip in self.unreachable:
            return
        self.unreachable[ip] = status
        if self.emit_icmp:
            label = None
            if self.label_for:
                # Like on_result, label_for is foreign code on the loop thread;
                # a fault must not abort the scan -- fall back to no label.
                try:
                    label = self.label_for(ip)
                except Exception:
                    label = None
            self._deliver({'ip': ip, 'subnet': label, 'status': status})

    # -- send helpers -------------------------------------------------------
    def _next_rid(self):
        self._rid = (self._rid + 1) & 0x7FFFFFFF
        if self._rid == 0:
            self._rid = 1
        return self._rid

    async def _send(self, packet, ip):
        drained = False
        while True:
            try:
                self.sock.sendto(packet, (ip, self.config.port))
                return True
            except BlockingIOError:
                await asyncio.sleep(0.001)
            except OSError:
                # On an IP_RECVERR socket a sendto can surface a queued ICMP
                # error from a PRIOR datagram to a different host -- it is not
                # attributable to this target. Drain the error queue once
                # (recording it against the right host) and retry before giving
                # up, so an unrelated reachable host is not dropped.
                if self.icmp_enabled and not drained:
                    drained = True
                    self._drain_errqueue()
                    continue
                self.send_errors += 1
                return False

    def _version_for(self, attempt):
        """SNMP version for a 1-based attempt number.

        Every attempt uses primary_version except, when a fallback is set, the
        final one (retries + 1) -- the v1-fallback piggyback. Fallback requires
        retries >= 1 so the primary version always gets at least one clean try.
        """
        if (self.fallback_version is not None and attempt == self.config.retries + 1):
            return self.fallback_version
        return self.primary_version

    def _reap(self, now):
        """Service arrived replies, then retransmit or drop timed-out requests.

        Draining first is deliberate: a reply that was sitting unserviced in
        the socket gets recorded (its host marked found) before we would time
        it out or fall through to the v1 attempt -- so servicing lag never
        causes a spurious timeout or version downgrade.
        """
        if self.icmp_enabled:
            self._drain_errqueue()
        self._drain_data()
        expired = [rid for rid, e in self.pending.items() if e.deadline <= now]
        for rid in expired:
            e = self.pending.pop(rid)
            if e.target.ip in self.responders or e.target.ip in self.unreachable:
                continue  # already resolved; don't retransmit
            if e.tries > self.config.retries:
                continue  # give up (no responder)
            attempt = e.tries + 1
            packet = assemble_get(rid, e.target.community, _SCAN_VARBINDS,
                                  self._version_for(attempt))
            try:
                self.sock.sendto(packet, (e.target.ip, self.config.port))
                self.retransmits += 1
                self.pending[rid] = e._replace(tries=attempt,
                                               deadline=now + self.config.timeout)
            except BlockingIOError:
                self.pending[rid] = e._replace(deadline=now)  # retry next reap
            except OSError:
                # As in _send: a surfaced queued ICMP error (IP_RECVERR) may be
                # from an unrelated host. Drain it and retry the retransmit once
                # before dropping this entry.
                if self.icmp_enabled:
                    self._drain_errqueue()
                    try:
                        self.sock.sendto(packet, (e.target.ip, self.config.port))
                        self.retransmits += 1
                        self.pending[rid] = e._replace(tries=attempt,
                                                       deadline=now +
                                                       self.config.timeout)
                        continue
                    except OSError:
                        pass
                self.send_errors += 1

    # -- main loop ----------------------------------------------------------
    async def scan(self, targets):
        """Probe an iterable of Target, pacing sends to the configured rate.

        Targets whose IP is already resolved (found, or shown closed via ICMP)
        are skipped. This is a best-effort per-host short-circuit for producers
        that group a host's targets together: it drops probes for a host that
        has already answered, but because a host's communities are sent
        back-to-back a later one may still go out before the first reply lands.
        Dedupe by host happens in _record regardless (see Target).
        """
        interval = self.config.burst / self.config.rate
        since_pause = 0
        for target in targets:
            if target.ip in self.responders or target.ip in self.unreachable:
                continue
            # Bound outstanding requests -> bounded memory, gentler on the net.
            while len(self.pending) >= self.config.max_inflight:
                self._reap(time.monotonic())
                await asyncio.sleep(0.001)
            rid = self._next_rid()
            packet = assemble_get(rid, target.community, _SCAN_VARBINDS,
                                  self._version_for(1))
            if await self._send(packet, target.ip):
                self.pending[rid] = Pending(target,
                                            time.monotonic() + self.config.timeout, 1)
                self.sent += 1
            since_pause += 1
            if since_pause >= self.config.burst:
                since_pause = 0
                self._reap(time.monotonic())
                await asyncio.sleep(interval)
        # Drain: keep reaping/awaiting until every request resolves or dies.
        while self.pending:
            self._reap(time.monotonic())
            await asyncio.sleep(0.01)


async def run(config, on_result, targets, label_for=None):
    """Open a Prober, scan the target stream, and close it.

    Returns (Prober, elapsed_seconds).
    """
    loop = asyncio.get_running_loop()
    prober = Prober(config, on_result, label_for)
    start = time.monotonic()
    try:
        # Inside the try so a failure in open() (e.g. bind/add_reader after the
        # socket fd exists) still hits close() and releases the fd.
        prober.open(loop)
        await prober.scan(targets)
    finally:
        prober.close(loop)
    return prober, time.monotonic() - start
