#!/usr/bin/env python3
"""
Test script for app_photomagnetic communication protocol.

Frame format (SimpleProtocal, per DTS config):
    H0(0x7E) H1(0xE7) CMD(1B) LEN(1B) DATA[0..N] CRC_LO CRC_HI
    CRC = CRC-16 Modbus over CMD+LEN+DATA, LSB-MSB order

Usage:
    python test.py /dev/ttyUSB0 --get-id              (request device ID)
    python test.py /dev/ttyUSB0 --state 1              (set running state)
    python test.py /dev/ttyUSB0 --monitor              (listen for temp data)
    python test.py /dev/ttyUSB0 --baud 115200 --loop   (poll ID in a loop)
"""

import argparse
import struct
import sys
import time

try:
    import serial
except ImportError:
    print("Please install pyserial: pip install pyserial", file=sys.stderr)
    sys.exit(1)

# ── Protocol constants ────────────────────────────────────────
HEADER = b"\x7e\xe7"  # 2-byte header (DTS: header = <0x7EE7>)
HEADER_SIZE = len(HEADER)

CMD_TEMP = 0x00
CMD_GET_ID = 0x01
CMD_RUNNING_STATE = 0x02

STATE_NAMES = {
    0x00: "Standby",
    0x01: "Running",
    0x02: "Pause",
    0x03: "Fault",
    0x04: "Upgrade",
    0x05: "EStop",
    0x06: "Exception",
}


# ── CRC-16 Modbus ─────────────────────────────────────────────
def crc16_modbus(data: bytes) -> int:
    """CRC-16 Modbus (polynomial 0xA001, init 0xFFFF)."""
    crc = 0xFFFF
    for byte in data:
        crc ^= byte
        for _ in range(8):
            if crc & 1:
                crc = (crc >> 1) ^ 0xA001
            else:
                crc >>= 1
    return crc & 0xFFFF


# ── Frame helpers ─────────────────────────────────────────────
def build_frame(cmd: int, payload: bytes = b"") -> bytes:
    """Build a SimpleProtocal frame with LSB-MSB CRC.
    Frame: H0 H1 CMD LEN DATA[0..N] CRC_LO CRC_HI"""
    body = bytes([cmd, len(payload)]) + payload
    crc = crc16_modbus(body)
    crc_lo = crc & 0xFF
    crc_hi = (crc >> 8) & 0xFF
    return HEADER + body + bytes([crc_lo, crc_hi])


def parse_frame(frame: bytes) -> tuple[int, bytes] | None:
    """Parse a SimpleProtocal frame with 2-byte header + LSB-MSB CRC.
    Returns (cmd, data) or None on error."""
    min_len = HEADER_SIZE + 2 + 2  # HDR(2) + CMD(1) + LEN(1) + CRC(2) = 6
    if len(frame) < min_len:
        return None
    if frame[:HEADER_SIZE] != HEADER:
        return None

    cmd = frame[HEADER_SIZE]
    length = frame[HEADER_SIZE + 1]
    expected = HEADER_SIZE + 2 + length + 2  # HDR + CMD+LEN + DATA + CRC
    if len(frame) != expected:
        return None

    data = frame[HEADER_SIZE + 2 : HEADER_SIZE + 2 + length]

    # CRC covers CMD + LEN + DATA (LSB-MSB order on wire)
    crc_body = frame[HEADER_SIZE : HEADER_SIZE + 2 + length]
    actual_crc = crc16_modbus(crc_body)
    crc_lo = frame[HEADER_SIZE + 2 + length]
    crc_hi = frame[HEADER_SIZE + 2 + length + 1]
    wire_crc = (crc_hi << 8) | crc_lo

    if actual_crc != wire_crc:
        return None

    return cmd, data


# ── Serial reader ─────────────────────────────────────────────
class ProtocolReader:
    """State-machine frame reader for 2-byte header protocol."""

    def __init__(self, ser: serial.Serial):
        self._ser = ser
        self._buf = bytearray()

    def read_frame(self, timeout: float = 1.0) -> bytes | None:
        """Read one complete frame. Returns raw frame bytes or None on timeout."""
        min_frame = HEADER_SIZE + 2 + 2  # HDR + CMD + LEN + CRC = 6

        deadline = time.time() + timeout
        while time.time() < deadline:
            # Look for header
            if len(self._buf) >= 2:
                idx = self._buf.find(HEADER)
                if idx > 0:
                    del self._buf[:idx]  # discard bytes before header
                elif idx < 0 and len(self._buf) > 0:
                    self._buf.clear()

            # Have a header → try to parse frame length
            if len(self._buf) >= HEADER_SIZE:
                if self._buf[:HEADER_SIZE] == HEADER:
                    if len(self._buf) >= min_frame:
                        length = self._buf[HEADER_SIZE + 1]
                        total = HEADER_SIZE + 2 + length + 2
                        if len(self._buf) >= total:
                            frame = bytes(self._buf[:total])
                            del self._buf[:total]
                            return frame
                else:
                    # Wrong header, discard first byte and retry
                    del self._buf[0]
                    continue

            # Read more bytes
            try:
                chunk = self._ser.read(self._ser.in_waiting or 1)
                if chunk:
                    self._buf.extend(chunk)
            except (serial.SerialTimeoutException, serial.SerialException):
                pass

        return None


# ── Command handlers ──────────────────────────────────────────
def cmd_get_id(ser: serial.Serial, timeout: float = 2.0):
    """Send GET_ID request and print response."""
    print("→ Sending GET_ID...")
    ser.write(build_frame(CMD_GET_ID))

    reader = ProtocolReader(ser)
    frame = reader.read_frame(timeout)
    if frame is None:
        print("✗ No response (timeout)")
        return

    result = parse_frame(frame)
    if result is None:
        print(f"✗ Invalid frame: {frame.hex()}")
        return

    cmd, data = result
    if cmd != CMD_GET_ID:
        print(f"✗ Unexpected cmd=0x{cmd:02X} (expected 0x{CMD_GET_ID:02X})")
        return

    print(f"✓ Device ID ({len(data)} bytes): {data.hex(' ')}")


def cmd_set_state(ser: serial.Serial, state: int):
    """Send RUNNING_STATE command."""
    name = STATE_NAMES.get(state, "Unknown")
    print(f"→ Setting state: {name} (0x{state:02X})")
    ser.write(build_frame(CMD_RUNNING_STATE, bytes([state])))
    print(f"✓ Sent")


def cmd_monitor(ser: serial.Serial):
    """Listen for incoming frames (TEMP data) and print them."""
    print("Monitoring... (Ctrl+C to stop)")
    reader = ProtocolReader(ser)
    try:
        while True:
            frame = reader.read_frame(timeout=0.5)
            if frame is None:
                continue

            result = parse_frame(frame)
            if result is None:
                print(f"⚠ Bad frame: {frame.hex()}")
                continue

            cmd, data = result
            if cmd == CMD_TEMP:
                if len(data) >= 2:
                    # sensor_value: val1 (int8), val2 (int8)
                    t = data[0]
                    frac = data[1] / 100.0
                    temp = t + frac
                    print(f"🌡  Temp: {temp:.2f} °C  (raw: {data.hex(' ')})")
                else:
                    print(f"⚠ TEMP with invalid data: {data.hex()}")
            elif cmd == CMD_GET_ID:
                print(f"🆔 Device ID response: {data.hex(' ')}")
            elif cmd == CMD_RUNNING_STATE:
                print(f"🔁 Running state echo: 0x{data.hex()}")
            else:
                print(f"📦 Unknown cmd=0x{cmd:02X} data={data.hex()}")
    except KeyboardInterrupt:
        print("\nDone.")


def cmd_loop(ser: serial.Serial, count: int = 0):
    """Send GET_ID in a loop."""
    i = 0
    print(f"Looping GET_ID... (Ctrl+C to stop)")
    try:
        while count == 0 or i < count:
            cmd_get_id(ser, timeout=1.0)
            i += 1
            time.sleep(0.5)
    except KeyboardInterrupt:
        print("\nDone.")


# ── Main ──────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(
        description="Test script for app_photomagnetic communication protocol"
    )
    parser.add_argument("port", help="Serial port (e.g. /dev/ttyUSB0, COM3)")
    parser.add_argument(
        "--baud", type=int, default=115200, help="Baud rate (default: 115200)"
    )
    parser.add_argument(
        "--timeout", type=float, default=0.1, help="Serial timeout in seconds"
    )
    parser.add_argument(
        "--get-id", action="store_true", help="Send GET_ID and print response"
    )
    parser.add_argument(
        "--state",
        type=lambda x: int(x, 0),
        metavar="N",
        help="Set running state (0=Standby, 1=Running, 2=Pause, ...)",
    )
    parser.add_argument(
        "--monitor", action="store_true", help="Listen for incoming frames"
    )
    parser.add_argument(
        "--loop",
        type=int,
        nargs="?",
        const=0,
        metavar="N",
        help="Send GET_ID in a loop (N times, or infinite if omitted)",
    )
    parser.add_argument(
        "--raw",
        type=lambda x: bytes.fromhex(x),
        metavar="HEX",
        help="Send raw payload (CMD DATA..., CRC auto-appended)",
    )
    parser.add_argument("--debug", action="store_true", help="Show raw frame hex")

    args = parser.parse_args()

    ser = serial.Serial(args.port, args.baud, timeout=args.timeout)
    print(f"Connected to {args.port} @ {args.baud} baud")

    try:
        if args.get_id:
            cmd_get_id(ser)
        elif args.state is not None:
            cmd_set_state(ser, args.state)
        elif args.monitor:
            cmd_monitor(ser)
        elif args.loop is not None:
            cmd_loop(ser, args.loop)
        elif args.raw is not None:
            payload = args.raw
            cmd, data = payload[0], payload[1:] if len(payload) > 1 else b""
            print(f"→ Sending CMD=0x{cmd:02X} data={data.hex()}")
            frame = build_frame(cmd, data)
            print(f"  Frame: {frame.hex()}")
            ser.write(frame)
            print(f"✓ Sent")
        else:
            print("No action. Use --get-id, --state N, --monitor, --loop, or --raw")
    finally:
        ser.close()


if __name__ == "__main__":
    main()