import asyncio
import os
import struct


HOST = os.getenv("STRESS_TARGET_HOST", "0.0.0.0")
PORT = int(os.getenv("STRESS_TARGET_PORT", "18080"))
MAX_FRAME_BYTES = 1_048_576


async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    try:
        while True:
            header = await reader.readexactly(4)
            size = struct.unpack("!I", header)[0]
            if size <= 0 or size > MAX_FRAME_BYTES:
                raise ValueError(f"invalid frame size: {size}")
            payload = await reader.readexactly(size)
            writer.write(struct.pack("!I", len(payload)) + payload)
            await writer.drain()
    except (asyncio.IncompleteReadError, ConnectionError, ValueError):
        pass
    finally:
        writer.close()
        try:
            await writer.wait_closed()
        except ConnectionError:
            pass


async def main() -> None:
    server = await asyncio.start_server(handle, HOST, PORT, backlog=4096, limit=2 * MAX_FRAME_BYTES)
    addresses = ", ".join(str(sock.getsockname()) for sock in server.sockets or [])
    print(f"TCP stress target listening on {addresses}", flush=True)
    async with server:
        await server.serve_forever()


if __name__ == "__main__":
    asyncio.run(main())