#!/usr/bin/env python3

import serial
import threading
import time

# ==========================
# Configuration
# ==========================

PORT = "/dev/ttyACM0"
BAUDRATE = 115200

COMMAND_PREFIX = 0x69

COMMAND_ACK = 0
COMMAND_NACK = 1
LED = 2

DEVICE_ID = 0

# Test parameters
TEST_DURATION = 2.0  # seconds per stage
LOSS_THRESHOLD = 1.0  # percent

# Rates to test (packets/second)
RATES = [
    100,
    200,
    500,
    1000,
    2000,
    4000,
    8000,
    16000,
]

# ==========================
# Statistics
# ==========================

lock = threading.Lock()

tx_packets = 0
tx_bytes = 0

rx_packets = 0
rx_ack = 0
rx_nack = 0
rx_error = 0


# ==========================
# Packet helpers
# ==========================

def calculate_crc(msg: bytes) -> int:
    s = sum(msg) & 0xFF
    return (-s) & 0xFF


def make_packet(command: int, data: bytes = b"") -> bytes:
    pkt = bytearray()

    pkt.append(COMMAND_PREFIX)
    pkt.append(len(data))
    pkt.append(DEVICE_ID)
    pkt.append(command)
    pkt.append(0)

    pkt.extend(data)

    pkt[4] = calculate_crc(pkt[:4] + pkt[5:])
    return bytes(pkt)


def verify_crc(packet: bytes) -> bool:
    crc = packet[4]
    calc = calculate_crc(packet[:4] + packet[5:])
    return crc == calc


def packet_size(buf: bytes):
    if len(buf) < 2:
        return None
    return 5 + buf[1]


# ==========================
# Receiver
# ==========================

def reader(ser):
    global rx_packets, rx_ack, rx_nack, rx_error

    rx = bytearray()

    while True:
        data = ser.read(4096)
        if not data:
            continue

        rx.extend(data)

        while True:
            if not rx:
                break

            # ASCII log output
            if rx[0] != COMMAND_PREFIX:
                nl = rx.find(b"\n")
                if nl == -1:
                    rx.clear()
                    break

                line = rx[: nl + 1]
                del rx[: nl + 1]

                try:
                    print("[LOG]", line.decode().rstrip())
                except UnicodeDecodeError:
                    pass

                continue

            size = packet_size(rx)

            if size is None or len(rx) < size:
                break

            pkt = bytes(rx[:size])
            del rx[:size]

            if not verify_crc(pkt):
                continue

            cmd = pkt[3]

            with lock:
                rx_packets += 1

                if cmd == COMMAND_ACK:
                    rx_ack += 1
                elif cmd == COMMAND_NACK:
                    rx_nack += 1
                elif cmd == COMMAND_ERROR:
                    rx_error += 1


# ==========================
# Benchmark
# ==========================

def run_stage(ser, packet, target_pps):
    global tx_packets, tx_bytes
    global rx_packets, rx_ack, rx_nack, rx_error

    with lock:
        tx_packets = 0
        tx_bytes = 0
        rx_packets = 0
        rx_ack = 0
        rx_nack = 0
        rx_error = 0

    interval = 1.0 / target_pps

    start = time.perf_counter()
    end = start + TEST_DURATION
    next_tx = start

    while True:
        now = time.perf_counter()

        if now >= end:
            break

        if now >= next_tx:
            ser.write(packet)

            with lock:
                tx_packets += 1
                tx_bytes += len(packet)

            next_tx += interval
        else:
            sleep = next_tx - now
            if sleep > 0:
                time.sleep(min(sleep, 0.0005))

    # Allow final ACKs to arrive
    time.sleep(0.25)

    elapsed = time.perf_counter() - start

    with lock:
        tx = tx_packets
        ack = rx_ack
        nack = rx_nack
        err = rx_error
        bytes_sent = tx_bytes

    loss = 0.0
    if tx:
        loss = (tx - ack) / tx * 100.0

    print()
    print("=" * 60)
    print(f"Target Rate : {target_pps:>7} pkt/s")
    print(f"Elapsed     : {elapsed:.3f} s")
    print(f"Sent        : {tx}")
    print(f"ACK         : {ack}")
    print(f"NACK        : {nack}")
    print(f"ERROR       : {err}")
    print(f"Loss        : {loss:.2f}%")
    print(f"Actual TX   : {tx / elapsed:.0f} pkt/s")
    print(f"Throughput  : {bytes_sent / elapsed / 1024:.2f} KiB/s")

    return loss


# ==========================
# Main
# ==========================

def main():
    packet = make_packet(LED)

    print(f"Opening {PORT} @ {BAUDRATE} baud...")
    ser = serial.Serial(PORT, BAUDRATE, timeout=0.01)

    threading.Thread(target=reader, args=(ser,), daemon=True).start()

    print("\nStarting communication benchmark...\n")

    previous_rate = None

    for rate in RATES:
        loss = run_stage(ser, packet, rate)

        if loss > LOSS_THRESHOLD:
            print()
            print(
                f"Link became unreliable (> {LOSS_THRESHOLD:.1f}% loss)."
            )

            if previous_rate is not None:
                print(f"Maximum reliable rate ≈ {previous_rate} pkt/s")

            break

        previous_rate = rate

    else:
        print()
        print("Completed all test stages.")
        print(f"Reliable up to at least {RATES[-1]} pkt/s.")

    ser.close()


if __name__ == "__main__":
    main()