"""Shared async SNMP transport: one UDP socket multiplexing requests by id.

A single unconnected UDP socket carries many concurrent SNMP request/response
exchanges to arbitrary hosts, matched back to their callers by request-id and
verified by source address. Unlike the fan-out scanner in snmp_probe, this
engine imposes no in-flight bound, burst, or rate limit -- concurrency is the
caller's to decide (e.g. asyncio.gather over many bulkwalk() coroutines). It is
operation-agnostic: any PDU the codec can build (GET, GETBULK, ...) rides on the
same socket, so higher-level operations share one transport. The socket is
intentionally separate from the probe's.

ICMP/IP_RECVERR error attribution is deliberately not wired in here: an
unreachable host simply times out, and callers that need to tell "dead" from
"slow" issue their own liveness probe over the engine (see snmp_bulkwalk).
"""
import asyncio
import socket

from preseem.snmp_wire import decode_response


class SnmpEngine:
    """One shared UDP socket; async request/response keyed by request-id.

    Open it once (``open`` / ``async with``), then call ``request`` from as many
    concurrent coroutines as desired. Replies are matched to the awaiting caller
    by request-id and source IP.

    Decoded replies carry OIDs as tuples of integer sub-identifiers (SNMP's
    native form), so walk/table consumers can order and slice them without
    re-parsing dotted strings; stringify with '.'.join(map(str, oid)) if needed.
    OCTET STRING / Opaque values come back as raw bytes -- the caller decodes
    text or formats a MAC as it sees fit -- since an arbitrary subtree walk sees
    binary values as often as display text.
    """

    def __init__(self):
        self.sock = None
        self._loop = None
        self._pending = {}  # request_id -> (Future, expected_ip)
        self._rid = 0
        self.sent = 0  # datagrams handed to the socket
        self.received = 0  # replies matched to a pending request
        self.dropped = 0  # replies matching no live/right-source request
        # Reused receive buffer: recvfrom_into fills it in place, so the hot
        # receive path allocates no per-datagram bytes object. decode_response
        # parses a memoryview of it and fully materialises its results, so the
        # buffer is safe to overwrite on the next datagram.
        self._recvbuf = bytearray(65535)
        self._recvmv = memoryview(self._recvbuf)

    # -- lifecycle ----------------------------------------------------------
    def open(self, loop=None):
        self._loop = loop or asyncio.get_event_loop()
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            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.
            self.sock.bind(('0.0.0.0', 0))
            self._loop.add_reader(self.sock.fileno(), self._on_readable)
        except BaseException:
            # Release the fd if we fail after creating the socket -- __aexit__
            # is not called when __aenter__ raises.
            self.close()
            raise

    def close(self):
        if self.sock is None:  # never opened / already closed
            return
        try:
            self._loop.remove_reader(self.sock.fileno())
        except (OSError, ValueError):
            pass
        self.sock.close()
        self.sock = None
        # Cancel anything still awaiting so callers don't hang after close.
        for fut, _ip in self._pending.values():
            if not fut.done():
                fut.cancel()
        self._pending.clear()

    async def __aenter__(self):
        self.open()
        return self

    async def __aexit__(self, *exc):
        self.close()

    # -- request-id allocation ---------------------------------------------
    def next_rid(self):
        """Allocate a request-id: monotonic (wraps 31-bit), never 0, and never
        one currently in flight, so concurrent requests can't collide."""
        while True:
            self._rid = (self._rid + 1) & 0x7FFFFFFF
            if self._rid and self._rid not in self._pending:
                return self._rid

    # -- receive path -------------------------------------------------------
    def _on_readable(self):
        sock, buf, mv, pending = (self.sock, self._recvbuf, self._recvmv, self._pending)
        while True:
            try:
                nbytes, addr = sock.recvfrom_into(buf)
            except (BlockingIOError, InterruptedError):
                return
            except ConnectionRefusedError:
                # Unconnected socket: a queued ICMP error can surface here. We
                # don't attribute it (no IP_RECVERR); keep draining.
                continue
            except OSError:
                return
            try:
                resp = decode_response(mv[:nbytes], oid_as_ints=True)
            except (IndexError, ValueError):
                continue  # malformed datagram -> ignore, not fatal
            if resp is None:
                continue
            entry = pending.get(resp.request_id)
            if entry is None:
                self.dropped += 1
                continue
            fut, expected_ip = entry
            # Peek, don't pop: a wrong-source or duplicate/late reply must not
            # resolve or remove the real pending entry (request() pops it).
            if addr[0] != expected_ip or fut.done():
                self.dropped += 1
                continue
            self.received += 1
            fut.set_result(resp)

    # -- send path ----------------------------------------------------------
    async def _send(self, packet, addr):
        while True:
            try:
                self.sock.sendto(packet, addr)
                self.sent += 1
                return
            except BlockingIOError:
                await asyncio.sleep(0.001)  # send buffer full; back off briefly
            except OSError:
                # Unattributable transport error; let the timeout/retry resend.
                return

    async def request(self, packet, addr, rid, *, timeout, retries=0):
        """Send ``packet`` (whose request-id is ``rid``) to ``addr`` and await
        the matching reply.

        Returns the decoded SnmpResponse, or None if no reply arrives within
        ``timeout`` across ``retries`` resends (retries + 1 attempts total).
        Resends reuse the same packet and request-id; a late reply to an earlier
        attempt still satisfies the current wait (same rid, same content).
        """
        expected_ip = addr[0]
        try:
            for _attempt in range(retries + 1):
                fut = self._loop.create_future()
                self._pending[rid] = (fut, expected_ip)
                await self._send(packet, addr)
                try:
                    return await asyncio.wait_for(fut, timeout)
                except asyncio.TimeoutError:
                    continue  # resend (a fresh future is armed next iteration)
            return None
        finally:
            self._pending.pop(rid, None)
