"""Unit tests for preseem.snmp_engine (the shared multiplexing UDP transport).

Drives SnmpEngine against in-process fake agents (create_datagram_endpoint) and
covers request/response matching by request-id + source, timeout/retry, dropping
of stray/duplicate/malformed replies, concurrency, and request-id allocation.
"""
import asyncio
import socket
import unittest

from preseem import snmp_engine as eng
from preseem import snmp_wire as w
from test import _snmp_fakes as fakes


def _uptime_reply(rid):
    # A minimal valid GetResponse echoing the request-id.
    return fakes.build_response(
        rid, [('1.3.6.1.2.1.1.3.0', w.T_TIMETICKS, w.ber_uint(42))])


def _echo(data):
    return _uptime_reply(fakes.parse_request(data).request_id)


class TestEngine(unittest.TestCase):

    def test_request_response_matches_rid(self):

        async def go():
            transport, port = await fakes.start_agent(_echo)
            try:
                async with eng.SnmpEngine() as e:
                    rid = e.next_rid()
                    pkt = w.build_get(rid, 'public', ['1.3.6.1.2.1.1.3.0'])
                    resp = await e.request(pkt, ('127.0.0.1', port), rid,
                                           timeout=1.0)
                    return rid, resp
            finally:
                transport.close()

        rid, resp = asyncio.run(go())
        self.assertIsNotNone(resp)
        self.assertEqual(resp.request_id, rid)

    def test_timeout_then_retries_return_none(self):

        async def go():
            transport, port = await fakes.start_agent(lambda data: None)
            try:
                async with eng.SnmpEngine() as e:
                    rid = e.next_rid()
                    pkt = w.build_get(rid, 'public', ['1.3.6.1.2.1.1.3.0'])
                    resp = await e.request(pkt, ('127.0.0.1', port), rid,
                                           timeout=0.1, retries=2)
                    return resp, e.sent
            finally:
                transport.close()

        resp, sent = asyncio.run(go())
        self.assertIsNone(resp)
        self.assertEqual(sent, 3)  # first attempt + 2 retries

    def test_concurrent_requests_resolved_by_rid(self):

        async def go():
            transport, port = await fakes.start_agent(_echo)
            try:
                async with eng.SnmpEngine() as e:

                    async def one():
                        rid = e.next_rid()
                        pkt = w.build_get(rid, 'public', ['1.3.6.1.2.1.1.3.0'])
                        r = await e.request(pkt, ('127.0.0.1', port), rid,
                                            timeout=1.0)
                        return rid, r.request_id

                    return await asyncio.gather(*[one() for _ in range(50)])
            finally:
                transport.close()

        pairs = asyncio.run(go())
        self.assertEqual(len(pairs), 50)
        for want, got in pairs:
            self.assertEqual(want, got)

    def test_wrong_source_reply_ignored_without_dropping_pending(self):

        async def go():
            async with eng.SnmpEngine() as e:
                port = e.sock.getsockname()[1]
                loop = asyncio.get_running_loop()
                fut = loop.create_future()
                rid = 12345
                e._pending[rid] = (fut, '10.9.9.9')  # a source we won't receive
                sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                sender.sendto(_uptime_reply(rid),
                              ('127.0.0.1', port))  # source is 127.0.0.1
                sender.close()
                await asyncio.sleep(0.05)
                res = (rid in e._pending, fut.done(), e.dropped)
                e._pending.pop(rid, None)
                return res

        pending, done, dropped = asyncio.run(go())
        self.assertTrue(pending)  # entry survived the mismatch
        self.assertFalse(done)  # future not resolved
        self.assertGreaterEqual(dropped, 1)

    def test_duplicate_reply_ignored(self):

        async def go():
            async with eng.SnmpEngine() as e:
                port = e.sock.getsockname()[1]
                loop = asyncio.get_running_loop()
                fut = loop.create_future()
                rid = 777
                e._pending[rid] = (fut, '127.0.0.1')
                pkt = _uptime_reply(rid)
                sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                sender.sendto(pkt, ('127.0.0.1', port))
                await asyncio.sleep(0.05)
                first = fut.done()
                sender.sendto(pkt, ('127.0.0.1', port))  # duplicate, fut done
                await asyncio.sleep(0.05)
                sender.close()
                res = (first, e.dropped)
                e._pending.pop(rid, None)
                return res

        first_done, dropped = asyncio.run(go())
        self.assertTrue(first_done)
        self.assertGreaterEqual(dropped, 1)

    def test_malformed_datagram_ignored(self):

        async def go():
            async with eng.SnmpEngine() as e:
                port = e.sock.getsockname()[1]
                sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                sender.sendto(b'\x30\x05abcde', ('127.0.0.1', port))
                sender.sendto(b'not-snmp-at-all', ('127.0.0.1', port))
                sender.close()
                await asyncio.sleep(0.05)
            return True  # reaching here means _on_readable did not crash

        self.assertTrue(asyncio.run(go()))

    def test_next_rid_skips_in_use_and_zero(self):
        e = eng.SnmpEngine()
        e._rid = 5
        e._pending[6] = (None, '1.2.3.4')
        self.assertEqual(e.next_rid(), 7)  # 6 is in use -> skipped
        e._rid = 0x7FFFFFFF  # wraps to 0 -> must skip 0
        self.assertEqual(e.next_rid(), 1)

    def test_close_tolerates_unopened(self):
        eng.SnmpEngine().close()  # must not raise


if __name__ == '__main__':
    unittest.main()
