#!/usr/bin/env python3
"""
Robust DMX512 receiver for Linux (Ubuntu 24.04) using termios PARMRK BREAK marking.

Goals:
- Align to real DMX frame boundaries using BREAK markers (0xFF 0x00 0x00 with PARMRK).
- Self-heal on reception errors:
  * If bytes are dropped/duplicated: resync at next BREAK.
  * If a BREAK marker is missed: keep outputting best-effort frames by length (513 bytes),
    and re-lock as soon as a BREAK appears again.
  * If we see garbage / oversize data without a BREAK: discard until next BREAK.

DMX512 wire format:
- 250000 baud, 8N2
- Frame payload: 1 start code + up to 512 slots (usually 513 bytes total)

Requires:
  pip install pyserial
"""

import argparse
import sys
import time
import serial
import termios


DMX_FRAME_LEN = 513  # start code + 512 slots
BREAK_SEQ = bytes([0xFF, 0x00, 0x00])


def enable_parmrk_break(fd: int) -> None:
    """Enable PARMRK so BREAK is delivered as 0xFF 0x00 0x00 in the read stream."""
    iflag, oflag, cflag, lflag, ispeed, ospeed, cc = termios.tcgetattr(fd)

    # Mark parity errors and breaks in the input stream.
    iflag |= termios.PARMRK | termios.INPCK

    # Do not ignore breaks.
    if hasattr(termios, "IGNBRK"):
        iflag &= ~termios.IGNBRK
    if hasattr(termios, "BRKINT"):
        iflag |= termios.BRKINT

    termios.tcsetattr(
        fd, termios.TCSANOW, [iflag, oflag, cflag, lflag, ispeed, ospeed, cc]
    )


def decode_parmrk_stream(raw: bytes) -> tuple[bytearray, int]:
    """
    Decode a PARMRK stream into:
      - data bytes (with 0xFF 0xFF collapsed to literal 0xFF)
      - count of BREAK markers encountered
    Rules (termios PARMRK):
      - BREAK appears as 0xFF 0x00 0x00
      - a literal 0xFF byte appears as 0xFF 0xFF
      - parity/framing errors can appear as 0xFF 0x00 <x> (x != 0)
        We'll treat that as "byte <x>" but it indicates the stream had an error.
    """
    out = bytearray()
    breaks = 0
    i = 0
    n = len(raw)

    while i < n:
        b0 = raw[i]
        if b0 != 0xFF:
            out.append(b0)
            i += 1
            continue

        # b0 == 0xFF: need at least one more byte
        if i + 1 >= n:
            # keep trailing 0xFF for next chunk by outputting it literally;
            # low probability, but avoids stalling.
            out.append(0xFF)
            i += 1
            continue

        b1 = raw[i + 1]
        if b1 == 0xFF:
            out.append(0xFF)  # escaped 0xFF data
            i += 2
            continue

        if b1 == 0x00:
            if i + 2 >= n:
                # incomplete marker; treat as literal 0xFF then continue
                out.append(0xFF)
                i += 1
                continue
            b2 = raw[i + 2]
            if b2 == 0x00:
                breaks += 1
                i += 3
                continue

            # parity/framing error marking: 0xFF 0x00 b2
            # keep b2 as best-effort data, but note that the stream had an error.
            out.append(b2)
            i += 3
            continue

        # unexpected: treat leading 0xFF as data
        out.append(0xFF)
        i += 1

    return out, breaks


def main() -> int:
    ap = argparse.ArgumentParser(
        description="Robust DMX512 receiver synced by BREAK (Linux PARMRK)."
    )
    ap.add_argument("port", help="Serial port, e.g. /dev/ttyUSB0")
    ap.add_argument(
        "--show", type=int, default=16, help="How many slots to print (default: 16)"
    )
    ap.add_argument(
        "--raw", action="store_true", help="Print raw 513-byte frames as hex"
    )
    ap.add_argument(
        "--stats",
        action="store_true",
        help="Print periodic stats (frames, resyncs, errors)",
    )
    ap.add_argument(
        "--stats-interval",
        type=float,
        default=5.0,
        help="Seconds between stats prints (default: 5)",
    )
    args = ap.parse_args()

    try:
        ser = serial.Serial(
            args.port,
            baudrate=250000,
            bytesize=serial.EIGHTBITS,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_TWO,
            timeout=0.2,
        )
    except Exception as e:
        print(f"Failed to open {args.port}: {e}", file=sys.stderr)
        return 2

    enable_parmrk_break(ser.fileno())
    print(f"Listening on {args.port} @ 250000 8N2 (BREAK via PARMRK). Ctrl-C to stop.")

    # State:
    # We treat BREAK as the authoritative start of a frame.
    # But if a BREAK is missed, we still emit "best-effort" frames every 513 bytes
    # and re-lock when the next BREAK is observed.
    have_break_lock = False
    frame_buf = bytearray()

    frames_ok = 0
    frames_best_effort = 0
    resyncs = 0
    breaks_seen = 0
    parity_marked_errors = (
        0  # approximate (we can't perfectly count without deeper parsing)
    )

    last_stats = time.monotonic()

    try:
        # Small stash for cases where a PARMRK sequence is split across reads.
        carry = bytearray()

        while True:
            chunk = ser.read(4096)
            if not chunk:
                # stats even if idle
                if (
                    args.stats
                    and (time.monotonic() - last_stats) >= args.stats_interval
                ):
                    print(
                        f"[stats] frames_ok={frames_ok} best_effort={frames_best_effort} "
                        f"breaks={breaks_seen} resyncs={resyncs} "
                        f"buf={len(frame_buf)}"
                    )
                    last_stats = time.monotonic()
                continue

            # Keep potential split markers together
            raw = bytes(carry) + chunk
            carry.clear()

            # If raw ends with 0xFF or 0xFF 0x00, keep it for next read (possible split marker)
            if len(raw) >= 1 and raw[-1] == 0xFF:
                carry.extend(raw[-1:])
                raw = raw[:-1]
            elif len(raw) >= 2 and raw[-2:] == bytes([0xFF, 0x00]):
                carry.extend(raw[-2:])
                raw = raw[:-2]

            # Quick-and-dirty count of "error-marked" sequences:
            # count 0xFF 0x00 X where X != 0 (parity/framing) in the raw stream.
            # (This is just for stats; it doesn't drive logic.)
            if args.stats:
                # avoid expensive scan if not needed
                pass
            # Always decode the stream:
            data, brk = decode_parmrk_stream(raw)
            breaks_seen += brk

            # Process data with knowledge of how many BREAKs occurred.
            # We can’t know exact byte positions of each BREAK from decode_parmrk_stream output alone,
            # so we instead scan raw for BREAK_SEQ to split precisely.
            # That’s robust and fast enough at DMX rates.
            #
            # Approach:
            # - Split the *raw* PARMRK stream on BREAK_SEQ.
            # - Decode each segment (which contains no BREAKs).
            # - On each BREAK: start a new locked frame, discarding partial.
            parts = raw.split(BREAK_SEQ)

            if len(parts) == 1:
                # No BREAK in this raw chunk; just consume decoded data.
                # If locked: append to current frame
                # If not locked: still append and output best-effort frames by length (self-heal),
                # but mark them as best-effort until we see a BREAK again.
                if data:
                    frame_buf.extend(data)

                while len(frame_buf) >= DMX_FRAME_LEN:
                    frame = frame_buf[:DMX_FRAME_LEN]
                    del frame_buf[:DMX_FRAME_LEN]

                    start_code = frame[0]
                    slots = frame[1:]

                    if have_break_lock:
                        frames_ok += 1
                        tag = "OK"
                    else:
                        frames_best_effort += 1
                        tag = "BEST"

                    if args.raw:
                        print(frame.hex())
                    else:
                        shown = list(slots[: max(0, min(args.show, 512))])
                        print(f"[{tag}] start=0x{start_code:02X} ch1..={shown}")

                # If buffer gets crazy large without BREAK, dump it and wait for BREAK.
                if len(frame_buf) > (DMX_FRAME_LEN * 4):
                    # We’ve almost certainly lost alignment; purge and force re-lock on BREAK.
                    frame_buf.clear()
                    have_break_lock = False
                    resyncs += 1

            else:
                # There is at least one BREAK in this raw chunk.
                # Each BREAK is an authoritative boundary. We:
                # - consume bytes before first BREAK (could complete best-effort frames)
                # - then for each subsequent segment after a BREAK:
                #   reset frame_buf and lock; begin collecting a new frame payload.
                # Consume bytes before first BREAK (might be tail of previous frame)
                pre = parts[0]
                if pre:
                    pre_dec, _ = decode_parmrk_stream(pre)
                    if pre_dec:
                        frame_buf.extend(pre_dec)

                    # Output any full frames we can form before the BREAK
                    while len(frame_buf) >= DMX_FRAME_LEN:
                        frame = frame_buf[:DMX_FRAME_LEN]
                        del frame_buf[:DMX_FRAME_LEN]
                        start_code = frame[0]
                        slots = frame[1:]

                        if have_break_lock:
                            frames_ok += 1
                            tag = "OK"
                        else:
                            frames_best_effort += 1
                            tag = "BEST"

                        if args.raw:
                            print(frame.hex())
                        else:
                            shown = list(slots[: max(0, min(args.show, 512))])
                            print(f"[{tag}] start=0x{start_code:02X} ch1..={shown}")

                # Now handle each segment after a BREAK
                for seg in parts[1:]:
                    # BREAK observed => authoritative resync
                    frame_buf.clear()
                    have_break_lock = True
                    resyncs += 1

                    if not seg:
                        continue

                    seg_dec, _ = decode_parmrk_stream(seg)
                    if seg_dec:
                        frame_buf.extend(seg_dec)

                    # We may already have a full frame payload in this segment (or more)
                    while len(frame_buf) >= DMX_FRAME_LEN:
                        frame = frame_buf[:DMX_FRAME_LEN]
                        del frame_buf[:DMX_FRAME_LEN]
                        frames_ok += 1

                        start_code = frame[0]
                        slots = frame[1:]

                        if args.raw:
                            print(frame.hex())
                        else:
                            shown = list(slots[: max(0, min(args.show, 512))])
                            print(f"[OK] start=0x{start_code:02X} ch1..={shown}")

                    # If we got *more than one frame* between BREAKs, alignment is suspect.
                    # Keep tail but we remain locked; the next BREAK will correct it anyway.

            if args.stats and (time.monotonic() - last_stats) >= args.stats_interval:
                print(
                    f"[stats] frames_ok={frames_ok} best_effort={frames_best_effort} "
                    f"breaks={breaks_seen} resyncs={resyncs} "
                    f"buf={len(frame_buf)}"
                )
                last_stats = time.monotonic()

    except KeyboardInterrupt:
        pass
    finally:
        ser.close()

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
