"""Asyncio activation probe. Use --source 127.0.0.1 only for local tests."""

import argparse
import asyncio
import socket

MAX_PAYLOAD = 1228
CHALLENGE = b"SHREDCAST-VERIFY"


class Receiver(asyncio.DatagramProtocol):
    def __init__(self, source: str):
        self.source = source
        self.packets = 0
        self.total_bytes = 0

    def datagram_received(self, data: bytes, peer: tuple[str, int]) -> None:
        if peer[0] != self.source or len(data) > MAX_PAYLOAD:
            return
        if data.startswith(CHALLENGE):
            print(f"verification={data.decode('utf-8', errors='replace')}", flush=True)
            return
        self.packets += 1
        self.total_bytes += len(data)

    def report(self) -> None:
        print(f"stream_packets={self.packets} stream_bytes={self.total_bytes}", flush=True)
        self.packets = self.total_bytes = 0
        asyncio.get_running_loop().call_later(1.0, self.report)


async def main(port: int, source: str) -> None:
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 8 * 1024 * 1024)
    sock.bind(("0.0.0.0", port))
    sock.setblocking(False)
    loop = asyncio.get_running_loop()
    transport, protocol = await loop.create_datagram_endpoint(
        lambda: Receiver(source), sock=sock
    )
    print(f"listening on udp://0.0.0.0:{port}; source={source}", flush=True)
    loop.call_later(1.0, protocol.report)
    try:
        await asyncio.Event().wait()
    finally:
        transport.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=9000)
    parser.add_argument("--source", default="203.23.178.228")
    args = parser.parse_args()
    try:
        asyncio.run(main(args.port, args.source))
    except KeyboardInterrupt:
        pass