#!/usr/bin/env python3
"""Exercise the io_uring fixed-buffer/SEND_ZC echo path with exact data."""
import socket
import sys
import threading
import time

port = int(sys.argv[1])
total = int(sys.argv[2]) if len(sys.argv) > 2 else 64 * 1024 * 1024
connections = int(sys.argv[3]) if len(sys.argv) > 3 else 1
chunk = bytes((i * 31 + 7) & 0xFF for i in range(128 * 1024))

start = time.monotonic()
errors = []

def connection_worker(connection_bytes):
    sock = socket.create_connection(("127.0.0.1", port), timeout=5)
    send_error = []

    def sender():
        try:
            sent = 0
            while sent < connection_bytes:
                payload = chunk[:min(len(chunk), connection_bytes - sent)]
                sock.sendall(payload)
                sent += len(payload)
            sock.shutdown(socket.SHUT_WR)
        except Exception as error:
            send_error.append(error)

    thread = threading.Thread(target=sender)
    thread.start()
    try:
        done = 0
        while done < connection_bytes:
            part = sock.recv(min(1024 * 1024, connection_bytes - done))
            if not part:
                raise RuntimeError("echo server closed early")
            offset = 0
            while offset < len(part):
                pattern_offset = (done + offset) % len(chunk)
                count = min(len(part) - offset, len(chunk) - pattern_offset)
                if part[offset:offset + count] != chunk[pattern_offset:pattern_offset + count]:
                    raise RuntimeError("echo content mismatch")
                offset += count
            done += len(part)
        thread.join()
        if send_error:
            raise send_error[0]
    except Exception as error:
        errors.append(error)
    finally:
        sock.close()

per_connection = total // connections
workers = [
    threading.Thread(target=connection_worker, args=(per_connection,))
    for _ in range(connections)
]
for worker in workers:
    worker.start()
for worker in workers:
    worker.join()
if errors:
    raise errors[0]
elapsed = time.monotonic() - start
print("ECHO %.2f Gbit/s (%d connections)" %
      (per_connection * connections * 8 / elapsed / 1e9, connections),
      flush=True)