#!/usr/bin/env python3
"""RS485 master and end-to-end test suite for the wiredsensor protocol.

This is deliberately an **independent** implementation of the wire format. The
framing and CRC below were written from the specification in README.md rather
than sharing code with `wiredsensor-core`. If both ends shared an
implementation, a bug in it would cancel out and every test here would pass
regardless — so the duplication is the point, not an oversight.

Talks to the bus through the USB-to-RS485 bridge firmware (`bridge/`), which is
protocol-agnostic and just moves bytes.

    ./wiredsensor.py --port /dev/ttyACM0 measure
    ./wiredsensor.py --port /dev/ttyACM0 monitor
    ./wiredsensor.py --port /dev/ttyACM0 test
"""

from __future__ import annotations

import argparse
import struct
import sys
import time
from dataclasses import dataclass

try:
    import serial
except ImportError:
    sys.exit("pyserial is required: pip install pyserial")

# ---------------------------------------------------------------- protocol ---

ADDR_BROADCAST = 0x00
ADDR_MIN, ADDR_MAX = 0x01, 0xF7

MAX_PAYLOAD = 64
HEADER_LEN = 3
CRC_LEN = 2
MIN_FRAME = HEADER_LEN + CRC_LEN
MAX_FRAME = HEADER_LEN + MAX_PAYLOAD + CRC_LEN

CMD_READ_MEASUREMENT = 0x01
CMD_READ_INFO = 0x02
CMD_READ_STATUS = 0x03
CMD_PING = 0x04
CMD_SET_ADDRESS = 0x10

ERROR_FLAG = 0x80

ERR_NAMES = {
    0x01: "ILLEGAL_COMMAND",
    0x02: "ILLEGAL_LENGTH",
    0x03: "SENSOR_UNAVAILABLE",
    0x04: "SENSOR_FAULT",
    0x05: "UNSUPPORTED",
}

FLAG_NAMES = [
    (1 << 0, "SENSOR_OK"),
    (1 << 1, "DATA_STALE"),
    (1 << 2, "SENSOR_FAULT"),
    (1 << 3, "UART_ERROR"),
    (1 << 4, "EVER_MEASURED"),
    (1 << 5, "SENSOR_RESET_SEEN"),
    (1 << 6, "HEATER_ON"),
]


def crc16(data: bytes) -> int:
    """CRC-16/MODBUS: reflected poly 0xA001, init 0xFFFF, no final XOR."""
    crc = 0xFFFF
    for byte in data:
        crc ^= byte
        for _ in range(8):
            crc = (crc >> 1) ^ 0xA001 if crc & 1 else crc >> 1
    return crc


def build_frame(addr: int, cmd: int, payload: bytes = b"") -> bytes:
    if len(payload) > MAX_PAYLOAD:
        raise ValueError(f"payload of {len(payload)} exceeds {MAX_PAYLOAD}")
    body = bytes([addr, cmd, len(payload)]) + payload
    return body + struct.pack("<H", crc16(body))


@dataclass
class Frame:
    addr: int
    cmd: int
    payload: bytes

    @property
    def is_error(self) -> bool:
        return bool(self.cmd & ERROR_FLAG)

    @property
    def error_name(self) -> str | None:
        if not self.is_error or len(self.payload) != 1:
            return None
        return ERR_NAMES.get(self.payload[0], f"unknown({self.payload[0]:#04x})")


class ProtocolError(Exception):
    pass


def parse_frame(raw: bytes) -> Frame:
    if len(raw) < MIN_FRAME:
        raise ProtocolError(f"short frame: {len(raw)} bytes ({raw.hex(' ')})")
    if len(raw) > MAX_FRAME:
        raise ProtocolError(f"long frame: {len(raw)} bytes")
    declared = raw[2]
    if HEADER_LEN + declared + CRC_LEN != len(raw):
        raise ProtocolError(
            f"LEN field {declared} disagrees with {len(raw)} received bytes "
            f"({raw.hex(' ')})"
        )
    body, tail = raw[:-CRC_LEN], raw[-CRC_LEN:]
    expected = struct.unpack("<H", tail)[0]
    actual = crc16(body)
    if actual != expected:
        raise ProtocolError(f"CRC {actual:#06x} != {expected:#06x} ({raw.hex(' ')})")
    return Frame(raw[0], raw[1], raw[HEADER_LEN : HEADER_LEN + declared])


# ------------------------------------------------------------------ decode ---


@dataclass
class Measurement:
    temp_c: float
    rh_pct: float
    age_ms: int

    @classmethod
    def decode(cls, p: bytes) -> Measurement:
        if len(p) != 10:
            raise ProtocolError(f"measurement payload is {len(p)} bytes, expected 10")
        t, rh, age = struct.unpack("<iiH", p)
        return cls(t / 1000.0, rh / 1000.0, age)

    def __str__(self) -> str:
        return f"{self.temp_c:+.3f} °C   {self.rh_pct:.3f} %RH   age {self.age_ms} ms"


@dataclass
class Info:
    fw: tuple[int, int, int]
    proto_version: int
    device_serial: int
    sensor_serial: int
    uptime_s: int

    @classmethod
    def decode(cls, p: bytes) -> Info:
        if len(p) != 16:
            raise ProtocolError(f"info payload is {len(p)} bytes, expected 16")
        maj, mnr, pat, proto, dev, sens, up = struct.unpack("<BBBBIII", p)
        return cls((maj, mnr, pat), proto, dev, sens, up)

    def __str__(self) -> str:
        return (
            f"fw {self.fw[0]}.{self.fw[1]}.{self.fw[2]}  proto v{self.proto_version}  "
            f"device_serial {self.device_serial:#010x}  "
            f"sensor_serial {self.sensor_serial:#010x}  uptime {self.uptime_s} s"
        )


@dataclass
class Status:
    flags: int
    sensor_status: int
    i2c_errors: int
    sensor_crc_errors: int
    frame_errors: int
    crc_errors: int

    @classmethod
    def decode(cls, p: bytes) -> Status:
        if len(p) != 11:
            raise ProtocolError(f"status payload is {len(p)} bytes, expected 11")
        return cls(*struct.unpack("<BHHHHH", p))

    @property
    def flag_names(self) -> list[str]:
        return [name for bit, name in FLAG_NAMES if self.flags & bit]

    def __str__(self) -> str:
        return (
            f"flags {self.flags:#04x} [{' '.join(self.flag_names) or '-'}]  "
            f"sht_status {self.sensor_status:#06x}  "
            f"err i2c={self.i2c_errors} sht_crc={self.sensor_crc_errors} "
            f"frame={self.frame_errors} bus_crc={self.crc_errors}"
        )


# ------------------------------------------------------------------ master ---


class Bus:
    """A master on the segment, reached through the USB-RS485 bridge."""

    def __init__(self, port: str, unit: int = 0x01, verbose: bool = False):
        # The bridge is a CDC device, so the baud rate here is ignored; the RS485
        # line rate is fixed in the bridge firmware.
        self.ser = serial.Serial(port, 115200, timeout=0)
        self.unit = unit
        self.verbose = verbose
        time.sleep(0.05)
        self.ser.reset_input_buffer()

    def close(self) -> None:
        self.ser.close()

    def send_raw(self, raw: bytes) -> None:
        """Put bytes on the bus verbatim, bypassing frame construction.

        One write becomes exactly one transmission by the bridge, which is what
        lets us inject deliberately malformed frames.
        """
        if self.verbose:
            print(f"  TX {raw.hex(' ')}", file=sys.stderr)
        self.ser.reset_input_buffer()
        self.ser.write(raw)
        self.ser.flush()

    def read_raw(self, timeout: float = 0.5, gap: float = 0.02) -> bytes:
        """Collect a reply, using an idle gap to decide it is complete.

        Mirrors how the node itself delimits frames, so a reply that arrives in
        several USB chunks is still reassembled into one frame.
        """
        deadline = time.monotonic() + timeout
        buf = bytearray()
        last = None
        while time.monotonic() < deadline:
            waiting = self.ser.in_waiting
            if waiting:
                buf.extend(self.ser.read(waiting))
                last = time.monotonic()
            elif buf and last is not None and time.monotonic() - last >= gap:
                break
            else:
                time.sleep(0.001)
        if self.verbose and buf:
            print(f"  RX {bytes(buf).hex(' ')}", file=sys.stderr)
        return bytes(buf)

    def request(self, cmd: int, payload: bytes = b"", addr: int | None = None) -> Frame:
        """Send a command and return the parsed reply, raising if none arrives."""
        raw = self.expect_raw(cmd, payload, addr)
        if not raw:
            raise ProtocolError(f"no reply to command {cmd:#04x}")
        return parse_frame(raw)

    def expect_raw(
        self, cmd: int, payload: bytes = b"", addr: int | None = None
    ) -> bytes:
        self.send_raw(build_frame(self.unit if addr is None else addr, cmd, payload))
        return self.read_raw()

    def measurement(self) -> Measurement:
        return Measurement.decode(self._ok(CMD_READ_MEASUREMENT).payload)

    def info(self) -> Info:
        return Info.decode(self._ok(CMD_READ_INFO).payload)

    def status(self) -> Status:
        return Status.decode(self._ok(CMD_READ_STATUS).payload)

    def ping(self, payload: bytes) -> bytes:
        return self._ok(CMD_PING, payload).payload

    def _ok(self, cmd: int, payload: bytes = b"") -> Frame:
        reply = self.request(cmd, payload)
        if reply.is_error:
            raise ProtocolError(f"command {cmd:#04x} failed: {reply.error_name}")
        if reply.cmd != cmd:
            raise ProtocolError(f"reply cmd {reply.cmd:#04x} != request {cmd:#04x}")
        if reply.addr != self.unit:
            raise ProtocolError(f"reply from {reply.addr:#04x} != {self.unit:#04x}")
        return reply


# ------------------------------------------------------------------- tests ---


class Results:
    def __init__(self) -> None:
        self.passed = 0
        self.failed: list[tuple[str, str]] = []

    def check(self, name: str, fn) -> None:
        try:
            detail = fn()
            self.passed += 1
            print(f"  \033[32mPASS\033[0m  {name}" + (f"  — {detail}" if detail else ""))
        except Exception as exc:  # noqa: BLE001 - a failed probe is a test result
            self.failed.append((name, str(exc)))
            print(f"  \033[31mFAIL\033[0m  {name}\n          {exc}")


def run_tests(bus: Bus) -> int:
    """Exercise the parts of the protocol only a real bus can verify.

    Framing and CRC arithmetic are already covered by host unit tests in
    `core/`. What cannot be tested off-hardware is everything timing-dependent:
    that the node delimits frames by an idle gap, that it stays off the wire for
    traffic it must not answer, and that its driver-enable turnaround leaves the
    reply intact.
    """
    r = Results()

    print("\nidentity and readings")

    def t_info():
        info = bus.info()
        assert info.proto_version == 1, f"protocol version {info.proto_version} != 1"
        assert info.sensor_serial != 0, "sensor serial is zero — SHT31 not identified"
        return str(info)

    r.check("READ_INFO", t_info)

    def t_measure():
        m = bus.measurement()
        assert -40 <= m.temp_c <= 125, f"temperature {m.temp_c} outside SHT31 range"
        assert 0 <= m.rh_pct <= 100, f"humidity {m.rh_pct} outside 0..100"
        assert m.age_ms <= 3000, f"reading is {m.age_ms} ms old"
        return str(m)

    r.check("READ_MEASUREMENT", t_measure)

    def t_status():
        s = bus.status()
        assert s.flags & 0x10, "EVER_MEASURED not set"
        assert not s.flags & 0x04, "SENSOR_FAULT is set"
        return str(s)

    r.check("READ_STATUS", t_status)

    print("\nround trip integrity")

    def t_ping_binary():
        # Bytes that would need escaping in any delimiter-based protocol.
        probe = bytes([0x00, 0xFF, 0x0A, 0x0D, 0x3A, 0x7E, 0x55, 0xAA])
        echo = bus.ping(probe)
        assert echo == probe, f"echo {echo.hex(' ')} != sent {probe.hex(' ')}"
        return f"{len(probe)} bytes binary-transparent"

    r.check("PING echoes arbitrary bytes", t_ping_binary)

    def t_ping_long():
        # 59 payload bytes -> a 64-byte frame, the largest the bridge can send as
        # a single USB packet and therefore as a single transmission.
        probe = bytes(range(59))
        echo = bus.ping(probe)
        assert echo == probe, f"echo of {len(echo)} bytes != sent {len(probe)}"
        return "59-byte payload, 64-byte frame"

    r.check("PING echoes a maximum single-packet frame", t_ping_long)

    def t_back_to_back():
        # Catches state leaking between frames — a stale receive buffer or a
        # gap-detection flag left set would show up here and nowhere else.
        for i in range(20):
            probe = bytes([i, 0xA5 ^ i])
            echo = bus.ping(probe)
            assert echo == probe, f"iteration {i}: {echo.hex(' ')} != {probe.hex(' ')}"
        return "20 consecutive requests, no state leak"

    r.check("back-to-back requests", t_back_to_back)

    print("\nsilence where silence is required")

    def t_other_unit():
        other = 0x02 if bus.unit != 0x02 else 0x03
        raw = bus.expect_raw(CMD_READ_MEASUREMENT, addr=other)
        assert not raw, f"answered a frame addressed to {other:#04x}: {raw.hex(' ')}"
        return f"ignored a frame for unit {other:#04x}"

    r.check("frame for another unit is ignored", t_other_unit)

    def t_broadcast():
        raw = bus.expect_raw(CMD_READ_MEASUREMENT, addr=ADDR_BROADCAST)
        assert not raw, f"answered a broadcast: {raw.hex(' ')}"
        return "ignored the broadcast address"

    r.check("broadcast is not answered", t_broadcast)

    def t_bad_crc():
        frame = bytearray(build_frame(bus.unit, CMD_READ_MEASUREMENT))
        frame[-1] ^= 0xFF
        bus.send_raw(bytes(frame))
        raw = bus.read_raw()
        assert not raw, f"answered a frame with a bad CRC: {raw.hex(' ')}"
        return "ignored a corrupted frame"

    r.check("bad CRC is ignored", t_bad_crc)

    def t_truncated():
        frame = build_frame(bus.unit, CMD_READ_MEASUREMENT)[:-1]
        bus.send_raw(frame)
        raw = bus.read_raw()
        assert not raw, f"answered a truncated frame: {raw.hex(' ')}"
        return "ignored a truncated frame"

    r.check("truncated frame is ignored", t_truncated)

    def t_length_lie():
        frame = bytearray(build_frame(bus.unit, CMD_PING, b"ab"))
        frame[2] = 5  # claim more payload than is present
        bus.send_raw(bytes(frame))
        raw = bus.read_raw()
        assert not raw, f"answered a frame with a bad LEN: {raw.hex(' ')}"
        return "ignored a frame whose LEN lies"

    r.check("inconsistent LEN is ignored", t_length_lie)

    print("\nerror replies")

    def t_unknown_cmd():
        reply = bus.request(0x7E)
        assert reply.cmd == 0x7E | ERROR_FLAG, f"cmd {reply.cmd:#04x}"
        assert reply.error_name == "ILLEGAL_COMMAND", reply.error_name
        return "0x7E -> ILLEGAL_COMMAND"

    r.check("unknown command is refused", t_unknown_cmd)

    def t_bad_length():
        reply = bus.request(CMD_READ_MEASUREMENT, b"\xaa")
        assert reply.error_name == "ILLEGAL_LENGTH", reply.error_name
        return "payload on READ_MEASUREMENT -> ILLEGAL_LENGTH"

    r.check("unexpected payload is refused", t_bad_length)

    def t_set_address():
        reply = bus.request(CMD_SET_ADDRESS, b"\x22")
        assert reply.error_name == "UNSUPPORTED", reply.error_name
        return "SET_ADDRESS -> UNSUPPORTED, not silently ignored"

    r.check("reserved command reports UNSUPPORTED", t_set_address)

    print("\ndiagnostic counters")

    def t_crc_counter():
        before = bus.status().crc_errors
        bad = bytearray(build_frame(bus.unit, CMD_READ_MEASUREMENT))
        bad[-1] ^= 0xFF
        bus.send_raw(bytes(bad))
        bus.read_raw(timeout=0.15)
        after = bus.status().crc_errors
        assert after == before + 1, (
            f"crc_errors went {before} -> {after}, expected +1"
        )
        return f"crc_errors {before} -> {after}"

    r.check("a corrupt frame increments crc_errors", t_crc_counter)

    total = r.passed + len(r.failed)
    print(f"\n{r.passed}/{total} passed")
    if r.failed:
        print("\nfailures:")
        for name, why in r.failed:
            print(f"  {name}: {why}")
        return 1
    return 0


# --------------------------------------------------------------------- CLI ---


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--port", default="/dev/ttyACM0", help="bridge serial port")
    ap.add_argument(
        "--unit", type=lambda s: int(s, 0), default=0x01, help="node address"
    )
    ap.add_argument("-v", "--verbose", action="store_true", help="dump bus traffic")
    sub = ap.add_subparsers(dest="cmd", required=True)

    sub.add_parser("measure", help="read temperature and humidity once")
    sub.add_parser("info", help="read firmware and serial numbers")
    sub.add_parser("status", help="read health flags and counters")
    sub.add_parser("test", help="run the end-to-end test suite")

    p_ping = sub.add_parser("ping", help="echo test")
    p_ping.add_argument("--bytes", type=int, default=8, help="payload length")

    p_mon = sub.add_parser("monitor", help="poll continuously")
    p_mon.add_argument(
        "--interval", type=float, default=1.0, help="seconds between polls"
    )

    args = ap.parse_args()

    try:
        bus = Bus(args.port, args.unit, args.verbose)
    except serial.SerialException as exc:
        return f"cannot open {args.port}: {exc}"

    try:
        if args.cmd == "measure":
            print(bus.measurement())
        elif args.cmd == "info":
            print(bus.info())
        elif args.cmd == "status":
            print(bus.status())
        elif args.cmd == "ping":
            probe = bytes(i & 0xFF for i in range(args.bytes))
            echo = bus.ping(probe)
            print("echo ok" if echo == probe else f"MISMATCH: {echo.hex(' ')}")
        elif args.cmd == "monitor":
            while True:
                try:
                    print(f"{time.strftime('%H:%M:%S')}  {bus.measurement()}")
                except ProtocolError as exc:
                    print(f"{time.strftime('%H:%M:%S')}  {exc}")
                time.sleep(args.interval)
        elif args.cmd == "test":
            return run_tests(bus)
    except KeyboardInterrupt:
        return 130
    except ProtocolError as exc:
        return f"protocol error: {exc}"
    finally:
        bus.close()

    return 0


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