#!/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_DURATION = 5.0
ACK_TIMEOUT = 1.0

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

def calculate_crc(msg: bytes) -> int:
    return (-sum(msg)) & 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:
    return packet[4] == calculate_crc(packet[:4] + packet[5:])


# ==========================================================
# ACK synchronization
# ==========================================================

ack_event = threading.Event()

ack_count = 0
nack_count = 0
error_count = 0
timeout_count = 0

running = True


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

def reader(ser):
    global ack_count
    global nack_count
    global error_count

    rx = bytearray()

    while running:

        data = ser.read(4096)

        if not data:
            continue

        rx.extend(data)

        while True:

            if len(rx) < 5:
                break

            if rx[0] != COMMAND_PREFIX:

                nl = rx.find(b"\n")

                if nl == -1:
                    rx.clear()
                    break

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

                del rx[: nl + 1]
                continue

            length = rx[1]
            size = 5 + length

            if len(rx) < size:
                break

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

            if not verify_crc(pkt):
                continue

            cmd = pkt[3]

            if cmd == COMMAND_ACK:
                ack_count += 1
                ack_event.set()

            elif cmd == COMMAND_NACK:
                nack_count += 1
                ack_event.set()

            elif cmd == COMMAND_ERROR:
                error_count += 1
                ack_event.set()


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

def benchmark(ser):

    global timeout_count

    packet = make_packet(LED)

    sent = 0

    rtt_sum = 0.0
    rtt_min = float("inf")
    rtt_max = 0.0

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

    while time.perf_counter() < end:

        ack_event.clear()

        t0 = time.perf_counter()

        ser.write(packet)

        if not ack_event.wait(ACK_TIMEOUT):
            timeout_count += 1
            continue

        t1 = time.perf_counter()

        rtt = t1 - t0

        sent += 1
        rtt_sum += rtt
        rtt_min = min(rtt_min, rtt)
        rtt_max = max(rtt_max, rtt)

    elapsed = time.perf_counter() - start

    print()
    print("=" * 60)
    print(f"Elapsed      : {elapsed:.3f} s")
    print(f"Sent         : {sent}")
    print(f"ACK          : {ack_count}")
    print(f"NACK         : {nack_count}")
    print(f"ERROR        : {error_count}")
    print(f"Timeouts     : {timeout_count}")
    print(f"Packet Rate  : {sent / elapsed:.1f} pkt/s")

    if sent:
        print(f"Mean RTT     : {1000*rtt_sum/sent:.3f} ms")
        print(f"Min RTT      : {1000*rtt_min:.3f} ms")
        print(f"Max RTT      : {1000*rtt_max:.3f} ms")


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

def main():

    global running

    print(f"Opening {PORT}")

    ser = serial.Serial(
        PORT,
        BAUDRATE,
        timeout=0.01,
    )

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

    print("Running benchmark...")

    benchmark(ser)

    running = False

    ser.close()


if __name__ == "__main__":
    main()