motian_mavbridge/motian_bridge/bridge.py

"""MQTT / serial bridge with optional raw forwarding and parsed telemetry."""

from dataclasses import dataclass
from datetime import datetime
import json
import logging
import re
import threading
from typing import Any, Callable

from .config import Config
from .telemetry import Telemetry


LOG = logging.getLogger(__name__)
_HEX = re.compile(r"(?:[0-9a-fA-F]{2})+")
MAVLINK_PASSTHROUGH = 526
MAX_MQTT_PAYLOAD = 65536


@dataclass(frozen=True)
class BaseMessage:
    data_type: int
    uav_type: int
    data: Any
    system_id: str
    timestamp: str


def decode_command(payload: bytes, config: Config) -> BaseMessage:
    """Validate the Base envelope. Data for future handlers is left untouched."""
    if len(payload) > MAX_MQTT_PAYLOAD:
        raise ValueError("MQTT payload too large")
    try:
        value = json.loads(payload.decode("utf-8"))
    except (UnicodeError, json.JSONDecodeError) as exc:
        raise ValueError("invalid JSON Base payload") from exc
    if not isinstance(value, dict):
        raise ValueError("Base payload must be an object")
    if value.get("SystemID") != config.system_id:
        raise ValueError("SystemID does not match this bridge")
    data_type = value.get("DataType")
    if type(data_type) is not int:
        raise ValueError("DataType must be an integer")
    uav_type = value.get("UAVType")
    if type(uav_type) is not int or uav_type != config.uav_type:
        raise ValueError("UAVType does not match this bridge")
    if "Data" not in value:
        raise ValueError("Data is required")
    timestamp = value.get("TimeStamp")
    if not isinstance(timestamp, str):
        raise ValueError("TimeStamp must be a string")
    return BaseMessage(data_type, uav_type, value["Data"], value["SystemID"], timestamp)


def decode_hex(data: Any, max_bytes: int) -> bytes:
    if not isinstance(data, str) or len(data) > 2 * max_bytes or not _HEX.fullmatch(data):
        raise ValueError("Data must be nonempty even-length HEX within the configured limit")
    return bytes.fromhex(data)


def make_report(config: Config, data_type: int, data: Any) -> bytes:
    if type(data_type) is not int:
        raise ValueError("DataType must be an integer")
    envelope = {
        "DataType": data_type,
        "UAVType": config.uav_type,
        "Data": data,
        "SystemID": config.system_id,
        "TimeStamp": datetime.now().astimezone().isoformat(timespec="milliseconds"),
    }
    return json.dumps(envelope, ensure_ascii=False, allow_nan=False,
                      separators=(",", ":")).encode("utf-8")


Handler = Callable[[BaseMessage, "Bridge"], None]


class Bridge:
    """Uses injected MQTT/serial objects; the CLI owns concrete dependencies."""

    def __init__(self, config: Config, mqtt_client: Any, serial_factory: Callable[..., Any]):
        self.config = config
        self.mqtt = mqtt_client
        self.serial_factory = serial_factory
        self._handlers: dict[int, Handler] = {}
        self._mqtt_connected = threading.Event()
        self._stop = threading.Event()
        self._serial_lock = threading.Lock()
        self._serial: Any = None
        self._thread: threading.Thread | None = None
        self._telemetry = Telemetry(config)
        self.mqtt.on_connect = self._on_connect
        self.mqtt.on_disconnect = self._on_disconnect
        self.mqtt.on_message = self._on_message

    def register_handler(self, data_type: int, handler: Handler) -> None:
        """Register a future business-command handler (526 remains reserved)."""
        if type(data_type) is not int or data_type == MAVLINK_PASSTHROUGH or not callable(handler):
            raise ValueError("handler requires a non-526 integer DataType and a callable")
        self._handlers[data_type] = handler

    def publish_report(self, data_type: int, data: Any, *, passthrough: bool = False) -> bool:
        """Publish future reports to /RECEIVE, or raw data to /PUBLISH."""
        if not self._mqtt_connected.is_set():
            return False
        topic = self.config.publish_topic if passthrough else self.config.receive_topic
        try:
            payload = make_report(self.config, data_type, data)
        except (TypeError, ValueError):
            LOG.exception("invalid report data; message discarded")
            return False
        try:
            result = self.mqtt.publish(topic, payload, qos=0, retain=False)
        except Exception:
            LOG.exception("MQTT publish failed")
            return False
        if result.rc != 0:
            LOG.warning("MQTT publish failed: rc=%s", result.rc)
            return False
        return True

    def start(self) -> None:
        self._stop.clear()
        self._thread = threading.Thread(target=self._serial_loop, name="serial-reader", daemon=True)
        self._thread.start()
        self.mqtt.connect_async(self.config.mqtt_host, self.config.mqtt_port, self.config.mqtt_keepalive)
        self.mqtt.loop_start()

    def stop(self) -> None:
        self._stop.set()
        self._mqtt_connected.clear()
        self.mqtt.disconnect()
        self.mqtt.loop_stop()
        with self._serial_lock:
            serial_port, self._serial = self._serial, None
        if serial_port is not None:
            serial_port.close()
        if self._thread is not None:
            self._thread.join(timeout=3)

    def _on_connect(self, client: Any, userdata: Any, flags: Any, reason_code: Any, properties: Any) -> None:
        if reason_code != 0:
            self._mqtt_connected.clear()
            LOG.error("MQTT connection refused: %s", reason_code)
            return
        result, _mid = client.subscribe(self.config.obtain_topic, qos=0)
        if result != 0:
            LOG.error("MQTT subscribe failed: rc=%s", result)
            self._mqtt_connected.clear()
            return
        self._mqtt_connected.set()
        LOG.info("MQTT connected; subscribed to %s", self.config.obtain_topic)

    def _on_disconnect(self, client: Any, userdata: Any, disconnect_flags: Any, reason_code: Any, properties: Any) -> None:
        self._mqtt_connected.clear()
        LOG.warning("MQTT disconnected: %s", reason_code)

    def _on_message(self, client: Any, userdata: Any, message: Any) -> None:
        if message.topic != self.config.obtain_topic or message.retain or not self._mqtt_connected.is_set():
            LOG.warning("ignored wrong-topic, retained, or disconnected MQTT command")
            return
        try:
            command = decode_command(message.payload, self.config)
            if command.data_type == MAVLINK_PASSTHROUGH:
                self._write_serial(decode_hex(command.data, self.config.max_command_bytes))
            else:
                handler = self._handlers.get(command.data_type)
                if handler is None:
                    LOG.warning("unsupported DataType=%s; command discarded", command.data_type)
                else:
                    handler(command, self)
        except ValueError as exc:
            LOG.warning("invalid MQTT command: %s", exc)
        except Exception:
            LOG.exception("command handler failed")

    def _write_serial(self, data: bytes) -> None:
        with self._serial_lock:
            if self._serial is None:
                LOG.warning("serial disconnected; command discarded")
                return
            try:
                count = self._serial.write(data)
                if count != len(data):
                    raise OSError(f"short serial write: {count}/{len(data)}")
            except Exception:
                LOG.exception("serial write failed; command discarded")
                self._serial.close()
                self._serial = None

    def _serial_loop(self) -> None:
        while not self._stop.is_set():
            with self._serial_lock:
                current = self._serial
            if current is None:
                try:
                    current = self.serial_factory(
                        self.config.serial_port,
                        self.config.serial_baudrate,
                        timeout=0.2,
                        write_timeout=0.5,
                    )
                    with self._serial_lock:
                        if self._stop.is_set():
                            current.close()
                            return
                        self._serial = current
                    # A fresh serial link can be another flight controller.
                    self._telemetry = Telemetry(self.config)
                    LOG.info("serial connected: %s", self.config.serial_port)
                except Exception as exc:
                    LOG.warning("serial connection failed: %s", exc)
                    self._stop.wait(self.config.reconnect_seconds)
                    continue
            try:
                data = current.read(self.config.max_chunk_bytes)
            except Exception as exc:
                LOG.warning("serial read failed: %s", exc)
                with self._serial_lock:
                    if self._serial is current:
                        self._serial = None
                        current.close()
                self._stop.wait(self.config.reconnect_seconds)
                continue
            try:
                if data and self.config.publish_raw_mavlink and self._mqtt_connected.is_set():
                    self.publish_report(MAVLINK_PASSTHROUGH, data.hex().upper(), passthrough=True)
                if data:
                    for data_type, report in self._telemetry.feed(data):
                        self.publish_report(data_type, report)
                for data_type, report in self._telemetry.tick():
                    self.publish_report(data_type, report)
            except Exception:
                LOG.exception("telemetry processing failed; serial link preserved")