#!/usr/bin/env python3
"""Send binary commands to a device over a serial port, optionally receiving and/or requiring an ACK.

Packet layout matches:

	struct command_message_t {
		uint8_t prefix;
		uint8_t length;
		uint8_t id;
		uint8_t command;
		int64_t tick;      // milliseconds, little-endian
		uint8_t crc;
		uint8_t data[COMMAND_DATA_SIZE];
	} __attribute__((packed));
"""

import argparse
import struct
import sys
import threading
import time

import serial

DEFAULT_PORT = "/dev/ttyACM0"
DEFAULT_BAUD = 115200
DEFAULT_DEVICE_ID = 0

COMMAND_PREFIX = 0x69

COMMAND_ACK = 0
COMMAND_NACK = 1

# name -> command id
COMMANDS = {
	"ACK": 0,
	"NACK": 1,
	"LED_TOGGLE": 2,
	"LED_SET": 3,
	"SERVO_SET": 4,
	"SERVO_SET_ALL": 5,
	"ADC_READ_RAW": 6,
	"ADC_READ_RAW_ALL": 7,
	"ADC_READ": 8,
	"ADC_READ_ALL": 9,
	"ADC_SET_READ": 10,
	"DIGITAL_IN": 11,
	"DIGITAL_LEG_OUT": 12,
}
COMMAND_NAMES = {v: k for k, v in COMMANDS.items()}

# Commands for which the data argument is OPTIONAL.
# Every command not listed here REQUIRES data.
# >>> Edit this list to match your firmware. <<<
COMMANDS_DATA_OPTIONAL = {
	"ACK",
	"NACK",
	"LED_TOGGLE",
	"ADC_READ_RAW_ALL",
	"ADC_READ_ALL",
}

# struct formats for typed data tokens (all little-endian)
TYPE_FORMATS = {
	"u8": "<B", "i8": "<b",
	"u16": "<H", "i16": "<h",
	"u32": "<I", "i32": "<i",
	"f32": "<f",
}

# Exit codes
EXIT_OK = 0
EXIT_ERROR = 1
EXIT_NO_ACK = 2
EXIT_NACK = 3

# Fixed header layout: prefix, length, id, command, tick(int64), 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

# Safety valve for the ASCII-log branch: a real firmware log line shouldn't
# be longer than this. If we've buffered more than this with no newline,
# we're desynced (likely sitting inside binary packet data), so drop a byte
# and retry instead of blocking forever waiting for a '\n' that won't come.
MAX_LOG_LINE = 256


# --------------------------------------------------------------------------
# Packet helpers
# --------------------------------------------------------------------------

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


def make_packet(command: int, data: bytes = b"", device_id: int = DEFAULT_DEVICE_ID,
                 tick: int = 0) -> bytes:
	pkt = bytearray()
	pkt.append(COMMAND_PREFIX)
	pkt.append(len(data))
	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]


# --------------------------------------------------------------------------
# Receiver thread
# --------------------------------------------------------------------------

class Receiver(threading.Thread):
	"""Parses incoming packets / ASCII logs.

	`verbose` gates whether anything is printed at all. `mode` picks the
	print format for received command packets:
	  normal    - name, tick, length, then the raw hex line (default)
	  debug     - the full raw hex bitstream of every packet/line, nothing else
	  data      - only the data payload bytes
	  timestamp - tick + data payload bytes ("-l")

	Stats for --telemetry are always collected, regardless of `verbose`.
	"""

	def __init__(self, ser, verbose=False, mode="normal"):
		super().__init__(daemon=True)
		self.ser = ser
		self.verbose = verbose
		self.mode = mode
		self.stop_event = threading.Event()
		self.ack_event = threading.Event()
		self.result = None  # "ACK" or "NACK" once received

		# telemetry
		self.start_time = None
		self.end_time = None
		self.command_counts = {}
		self.ack_count = 0
		self.nack_count = 0
		self.bad_crc_count = 0
		self.log_line_count = 0
		self.data_byte_count = 0
		self.packet_count = 0
		# device-clock (tick, ms) timestamps of every CRC-valid packet, in
		# arrival order, used to compute real RX frequency. Kept per name
		# and combined across all types.
		self.tick_history = {}
		self.all_ticks = []

	def _print(self, *args):
		if self.verbose:
			print(*args)

	def _print_packet(self, pkt, length, dev_id, cmd, tick, payload):
		name = COMMAND_NAMES.get(cmd, str(cmd))

		if self.mode == "debug":
			self._print(f"<-- RAW: {pkt.hex(' ')}")
		elif self.mode == "data":
			self._print(f"<-- DATA: {payload.hex(' ')}")
		elif self.mode == "timestamp":
			self._print(f"<-- t={tick}ms  DATA: {payload.hex(' ')}")
		else:  # normal
			self._print(f"<-- {name} (device={dev_id}, tick={tick}ms) len={length}")
			self._print(f"<-- RX: {pkt.hex(' ')}")

	def _print_ack_nack(self, pkt, dev_id, tick, name):
		if self.mode == "debug":
			self._print(f"<-- RAW: {pkt.hex(' ')}")
		elif self.mode == "data":
			pass  # no data payload to show
		elif self.mode == "timestamp":
			self._print(f"<-- t={tick}ms  {name}")
		else:  # normal
			self._print(f"<-- {name} (device={dev_id}, tick={tick}ms)")

	def run(self):
		self.start_time = time.time()
		rx = bytearray()

		while not self.stop_event.is_set():
			data = self.ser.read(64)
			if not data:
				continue

			rx.extend(data)

			while rx:
				if rx[0] == COMMAND_PREFIX:
					size = packet_size(rx)
					if size is None or len(rx) < size:
						break

					pkt = bytes(rx[:size])

					if not verify_crc(pkt):
						self.bad_crc_count += 1
						# The CRC is what's wrong, not necessarily the tick
						# field's position, so still surface it - useful for
						# spotting exactly this kind of firmware bug.
						bad_tick = unpack_tick(pkt)
						if self.mode == "debug":
							self._print(f"RX: Bad CRC: {pkt.hex(' ')}")
						else:
							self._print(f"RX: Bad CRC (tick={bad_tick}ms): {pkt.hex(' ')}")
						# The guessed size was based on a length byte that
						# might itself be garbage (stale/misaligned bytes).
						# Drop just the leading prefix byte and re-scan,
						# rather than eating `size` bytes and risking a
						# cascade of misalignment.
						del rx[:1]
						continue

					del rx[:size]
					length = pkt[1]
					dev_id = pkt[2]
					cmd = pkt[3]
					tick = unpack_tick(pkt)
					payload = pkt[HEADER_SIZE:HEADER_SIZE + length]

					self.packet_count += 1

					if cmd == COMMAND_ACK:
						self.ack_count += 1
						self.result = "ACK"
						self._print_ack_nack(pkt, dev_id, tick, "ACK")
						self.tick_history.setdefault("ACK", []).append(tick)
						self.all_ticks.append(tick)
						self.ack_event.set()
					elif cmd == COMMAND_NACK:
						self.nack_count += 1
						self.result = "NACK"
						self._print_ack_nack(pkt, dev_id, tick, "NACK")
						self.tick_history.setdefault("NACK", []).append(tick)
						self.all_ticks.append(tick)
						self.ack_event.set()
					else:
						name = COMMAND_NAMES.get(cmd, str(cmd))
						self.command_counts[name] = self.command_counts.get(name, 0) + 1
						self.data_byte_count += length
						self._print_packet(pkt, length, dev_id, cmd, tick, payload)
						self.tick_history.setdefault(name, []).append(tick)
						self.all_ticks.append(tick)
				else:
					idx = rx.find(b"\n")
					if idx == -1:
						if len(rx) > MAX_LOG_LINE:
							# Not really a log line: desynced, likely sitting
							# inside binary packet data. Shed a byte and
							# keep looking for the next prefix instead of
							# stalling forever waiting for '\n'.
							del rx[:1]
							continue
						break

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

					self.log_line_count += 1

					if self.mode == "data":
						continue  # data-only stream: skip firmware log text

					try:
						decoded = line.decode().rstrip()
						if self.mode == "debug":
							self._print(f"[RAW] {line.hex(' ')}")
						else:
							self._print("[LOG]", decoded)
					except UnicodeDecodeError:
						self._print("[RAW]", line.hex())

		self.end_time = time.time()

	def print_telemetry(self):
		start = self.start_time or time.time()
		end = self.end_time or time.time()
		elapsed = max(end - start, 1e-9)
		total = self.packet_count + self.log_line_count

		def hz(ticks):
			"""Frequency in Hz from device-clock (tick, ms) deltas between
			the first and last sample. None if there aren't enough points
			or the ticks didn't advance (e.g. all identical / rolled over)."""
			if len(ticks) < 2:
				return None
			span_ms = ticks[-1] - ticks[0]
			if span_ms <= 0:
				return None
			return (len(ticks) - 1) * 1000.0 / span_ms

		print("--- Telemetry ---")
		print(f"elapsed:        {elapsed:.3f} s")
		print(f"packets:        {self.packet_count} ({self.packet_count / elapsed:.2f} pkt/s host-side)")
		print(f"data bytes:     {self.data_byte_count} ({self.data_byte_count / elapsed:.2f} B/s)")
		print(f"log lines:      {self.log_line_count}")
		print(f"ACK / NACK:     {self.ack_count} / {self.nack_count}")
		print(f"bad CRC:        {self.bad_crc_count}")
		print(f"total frames:   {total} ({total / elapsed:.2f} frames/s)")

		overall_hz = hz(self.all_ticks)
		if overall_hz is not None:
			print(f"RX frequency:   {overall_hz:.2f} Hz (from device tick deltas, all valid packets)")
		else:
			print("RX frequency:   n/a (need >=2 CRC-valid packets with advancing tick)")

		if self.command_counts:
			print("by command (count, Hz from tick deltas):")
			for name, count in sorted(self.command_counts.items()):
				f = hz(self.tick_history.get(name, []))
				f_str = f"{f:.2f} Hz" if f is not None else "n/a"
				print(f"  {name:<18} {count:>6}   {f_str}")


# --------------------------------------------------------------------------
# Argument parsing
# --------------------------------------------------------------------------

def parse_command(text: str):
	"""Accept a command name (case-insensitive) or a number. Returns (id, name)."""
	key = text.upper()
	if key in COMMANDS:
		return COMMANDS[key], key

	try:
		val = int(text, 0)
	except ValueError:
		raise argparse.ArgumentTypeError(
			f"unknown command '{text}' (use --list-commands to see them)"
		)
	if not 0 <= val <= 255:
		raise argparse.ArgumentTypeError("command number must be 0-255")
	return val, COMMAND_NAMES.get(val, f"#{val}")


def parse_data_token(tok: str) -> bytes:
	"""
	Formats:
	  5            -> u8
	  u16:500      -> unsigned 16-bit little-endian
	  i32:-3       -> signed 32-bit little-endian
	  f32:1.5      -> float32 little-endian
	  hex:0a0b0c   -> raw bytes
	"""
	if ":" in tok:
		kind, value = tok.split(":", 1)
		kind = kind.lower()
	else:
		kind, value = "u8", tok

	if kind == "hex":
		return bytes.fromhex(value)

	fmt = TYPE_FORMATS.get(kind)
	if fmt is None:
		raise ValueError(f"unknown type '{kind}' (use u8/i8/u16/i16/u32/i32/f32/hex)")

	num = float(value) if kind == "f32" else int(value, 0)
	return struct.pack(fmt, num)


def parse_tick(text: str) -> int:
	if text.lower() == "now":
		return int(time.time() * 1000)
	try:
		return int(text, 0)
	except ValueError:
		raise argparse.ArgumentTypeError("tick must be an integer (ms) or 'now'")


def build_parser():
	commands_help = ", ".join(f"{name}={cid}" for name, cid in COMMANDS.items())
	optional_help = ", ".join(sorted(COMMANDS_DATA_OPTIONAL)) or "(none)"

	p = argparse.ArgumentParser(
		description="Send a binary command to a serial device.",
		formatter_class=argparse.RawDescriptionHelpFormatter,
		epilog=f"""\
commands:
  {commands_help}

commands where data is optional:
  {optional_help}

data tokens (all multi-byte values are little-endian):
  5           u8 (default type)
  u16:500     also: u8 i8 u16 i16 u32 i32 f32
  hex:0a0b    raw bytes

RX display modes (pick at most one; require --rx, --debug, --data or -l to receive at all):
  (none)      name, tick, length, and raw hex per packet
  --debug     raw hex bitstream of every packet/line, nothing decoded
  --data      only the data payload bytes of received commands
  -l          tick timestamp + data payload bytes

examples:
  %(prog)s LED_TOGGLE
  %(prog)s ADC_SET_READ u32:1 --require-ack
  %(prog)s SERVO_SET 0 u16:1500 --rx --rx-time 5
  %(prog)s ADC_READ_ALL -l --telemetry
  %(prog)s -p /dev/ttyUSB0 -b 9600 ADC_READ_ALL --debug
""",
	)

	p.add_argument("command", nargs="?", type=parse_command,
	               help="command name or number")
	p.add_argument("data", nargs="*", help="command data tokens (see below)")

	p.add_argument("-p", "--port", default=DEFAULT_PORT,
	               help="serial port (default: %(default)s)")
	p.add_argument("-b", "--baud", type=int, default=DEFAULT_BAUD,
	               help="baud rate (default: %(default)s)")
	p.add_argument("-d", "--device-id", type=lambda s: int(s, 0), default=DEFAULT_DEVICE_ID,
	               help="device id in the packet header (default: %(default)s)")
	p.add_argument("-t", "--tick", type=parse_tick, default=0,
	               help="tick value in ms for the outgoing packet header, or 'now' (default: %(default)s)")

	p.add_argument("--rx", action="store_true",
	               help="receive and print incoming packets/logs in the default format (default: off)")
	p.add_argument("--rx-time", type=float, default=10.0,
	               help="seconds to keep listening after sending, when receiving (default: %(default)s)")

	rx_format = p.add_mutually_exclusive_group()
	rx_format.add_argument("--debug", action="store_true",
	                        help="RX mode: print the raw hex bitstream of every packet/line (default: off)")
	rx_format.add_argument("--data", dest="data_only", action="store_true",
	                        help="RX mode: print only the data payload of received commands (default: off)")
	rx_format.add_argument("-l", "--timestamp", action="store_true",
	                        help="RX mode: print data payload with its tick timestamp (default: off)")

	p.add_argument("--require-ack", action="store_true",
	               help="wait for an ACK; exit 2 on timeout, 3 on NACK (default: off)")
	p.add_argument("--ack-timeout", type=float, default=1.0,
	               help="seconds to wait for ACK (default: %(default)s)")

	p.add_argument("--boot-delay", type=float, default=0.0,
	               help="seconds to wait after opening the port before flushing/sending, "
	                    "for boards that reset on connect (default: %(default)s, off)")

	p.add_argument("--telemetry", action="store_true",
	               help="print an RX telemetry summary (packet/byte rates, counts) at exit (default: off)")

	p.add_argument("--list-commands", action="store_true",
	               help="list known commands and exit")
	return p


# --------------------------------------------------------------------------
# Main
# --------------------------------------------------------------------------

def main():
	parser = build_parser()
	args = parser.parse_args()

	if args.list_commands:
		for name, cid in COMMANDS.items():
			opt = "data optional" if name in COMMANDS_DATA_OPTIONAL else "data required"
			print(f"{cid:3d}  {name:<18} {opt}")
		return EXIT_OK

	if args.command is None:
		parser.error("the following arguments are required: command")

	cmd_id, cmd_name = args.command

	if not args.data and cmd_name not in COMMANDS_DATA_OPTIONAL:
		parser.error(f"command {cmd_name} requires data")

	data = bytearray()
	try:
		for tok in args.data:
			data.extend(parse_data_token(tok))
	except (ValueError, struct.error) as e:
		parser.error(f"bad data: {e}")

	if len(data) > 255:
		parser.error(f"data too long ({len(data)} bytes, max 255)")

	packet = make_packet(cmd_id, bytes(data), args.device_id, args.tick)

	# --debug/--data/-l/--telemetry imply listening for the rx-time window,
	# same as --rx, even if --rx itself wasn't passed.
	listen = args.rx or args.debug or args.data_only or args.timestamp or args.telemetry
	need_receiver = listen or args.require_ack

	if args.debug:
		mode = "debug"
	elif args.data_only:
		mode = "data"
	elif args.timestamp:
		mode = "timestamp"
	else:
		mode = "normal"

	try:
		ser = serial.Serial(args.port, args.baud, timeout=0.05)
		if args.boot_delay > 0:
			time.sleep(args.boot_delay)  # let a board that resets on connect finish booting
		ser.reset_input_buffer()  # discard stale/boot-garbage bytes before we start
	except serial.SerialException as e:
		print(f"ERROR: could not open {args.port}: {e}", file=sys.stderr)
		return EXIT_ERROR

	receiver = None
	exit_code = EXIT_OK

	try:
		# Start the receiver BEFORE sending so a fast ACK isn't missed.
		if need_receiver:
			receiver = Receiver(ser, verbose=listen, mode=mode)
			receiver.start()

		print(f"--> TX {cmd_name} (tick={args.tick}ms): {packet.hex(' ')}")
		ser.write(packet)
		ser.flush()

		if args.require_ack:
			if not receiver.ack_event.wait(args.ack_timeout):
				print(f"ERROR: no ACK within {args.ack_timeout}s", file=sys.stderr)
				exit_code = EXIT_NO_ACK
			elif receiver.result == "NACK":
				print("ERROR: device replied NACK", file=sys.stderr)
				exit_code = EXIT_NACK
			else:
				print("ACK received.")

		if listen:
			time.sleep(args.rx_time)

	except KeyboardInterrupt:
		pass
	finally:
		if receiver:
			receiver.stop_event.set()
			receiver.join(timeout=1)
		ser.close()

	if args.telemetry and receiver:
		receiver.print_telemetry()

	print("Done.")
	return exit_code


if __name__ == "__main__":
	sys.exit(main())