"""T3R binary serial protocol — host-side codec.

A faithful Python port of the wire protocol implemented in ``main/protocol.c``.
Pure standard library (no PyQt / pyserial), so it can be unit-tested on its own.

Frame layout (all multi-byte fields little-endian):

    +------+------+------+----------------+------+
    | SOF  | CMD  | LEN  |  PAYLOAD[LEN]   | CRC8 |
    +------+------+------+----------------+------+
      0xA5   1 B    1 B       LEN bytes      1 B

CRC8 is CRC-8/SMBUS (poly 0x07, init 0x00) over CMD, LEN and PAYLOAD.
See PROTOCOL.md for the full catalogue.
"""

from __future__ import annotations

import struct
from dataclasses import dataclass

SOF = 0xA5
PROTO_VERSION = 1

# ---- Requests (host -> device) --------------------------------------------
CMD_PING = 0x01
CMD_SET_MICROSTEP = 0x10
CMD_SET_CURRENT = 0x11
CMD_ENABLE = 0x12
CMD_DISABLE = 0x13
CMD_ENABLE_MASK = 0x14
CMD_MOVE = 0x20
CMD_JOG = 0x21
CMD_STOP = 0x22
CMD_STOP_ALL = 0x23
CMD_MOVE_GROUP = 0x24
CMD_JOG_GROUP = 0x25
CMD_GET_INFO = 0x30
CMD_GET_DRV_STATUS = 0x31
CMD_GET_POSITION = 0x32
CMD_SET_POSITION = 0x33
CMD_READ_REG = 0x40
CMD_WRITE_REG = 0x41

# ---- Responses (device -> host) -------------------------------------------
RSP_ACK = 0x81
RSP_PONG = 0x82
RSP_INFO = 0x83
RSP_DRV_STATUS = 0x84
RSP_POSITION = 0x85
RSP_REG = 0x86

# ---- Asynchronous events (device -> host) ---------------------------------
EVT_MOTION_DONE = 0xE0
EVT_STOPPED = 0xE1
EVT_FAULT = 0xE2

# Human-readable names, used by the log/console.
CMD_NAMES = {
    CMD_PING: "PING", CMD_SET_MICROSTEP: "SET_MICROSTEP",
    CMD_SET_CURRENT: "SET_CURRENT", CMD_ENABLE: "ENABLE", CMD_DISABLE: "DISABLE",
    CMD_ENABLE_MASK: "ENABLE_MASK",
    CMD_MOVE: "MOVE", CMD_JOG: "JOG", CMD_STOP: "STOP", CMD_STOP_ALL: "STOP_ALL",
    CMD_MOVE_GROUP: "MOVE_GROUP", CMD_JOG_GROUP: "JOG_GROUP",
    CMD_GET_INFO: "GET_INFO", CMD_GET_DRV_STATUS: "GET_DRV_STATUS",
    CMD_GET_POSITION: "GET_POSITION", CMD_SET_POSITION: "SET_POSITION",
    CMD_READ_REG: "READ_REG", CMD_WRITE_REG: "WRITE_REG",
    RSP_ACK: "ACK", RSP_PONG: "PONG", RSP_INFO: "INFO",
    RSP_DRV_STATUS: "DRV_STATUS", RSP_POSITION: "POSITION", RSP_REG: "REG",
    EVT_MOTION_DONE: "MOTION_DONE", EVT_STOPPED: "STOPPED", EVT_FAULT: "FAULT",
}

# ---- ACK status codes ------------------------------------------------------
STATUS_NAMES = {
    0x00: "OK", 0x02: "bad length", 0x03: "bad channel", 0x04: "bad parameter",
    0x05: "busy", 0x06: "SPI comms fault", 0x07: "unknown command",
}

# ---- INFO.state ------------------------------------------------------------
STATE_NAMES = {0: "IDLE", 1: "MOVING", 2: "JOGGING", 3: "STOPPING"}

# ---- Fault mask bits -------------------------------------------------------
FAULT_BITS = [
    (1 << 0, "SHORT_GND_A"),
    (1 << 1, "SHORT_GND_B"),
    (1 << 2, "OVERTEMP"),
    (1 << 3, "OVERTEMP_WARN"),
    (1 << 4, "OPEN_LOAD_A"),
    (1 << 5, "OPEN_LOAD_B"),
]

MICROSTEPS = [1, 2, 4, 8, 16, 32, 64, 128, 256]
MAX_VELOCITY = 200_000      # steps/s  (MOTOR_MAX_VELOCITY)
MAX_ACCEL = 2_000_000       # steps/s2 (MOTOR_MAX_ACCEL)
NUM_CHANNELS = 4


def fault_names(mask: int) -> str:
    """Render a fault mask as a comma-separated list, or 'none'."""
    names = [name for bit, name in FAULT_BITS if mask & bit]
    return ", ".join(names) if names else "none"


# ---------------------------------------------------------------------------
# CRC-8 (poly 0x07, init 0x00) — matches the firmware reference.
# ---------------------------------------------------------------------------

def crc8(data: bytes) -> int:
    crc = 0
    for byte in data:
        crc ^= byte
        for _ in range(8):
            crc = ((crc << 1) ^ 0x07) & 0xFF if crc & 0x80 else (crc << 1) & 0xFF
    return crc


def build_frame(cmd: int, payload: bytes = b"") -> bytes:
    body = bytes([cmd, len(payload)]) + payload
    return bytes([SOF]) + body + bytes([crc8(body)])


# ---------------------------------------------------------------------------
# Request encoders — each returns a complete frame ready to write.
# ---------------------------------------------------------------------------

def ping() -> bytes:
    return build_frame(CMD_PING)


def set_microstep(ch: int, microsteps: int) -> bytes:
    return build_frame(CMD_SET_MICROSTEP, struct.pack("<BH", ch, microsteps))


def set_current(ch: int, run_ma: int, hold_ma: int, ihold_delay: int) -> bytes:
    return build_frame(CMD_SET_CURRENT,
                       struct.pack("<BHHB", ch, run_ma, hold_ma, ihold_delay))


def enable(ch: int) -> bytes:
    return build_frame(CMD_ENABLE, bytes([ch]))


def disable(ch: int) -> bytes:
    return build_frame(CMD_DISABLE, bytes([ch]))


def enable_mask(mask: int) -> bytes:
    """Energise exactly the channels in `mask` (bit i = channel i); 0 = all off."""
    return build_frame(CMD_ENABLE_MASK, bytes([mask & 0x0F]))


def move(ch: int, steps: int, velocity: int, accel: int) -> bytes:
    return build_frame(CMD_MOVE, struct.pack("<BiII", ch, steps, velocity, accel))


def jog(ch: int, velocity: int, accel: int) -> bytes:
    return build_frame(CMD_JOG, struct.pack("<BiI", ch, velocity, accel))


def move_group(mask: int, steps: int, velocity: int, accel: int) -> bytes:
    """Ganged move: every channel in `mask` steps together in lockstep."""
    return build_frame(CMD_MOVE_GROUP,
                       struct.pack("<BiII", mask & 0x0F, steps, velocity, accel))


def jog_group(mask: int, velocity: int, accel: int) -> bytes:
    """Ganged jog: every channel in `mask` runs together (velocity 0 = stop)."""
    return build_frame(CMD_JOG_GROUP,
                       struct.pack("<BiI", mask & 0x0F, velocity, accel))


def stop(ch: int, hard: bool) -> bytes:
    return build_frame(CMD_STOP, struct.pack("<BB", ch, 1 if hard else 0))


def stop_all() -> bytes:
    return build_frame(CMD_STOP_ALL)


def get_info(ch: int) -> bytes:
    return build_frame(CMD_GET_INFO, bytes([ch]))


def get_drv_status(ch: int) -> bytes:
    return build_frame(CMD_GET_DRV_STATUS, bytes([ch]))


def get_position(ch: int) -> bytes:
    return build_frame(CMD_GET_POSITION, bytes([ch]))


def set_position(ch: int, position: int) -> bytes:
    return build_frame(CMD_SET_POSITION, struct.pack("<Bi", ch, position))


# ---------------------------------------------------------------------------
# Response / event decoders. Each returns a dataclass (or None on bad length).
# ---------------------------------------------------------------------------

@dataclass
class Pong:
    proto_ver: int
    fw_version: int
    num_channels: int


@dataclass
class Ack:
    req_cmd: int
    status: int

    @property
    def ok(self) -> bool:
        return self.status == 0x00


@dataclass
class Info:
    ch: int
    state: int
    position: int
    velocity: int
    microsteps: int
    run_ma: int
    hold_ma: int
    enabled: bool
    comms_ok: bool
    fault_mask: int


@dataclass
class DrvStatus:
    ch: int
    raw: int
    fault_mask: int

    @property
    def standstill(self) -> bool:
        return bool(self.raw & (1 << 31))

    @property
    def cs_actual(self) -> int:
        return (self.raw >> 16) & 0x1F

    @property
    def sg_result(self) -> int:
        return self.raw & 0x3FF


@dataclass
class Position:
    ch: int
    position: int


def decode_pong(p: bytes):
    if len(p) < 4:
        return None
    proto, fw, nch = struct.unpack_from("<BHB", p, 0)
    return Pong(proto, fw, nch)


def decode_ack(p: bytes):
    if len(p) < 2:
        return None
    return Ack(p[0], p[1])


def decode_info(p: bytes):
    if len(p) < 18:
        return None
    (ch, state, pos, vel, micro, run_ma, hold_ma, flags, fault) = \
        struct.unpack_from("<BBiIHHHBB", p, 0)
    return Info(ch, state, pos, vel, micro, run_ma, hold_ma,
                bool(flags & 0x01), bool(flags & 0x02), fault)


def decode_drv_status(p: bytes):
    if len(p) < 6:
        return None
    ch, raw, fault = struct.unpack_from("<BIB", p, 0)
    return DrvStatus(ch, raw, fault)


def decode_position(p: bytes):
    if len(p) < 5:
        return None
    ch, pos = struct.unpack_from("<Bi", p, 0)
    return Position(ch, pos)


def decode_event_position(p: bytes):
    """MOTION_DONE / STOPPED share the (ch, position) layout."""
    return decode_position(p)


def decode_fault(p: bytes):
    if len(p) < 2:
        return None
    return Position(p[0], p[1])  # reuse (ch, value); value is the fault mask


# ---------------------------------------------------------------------------
# Incremental frame parser — mirrors the firmware byte-wise state machine.
# ---------------------------------------------------------------------------

class FrameParser:
    """Feed raw bytes, get back a list of (cmd, payload) for each valid frame.

    Bad-CRC and oversized frames are dropped silently; the parser resyncs on
    the next SOF, exactly like the device-side parser.
    """

    MAX_PAYLOAD = 64

    _SOF, _CMD, _LEN, _PAYLOAD, _CRC = range(5)

    def __init__(self):
        self._state = self._SOF
        self._cmd = 0
        self._len = 0
        self._payload = bytearray()

    def feed(self, data: bytes):
        out = []
        for b in data:
            if self._state == self._SOF:
                if b == SOF:
                    self._state = self._CMD
            elif self._state == self._CMD:
                self._cmd = b
                self._state = self._LEN
            elif self._state == self._LEN:
                self._len = b
                self._payload.clear()
                if b > self.MAX_PAYLOAD:
                    self._state = self._SOF        # too big: resync
                elif b == 0:
                    self._state = self._CRC
                else:
                    self._state = self._PAYLOAD
            elif self._state == self._PAYLOAD:
                self._payload.append(b)
                if len(self._payload) >= self._len:
                    self._state = self._CRC
            elif self._state == self._CRC:
                body = bytes([self._cmd, self._len]) + bytes(self._payload)
                if crc8(body) == b:
                    out.append((self._cmd, bytes(self._payload)))
                # else bad CRC: drop, resync
                self._state = self._SOF
        return out