#!/usr/bin/env python3

import serial
import struct
import threading
import time

PORT = "/dev/ttyACM0"
BAUDRATE = 115200

COMMAND_PREFIX = 0x69

COMMAND_ACK = 0
COMMAND_NACK = 1
LED_TOGGLE = 2
LED_SET = 3
SERVO_SET = 4
SERVO_SET_ALL = 5
DIGITAL_LEG_OUT = 12

DEVICE_ID = 0

# Header layout: prefix, length, id, command, tick(int64 LE), crc
TICK_FORMAT = "<q"                         # signed 64-bit little-endian, milliseconds
TICK_SIZE = struct.calcsize(TICK_FORMAT)   # 8
CRC_OFFSET = 4 + TICK_SIZE                 # index of the crc byte -> 12
HEADER_SIZE = CRC_OFFSET + 1               # bytes before data -> 13


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


def make_packet(command: int, data: bytes = b"", tick: int = 0) -> bytes:
	length = len(data)

	pkt = bytearray()
	pkt.append(COMMAND_PREFIX)
	pkt.append(length)
	pkt.append(DEVICE_ID)
	pkt.append(command)
	pkt.extend(struct.pack(TICK_FORMAT, tick))
	pkt.append(0)			  # CRC placeholder
	pkt.extend(data)

	pkt[CRC_OFFSET] = calculate_crc(pkt[:CRC_OFFSET] + pkt[CRC_OFFSET + 1:])

	return bytes(pkt)


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


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


def unpack_tick(packet: bytes) -> int:
	return struct.unpack(TICK_FORMAT, packet[4:4 + TICK_SIZE])[0]


def reader(ser):
	rx = bytearray()

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

		rx.extend(data)

		while rx:
			# Binary packet?
			if rx[0] == COMMAND_PREFIX:
				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):
					print("RX: Bad CRC:", pkt.hex())
					continue

				length = pkt[1]
				dev_id = pkt[2]
				cmd = pkt[3]
				tick = unpack_tick(pkt)

				if cmd == COMMAND_ACK:
					print(f"<-- ACK (device={dev_id}, tick={tick}ms)")

				elif cmd == COMMAND_NACK:
					print(f"<-- NACK (device={dev_id}, tick={tick}ms)")

				else:
					print(f"<-- Command {cmd} (tick={tick}ms) len={length}")

			else:
				# ASCII log output
				idx = rx.find(b'\n')
				if idx == -1:
					break

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

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


def main():
	ser = serial.Serial(PORT, BAUDRATE, timeout=0.05)

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

	for i in range(0x3f):
		data = struct.pack("<B", i)  # explicitly uint8_t, raises if i is out of 0-255 range

		ser.write(make_packet(DIGITAL_LEG_OUT, data, tick=0))
		time.sleep(0.1)

	print("Done.")
	time.sleep(2)


if __name__ == "__main__":
	main()