#!/usr/bin/env python3
"""
UART bridge test with fpga_test/loopback loaded into the FPGA (uart_tx = uart_rx).

Writes random data at several baud rates and checks that exactly the same
bytes come back; prints the sustained full-duplex throughput.

    python tests/uart_loopback_test.py --port COM12
"""
import argparse
import os
import sys
import threading
import time

import serial

BAUDS = [9600, 115200, 460800, 921600, 1000000, 2000000, 3000000, 4000000,
         5000000, 6000000, 8000000, 10000000, 12000000]


def run(port, baud, nbytes, fmt):
    data = os.urandom(nbytes)
    rx = bytearray()
    bytesize, parity, stopbits = fmt
    with serial.Serial(port, baud, bytesize=bytesize, parity=parity,
                       stopbits=stopbits, timeout=0.05) as s:
        if bytesize < 8:
            data = bytes(b & ((1 << bytesize) - 1) for b in data)
        time.sleep(0.05)
        s.reset_input_buffer()
        done = threading.Event()

        def reader():
            idle_since = None
            while len(rx) < len(data):
                chunk = s.read(65536)
                if chunk:
                    rx.extend(chunk)
                    idle_since = None
                elif done.is_set():
                    idle_since = idle_since or time.time()
                    if time.time() - idle_since > 0.5:
                        break

        t = threading.Thread(target=reader)
        t0 = time.time()
        t.start()
        for i in range(0, len(data), 4096):
            s.write(data[i:i + 4096])
        s.flush()
        done.set()
        t.join(timeout=60 + nbytes * 12 / baud)
        dt = time.time() - t0
    ok = bytes(rx) == data
    first_bad = next((i for i, (a, b) in enumerate(zip(rx, data)) if a != b), None)
    return ok, len(rx), dt, first_bad


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--port', required=True)
    ap.add_argument('--bytes', type=int, default=0, help='bytes per test (default: ~1 s of data)')
    ap.add_argument('--baud', type=int, action='append', help='test only these baud rates')
    args = ap.parse_args()

    fails = 0
    for baud in args.baud or BAUDS:
        n = args.bytes or max(4096, min(2_000_000, baud // 10))
        ok, got, dt, bad = run(args.port, baud, n, (8, serial.PARITY_NONE, serial.STOPBITS_ONE))
        rate = got / dt / 1024
        status = 'OK  ' if ok else 'FAIL'
        extra = '' if ok else f' (received {got}/{n}, first mismatch at {bad})'
        print(f'{status} {baud:>9} baud 8N1  {n:>8} bytes  {rate:8.1f} KiB/s{extra}', flush=True)
        fails += not ok

    for fmt, name in [((7, serial.PARITY_EVEN, serial.STOPBITS_ONE), '7E1'),
                      ((8, serial.PARITY_ODD, serial.STOPBITS_TWO), '8O2')]:
        ok, got, dt, bad = run(args.port, 115200, 4096, fmt)
        print(f'{"OK  " if ok else "FAIL"}    115200 baud {name}', flush=True)
        fails += not ok

    print('ALL PASSED' if fails == 0 else f'{fails} FAILED')
    sys.exit(1 if fails else 0)


if __name__ == '__main__':
    main()