motian_mavbridge/motian_bridge/aircraft.py

"""Per-flight-controller state grouped into a single aircraft object.

An :class:`Aircraft` is created only after the first valid flight-controller
HEARTBEAT is locked by :class:`~motian_bridge.telemetry.Telemetry`. It owns the
identity (system/component), the latest raw observations, the vendor battery
packs and the downloaded parameter table. A serial reconnect discards the whole
object, exactly like :class:`Telemetry` itself.
"""

from __future__ import annotations

from dataclasses import dataclass
import math
from typing import Any

from .config import Config
from .dialects.v20 import VKFly


# Observations older than this are never used as current flight state.
STALE_SECONDS = 5.0
# VKFly does not define a parameter index for unknown entries.
UNKNOWN_PARAM_INDEX = 0xFFFF


def enum_label(enum: str, value: int, prefix: str) -> str:
    entry = VKFly.enums.get(enum, {}).get(value)
    return entry.name.removeprefix(prefix) if entry else f"UNKNOWN_{value}"


def distance_m(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
    r1, r2 = math.radians(lat1), math.radians(lat2)
    dr = r2 - r1
    dl = math.radians(lon2 - lon1)
    a = math.sin(dr / 2) ** 2 + math.cos(r1) * math.cos(r2) * math.sin(dl / 2) ** 2
    return 6371000.0 * 2 * math.asin(min(1.0, math.sqrt(a)))


def _put_finite(data: dict[str, Any], name: str, value: float) -> None:
    if math.isfinite(value):
        data[name] = value


@dataclass
class ParamEntry:
    """One downloaded parameter, with the metadata VKFly reports."""

    value: float
    param_type: int | None = None
    index: int | None = None


class ParamTable:
    """The full parameter table of one aircraft, filled from PARAM_VALUE frames."""

    def __init__(self) -> None:
        self.entries: dict[str, ParamEntry] = {}
        # param_count from the flight controller, once it has been seen.
        self.count: int | None = None
        # param_index values already received, used to detect a complete download.
        self.received: set[int] = set()

    def update(self, message: Any) -> None:
        """Record one PARAM_VALUE message."""
        param_id = message.param_id
        if isinstance(param_id, bytes):
            param_id = param_id.decode("ascii", "replace")
        param_id = param_id.rstrip("\x00")
        index = int(message.param_index)
        self.entries[param_id] = ParamEntry(
            value=float(message.param_value),
            param_type=int(message.param_type),
            index=index,
        )
        count = int(message.param_count)
        if count > 0:
            self.count = count
        if index != UNKNOWN_PARAM_INDEX:
            self.received.add(index)

    def __setitem__(self, param_id: str, value: float | ParamEntry) -> None:
        entry = value if isinstance(value, ParamEntry) else ParamEntry(float(value))
        self.entries[param_id] = entry

    def __contains__(self, param_id: object) -> bool:
        return param_id in self.entries

    def __len__(self) -> int:
        return len(self.entries)

    def get(self, param_id: str, default: float | None = None) -> float | None:
        entry = self.entries.get(param_id)
        return entry.value if entry is not None else default

    def complete(self) -> bool:
        """True once every indexed parameter has been received."""
        return self.count is not None and self.count > 0 and len(self.received) >= self.count


class Aircraft:
    """State that belongs to a single flight controller."""

    def __init__(self, config: Config, system: int, component: int,
                 autopilot: int, uav_type: int) -> None:
        self.config = config
        self.system = system
        self.component = component
        self.autopilot = autopilot
        self.uav_type = uav_type
        self.latest: dict[str, tuple[Any, float]] = {}
        self.batteries: dict[int, tuple[Any, float]] = {}
        self.params = ParamTable()
        # Automatic parameter download bookkeeping.
        self.param_sync_requested = False
        self.param_sync_at: float | None = None
        self.param_request_attempts = 0

    @property
    def source(self) -> tuple[int, int]:
        return (self.system, self.component)

    def fresh(self, kind: str, now: float) -> Any | None:
        entry = self.latest.get(kind)
        return entry[0] if entry and 0 <= now - entry[1] <= STALE_SECONDS else None

    def param_value(self, param_id: str) -> float | None:
        return self.params.get(param_id)

    def flight_data(self, now: float, heartbeat: Any) -> dict[str, Any]:
        data: dict[str, Any] = {
            "SystemID": self.config.system_id,
            "Armed": bool(heartbeat.base_mode & VKFly.MAV_MODE_FLAG_SAFETY_ARMED),
            "IsFlying": bool(heartbeat.base_mode & VKFly.MAV_MODE_FLAG_SAFETY_ARMED),
            "DroneAutopilot": int(heartbeat.autopilot),
            "UavType": int(heartbeat.type),
            "FlightMode": enum_label(
                "VKFLY_CUSTOM_MODE", heartbeat.custom_mode, "VKFLY_CUSTOM_MODE_"
            ),
        }
        attitude = self.fresh("ATTITUDE", now)
        if attitude:
            _put_finite(data, "Roll", math.degrees(attitude.roll))
            _put_finite(data, "Pitch", math.degrees(attitude.pitch))
            _put_finite(data, "Yaw", math.degrees(attitude.yaw) % 360)
        position = self.fresh("GLOBAL_POSITION_INT", now)
        if position:
            if abs(position.lat) <= 900000000 and abs(position.lon) <= 1800000000:
                data.update(Latitude=position.lat / 1e7, Longitude=position.lon / 1e7)
            data.update(
                Altitude=position.relative_alt / 1000,
                VelocityX=position.vx / 100,
                VelocityY=position.vy / 100,
                VelocityZ=position.vz / 100,
            )
        hud = self.fresh("VFR_HUD", now)
        if hud:
            for key, value in (
                ("AirSpeed", hud.airspeed),
                ("GroundSpeed", hud.groundspeed),
                ("ClimbRate", hud.climb),
            ):
                _put_finite(data, key, value)
        home = self.fresh("HOME_POSITION", now)
        if (
            home
            and abs(home.latitude) <= 900000000
            and abs(home.longitude) <= 1800000000
        ):
            data.update(
                HomeLatitude=home.latitude / 1e7,
                HomeLongitude=home.longitude / 1e7,
                HomeAltitude=home.altitude / 1000,
                IsHomeLocationSet=True,
            )
            if position and "Latitude" in data:
                distance = distance_m(
                    data["Latitude"],
                    data["Longitude"],
                    data["HomeLatitude"],
                    data["HomeLongitude"],
                )
                data["DistanceToHome"] = distance
                data["SpaceDistToHome"] = math.hypot(
                    distance, position.alt / 1000 - data["HomeAltitude"]
                )
        gps = self.fresh("GPS_RAW_INT", now)
        if gps:
            info: dict[str, Any] = {
                "Type": 15,
                "FixType": enum_label("GPS_FIX_TYPE", gps.fix_type, "GPS_FIX_TYPE_"),
            }
            if gps.fix_type >= 2:
                info["AbsoluteAlt"] = gps.alt / 1000
            if gps.eph != 65535:
                info["LevelDivisor"] = gps.eph / 100
            if gps.epv != 65535:
                info["VerticalDivisor"] = gps.epv / 100
            if gps.satellites_visible != 255:
                info["SatelliteCount"] = gps.satellites_visible
            data["Gps"] = info
        return data

    def battery_data(self, now: float) -> dict[str, Any] | None:
        packs = {
            id_: item[0]
            for id_, item in self.batteries.items()
            if 0 <= now - item[1] <= STALE_SECONDS
        }
        if not packs:
            return None
        data: dict[str, Any] = {"SystemID": self.config.system_id}
        for id_, pack in packs.items():
            side = "Left" if id_ == 0 else "Right"
            if pack.voltage:
                data["Battery_Voltage" + side] = pack.voltage / 1000
            data["Battery_Current" + side] = pack.current / 100
            if 0 <= pack.cap_percent <= 100:
                data["Battery_Remaining" + side] = pack.cap_percent
        if len(packs) == 1:
            pack = next(iter(packs.values()))
            if pack.voltage:
                data["Battery_Voltage"] = pack.voltage / 1000
            data["Battery_Current"] = pack.current / 100
            if 0 <= pack.cap_percent <= 100:
                data["Battery_Remaining"] = pack.cap_percent
            data["Battery_Temperature"] = pack.temperature
            data["cycleIndex"] = pack.cyc_cnt
            data["flags"] = pack.err_code
        return data