#!/usr/bin/env python3
"""Bounded single-connection Linkalyst throughput peer; synthetic payload only.

Default: python3 tools/throughput_peer.py (127.0.0.1:5202).
Use --host with an explicitly selected numeric address for a user-operated LAN peer.
This is a plain TCP test protocol, not HTTP, TLS, iperf, or an Internet speed service.

Request: >4sB3xQI = NST1, direction (0 download / 1 upload), payload bytes, total ms.
Ready: >4sQ = NSR1, requested bytes. Upload receipt: >4sQ = NSC1, bytes received.
Only payload contributes to byte counts; no payload is saved or forwarded.
"""

import argparse
import ipaddress
import socket
import struct
import time

MAX_BYTES = 64 * 1024 * 1024
MAX_SECONDS = 60
CHUNK_BYTES = 64 * 1024
PAYLOAD = b"\xa5" * CHUNK_BYTES
REQUEST = struct.Struct(">4sB3xQI")
RESPONSE = struct.Struct(">4sQ")


def set_remaining_timeout(connection, deadline):
    remaining = deadline - time.monotonic()
    if remaining <= 0:
        raise TimeoutError("total request deadline reached")
    connection.settimeout(remaining)


def receive_exact(connection, count, deadline):
    data = bytearray()
    while len(data) < count:
        set_remaining_timeout(connection, deadline)
        chunk = connection.recv(count - len(data))
        if not chunk:
            raise ConnectionError("connection closed before complete frame")
        data.extend(chunk)
    return bytes(data)


def send_all(connection, data, deadline):
    set_remaining_timeout(connection, deadline)
    connection.sendall(data)


def serve_one(connection, accepted_at=None):
    """Serve exactly one bounded request; caller owns and closes the socket.

    Return (direction, payload_bytes) only on a complete transfer. An upload receipt
    is sent only after all requested bytes have actually reached this process.
    """
    started = time.monotonic() if accepted_at is None else accepted_at
    deadline = started + MAX_SECONDS
    raw = receive_exact(connection, REQUEST.size, deadline)
    magic, direction, count, milliseconds = REQUEST.unpack(raw)
    if (magic != b"NST1" or direction not in (0, 1) or raw[5:8] != b"\0\0\0"
            or not 1 <= count <= MAX_BYTES or not 1 <= milliseconds <= 60_000):
        raise ValueError("invalid protocol request or byte/time limit")
    deadline = min(deadline, started + milliseconds / 1000)
    send_all(connection, RESPONSE.pack(b"NSR1", count), deadline)
    transferred = 0
    while transferred < count:
        size = min(CHUNK_BYTES, count - transferred)
        if direction == 0:
            send_all(connection, PAYLOAD[:size], deadline)
            transferred += size
        else:
            set_remaining_timeout(connection, deadline)
            chunk = connection.recv(size)
            if not chunk:
                raise ConnectionError(f"upload closed after {transferred} of {count} payload bytes")
            transferred += len(chunk)
    if direction == 1:
        send_all(connection, RESPONSE.pack(b"NSC1", transferred), deadline)
    return ("download" if direction == 0 else "upload", transferred)


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--host", default="127.0.0.1", help="numeric bind address (default: loopback only)")
    parser.add_argument("--port", type=int, default=5202, help="TCP port (default: 5202; 0 chooses a local free port)")
    parser.add_argument("--once", action="store_true", help="exit after one connection, successful or otherwise")
    args = parser.parse_args()
    try:
        address = ipaddress.ip_address(args.host)
    except ValueError:
        parser.error("--host must be a numeric IPv4 or IPv6 address")
    if not 0 <= args.port <= 65535:
        parser.error("--port must be between 0 and 65535")
    family = socket.AF_INET6 if address.version == 6 else socket.AF_INET
    with socket.socket(family, socket.SOCK_STREAM) as listener:
        listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        listener.bind((args.host, args.port))
        listener.listen(1)
        print(f"Listening on {args.host}:{listener.getsockname()[1]} (single connection; 64 MiB / 60 s maximum)", flush=True)
        while True:
            connection, peer = listener.accept()
            accepted_at = time.monotonic()
            with connection:
                connection.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
                try:
                    direction, count = serve_one(connection, accepted_at)
                    print(f"{peer[0]} {direction}: {count} payload bytes complete", flush=True)
                except (OSError, ValueError) as error:
                    print(f"{peer[0]} request ended: {error}", flush=True)
            if args.once:
                break


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass