import asyncio
import json
import subprocess
from urllib.parse import parse_qs, urlsplit

async def echo(reader, writer):
    try:
        while message := await reader.readline():
            print(f"CHAT {writer.get_extra_info('peername')} {message.decode().strip()}", flush=True)
            writer.write(message)
            await writer.drain()
    except (ConnectionError, asyncio.LimitOverrunError):
        pass
    finally:
        writer.close()
        await writer.wait_closed()

class SensorSink(asyncio.DatagramProtocol):
    def datagram_received(self, data, addr):
        try:
            print(f"IOT {addr} {json.loads(data)}", flush=True)
        except (ValueError, UnicodeDecodeError):
            pass

class ContinuityEcho(asyncio.DatagramProtocol):
    def connection_made(self, transport):
        self.transport = transport

    def datagram_received(self, data, addr):
        if len(data) <= 1024 and data.startswith(b"mobility:"):
            self.transport.sendto(data, addr)

async def stream(reader, writer):
    try:
        request = await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), timeout=10)
        target = request.split(b" ")[1].decode()
        query = parse_qs(urlsplit(target).query)
        mbps = min(20, max(0.1, float(query.get("mbps", ["4"])[0])))
        duration = min(600, max(1, int(query.get("duration", ["180"])[0])))
        writer.write(b"HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream\r\nConnection: close\r\n\r\n")
        loop = asyncio.get_running_loop()
        start = loop.time()
        chunk = b"5G" * 4096
        period = len(chunk) * 8 / (mbps * 1e6)
        next_send = start
        while loop.time() - start < duration:
            writer.write(chunk)
            await writer.drain()
            next_send = max(next_send + period, loop.time())
            await asyncio.sleep(max(0, next_send - loop.time()))
    except (ConnectionError, ValueError, TimeoutError, asyncio.IncompleteReadError):
        pass
    finally:
        writer.close()
        try:
            await writer.wait_closed()
        except ConnectionError:
            pass

async def main():
    iperf = subprocess.Popen(["iperf3", "-s"])
    loop = asyncio.get_running_loop()
    transport, _ = await loop.create_datagram_endpoint(SensorSink, local_addr=("0.0.0.0", 9001))
    continuity, _ = await loop.create_datagram_endpoint(ContinuityEcho, local_addr=("0.0.0.0", 9003))
    servers = [await asyncio.start_server(echo, "0.0.0.0", 9000), await asyncio.start_server(stream, "0.0.0.0", 8080)]
    print("N6 server ready: chat :9000, IoT :9001, stream :8080, iperf3 :5201", flush=True)
    try:
        await asyncio.gather(*(server.serve_forever() for server in servers))
    finally:
        transport.close()
        continuity.close()
        iperf.terminate()

asyncio.run(main())