jmri_nce.consist

  1"""``NceConsistManager`` — hardware consist manager for NCE PowerPro.
  2
  3NCE tracks command-station consists internally (1..127, six-member max:
  4one LEAD + one REAR + up to four MID). The CS handles speed and
  5direction fan-out on its own once a consist is built; we send add /
  6remove commands using the existing ``0xA2`` loco command with
  7consist-specific sub-command bytes.
  8
  9Wire ops (all sub-commands of ``OP_NCE_LOCO_CMD = 0xA2`` — see
 10:mod:`jmri_nce.opcodes` for the numeric constants):
 11
 12* ``LOCO_SUB_CS_ADD_LEAD_FWD / _REV`` — install the first (lead) loco.
 13* ``LOCO_SUB_CS_ADD_REAR_FWD / _REV`` — install the second (rear) loco.
 14* ``LOCO_SUB_CS_ADD_MID_FWD / _REV`` — install a mid loco (positions 3..6).
 15* ``LOCO_SUB_CS_REMOVE`` — remove a *mid* loco. LEAD/REAR are rejected
 16  by the CS; the manager guards this at the bean level.
 17* ``LOCO_SUB_CS_KILL`` — dispose the whole consist; sent using the
 18  LEAD's DCC address so the CS knows which consist to tear down.
 19
 20All values were cross-checked against JMRI Java ``jmri.jmrix.nce.NceConsist``
 21(``add()``, ``removeLocoFromConsist``, ``killConsist``) before being
 22committed here — see the notes in :mod:`jmri_nce.opcodes`.
 23
 24Semantic constraints:
 25
 26* **Address range 1..127** — matches ``CONSIST_ADDRESS_MAX`` (already
 27  enforced by the ``Consist`` model).
 28* **Size limit 6 members** — enforced here so add-past-six raises a
 29  ``ValueError`` that surfaces to the UI as a 400 rather than a silent
 30  no-op.
 31* **LEAD and REAR can't be removed individually** — the CS design
 32  requires either killing the whole consist or leaving them alone.
 33  ``remove_member`` rejects those positions with a clear error.
 34* **Direction is baked in at ADD time**. Flipping a member's
 35  ``direction_normal`` after add doesn't reprogram the CS in this
 36  first pass; the workaround is delete + recreate. Consistent with
 37  JMRI Java's behaviour.
 38
 39Base :class:`~jmri_core.beans.consist_manager.ConsistManager` remains
 40usable for ``kind="software"`` consists on an NCE connection — those
 41bypass the CS and fan out via throttle commands, useful when a caller
 42wants inverted-direction followers (not supported by NCE CS consists in
 43this pass) or more than six members.
 44"""
 45from __future__ import annotations
 46
 47import logging
 48from typing import TYPE_CHECKING
 49
 50from jmri_core.beans.consist import Consist
 51from jmri_core.beans.consist_manager import ConsistManager
 52from jmri_core.beans.throttle_manager import ThrottleManager
 53from jmri_core.state import ThrottleDirection
 54
 55from jmri_nce.messages import nce_loco_cmd
 56from jmri_nce.opcodes import (
 57    LOCO_SUB_CS_ADD_LEAD_FWD,
 58    LOCO_SUB_CS_ADD_LEAD_REV,
 59    LOCO_SUB_CS_ADD_MID_FWD,
 60    LOCO_SUB_CS_ADD_MID_REV,
 61    LOCO_SUB_CS_ADD_REAR_FWD,
 62    LOCO_SUB_CS_ADD_REAR_REV,
 63    LOCO_SUB_CS_KILL,
 64    LOCO_SUB_CS_REMOVE,
 65)
 66
 67if TYPE_CHECKING:
 68    from jmri_core.beans.programmer import ProgrammerManager
 69
 70    from jmri_nce.traffic import NceTrafficController
 71
 72_log = logging.getLogger(__name__)
 73
 74# CS-managed consists on NCE PowerPro top out at six members (one LEAD,
 75# one REAR, up to four MIDs). Enforced in ``add_member`` — the CS would
 76# silently drop the seventh anyway but a clean error surfaces to the UI.
 77NCE_CS_CONSIST_SIZE_LIMIT = 6
 78
 79
 80class NceConsistManager(ConsistManager):
 81    """NCE consist manager. Default kind is ``command_station``.
 82
 83    Software consists still work if the caller asks for
 84    ``kind="software"`` at create time — the base fan-out then drives
 85    each member's throttle directly. Advanced consists (``kind="advanced"``)
 86    also still work, going through the base's CV19 ops-mode write.
 87    """
 88
 89    default_kind = "command_station"
 90
 91    def __init__(
 92        self,
 93        throttles: ThrottleManager,
 94        system_prefix: str,
 95        programmers: ProgrammerManager | None = None,
 96        *,
 97        tc: NceTrafficController,
 98    ) -> None:
 99        super().__init__(throttles, system_prefix, programmers)
100        self._tc = tc
101        # Members already installed on the CS side for each consist,
102        # keyed by consist address. Compared against ``consist.locos``
103        # in :meth:`_after_membership_change` to compute add / remove
104        # deltas. A member's tuple carries ``(address, long_address,
105        # direction_normal)`` so we can rebuild the exact ADD command
106        # if we ever need to re-send after a kill.
107        self._on_cs: dict[int, list[tuple[int, bool, bool]]] = {}
108
109    # ---- overrides ------------------------------------------------------
110
111    async def add_member(
112        self,
113        consist_address: int,
114        member_address: int,
115        *,
116        long_address: bool = False,
117        direction_normal: bool = True,
118    ) -> object:
119        """Enforce the CS-consist size cap before delegating to the
120        base add. Software consists on an NCE connection bypass the
121        cap since they don't touch the CS.
122        """
123        consist = self._consists.get(consist_address)
124        if (
125            consist is not None
126            and consist.kind == "command_station"
127            and not self._is_existing_member(consist, member_address)
128            and len(consist.locos) >= NCE_CS_CONSIST_SIZE_LIMIT
129        ):
130            raise ValueError(
131                f"NCE CS consist {consist_address} is full "
132                f"({NCE_CS_CONSIST_SIZE_LIMIT} members max); "
133                f"remove a member or delete the consist first"
134            )
135        return await super().add_member(
136            consist_address, member_address,
137            long_address=long_address, direction_normal=direction_normal,
138        )
139
140    async def remove_member(
141        self, consist_address: int, member_address: int,
142    ) -> bool:
143        """Reject LEAD (position 0) and REAR (position 1) removals for
144        CS-managed consists — the CS silently ignores those anyway,
145        and the JMRI Java implementation rejects them at the bean
146        level for the same reason. The escape hatch is to delete the
147        whole consist, which sends KILL.
148        """
149        consist = self._consists.get(consist_address)
150        if consist is not None and consist.kind == "command_station":
151            for loco in consist.locos:
152                if loco.address == member_address and loco.position < 2:
153                    role = "LEAD" if loco.position == 0 else "REAR"
154                    raise ValueError(
155                        f"cannot remove {role} loco {member_address} from "
156                        f"NCE CS consist {consist_address}; "
157                        f"delete the consist to disband it"
158                    )
159        return await super().remove_member(consist_address, member_address)
160
161    async def _after_membership_change(self, consist: Consist) -> None:
162        """Diff the CS-side member list against the consist's current
163        locos and send the ADD / REMOVE deltas.
164
165        Direction is baked in at ADD time — flipping ``direction_normal``
166        on an already-added member does *not* re-send the add. To change
167        orientation, remove the member (only MID positions are removable)
168        and re-add, or dispose the whole consist and rebuild.
169        """
170        if consist.kind != "command_station":
171            return
172        cs_list = self._on_cs.setdefault(consist.address, [])
173        cs_addrs = {t[0] for t in cs_list}
174        want_addrs = {m.address for m in consist.locos}
175
176        # Adds first, then removes. On adds, position is inferred from
177        # the member's index in ``consist.locos`` — the base's
178        # ``add_member`` sets ``position`` to the pre-append length, so
179        # it aligns with our LEAD/REAR/MID sub-command mapping.
180        for member in consist.locos:
181            if member.address in cs_addrs:
182                continue
183            await self._send_add(
184                consist.address, member.address,
185                long_address=member.long_address,
186                direction_normal=member.direction_normal,
187                position=member.position,
188            )
189            cs_list.append((
190                member.address, member.long_address, member.direction_normal,
191            ))
192
193        # Removes: anything on the CS side but not in the desired
194        # membership. In practice only MID positions ever land here
195        # because the ``remove_member`` override blocks LEAD/REAR.
196        stale = [t for t in cs_list if t[0] not in want_addrs]
197        for addr, long_addr, _dir in stale:
198            await self._send_remove(addr, long_addr)
199            cs_list.remove((addr, long_addr, _dir))
200
201    async def _before_delete(self, consist: Consist) -> None:
202        """Send KILL to the LEAD address so the CS wipes its
203        consist-table entry for this consist. If there's no LEAD (the
204        consist was empty), there's nothing to kill — just drop our
205        bookkeeping."""
206        if consist.kind != "command_station":
207            return
208        cs_list = self._on_cs.pop(consist.address, None)
209        if not cs_list:
210            return
211        lead_addr, lead_long, _dir = cs_list[0]
212        await self._send_kill(lead_addr, lead_long)
213
214    # ---- fan-out overrides — CS handles speed / direction internally ----
215
216    async def set_speed(self, consist_address: int, speed: float) -> None:
217        """Send to the LEAD only — the CS fans out. Falls through to
218        the base per-loco fan-out for ``kind="software"``.
219        """
220        consist = self._require(consist_address)
221        if consist.kind != "command_station":
222            await super().set_speed(consist_address, speed)
223            return
224        if not consist.locos:
225            return
226        lead = consist.locos[0]
227        throttle = await self._throttles.acquire(
228            lead.address, long_address=lead.long_address,
229        )
230        await throttle.set_speed(speed)
231        # Preserve the UI-side "last commanded" state that the base
232        # tracks — the /consists slider reads this back on re-render.
233        self._speeds[consist_address] = speed
234        await self.updated.emit(consist)
235
236    async def set_direction(
237        self, consist_address: int, direction: ThrottleDirection,
238    ) -> None:
239        """Send to the LEAD only. Per-loco orientation is already baked
240        in on each member's slot (from the ADD command's FWD/REV
241        variant) so flipping the whole consist's direction here is a
242        single-address operation."""
243        consist = self._require(consist_address)
244        if consist.kind != "command_station":
245            await super().set_direction(consist_address, direction)
246            return
247        if not consist.locos:
248            return
249        lead = consist.locos[0]
250        throttle = await self._throttles.acquire(
251            lead.address, long_address=lead.long_address,
252        )
253        await throttle.set_direction(direction)
254        self._directions[consist_address] = direction
255        await self.updated.emit(consist)
256
257    async def emergency_stop(self, consist_address: int) -> None:
258        """E-stop the LEAD; CS fans out. Belt-and-braces: we also
259        e-stop every member throttle individually so a decoder that's
260        momentarily out of sync with the CS still stops (matches JMRI
261        Java's ``NceConsistManager`` e-stop behaviour)."""
262        consist = self._require(consist_address)
263        if consist.kind != "command_station":
264            await super().emergency_stop(consist_address)
265            return
266        for member in consist.locos:
267            throttle = await self._throttles.acquire(
268                member.address, long_address=member.long_address,
269            )
270            await throttle.emergency_stop()
271        self._speeds[consist_address] = 0.0
272        await self.updated.emit(consist)
273
274    # ``set_function`` intentionally *not* overridden. NCE CS consists,
275    # like LocoNet's, don't fan functions — the base's per-loco loop is
276    # exactly what a user wants (a headlight command should hit every
277    # loco separately).
278
279    # ---- wire helpers ---------------------------------------------------
280
281    async def _send_add(
282        self,
283        consist_address: int,
284        member_address: int,
285        *,
286        long_address: bool,
287        direction_normal: bool,
288        position: int,
289    ) -> None:
290        """Send the appropriate ``0xA2`` sub-command based on position
291        and direction. Position 0 → LEAD, 1 → REAR, else MID."""
292        if position == 0:
293            sub = (
294                LOCO_SUB_CS_ADD_LEAD_FWD if direction_normal
295                else LOCO_SUB_CS_ADD_LEAD_REV
296            )
297            role = "LEAD"
298        elif position == 1:
299            sub = (
300                LOCO_SUB_CS_ADD_REAR_FWD if direction_normal
301                else LOCO_SUB_CS_ADD_REAR_REV
302            )
303            role = "REAR"
304        else:
305            sub = (
306                LOCO_SUB_CS_ADD_MID_FWD if direction_normal
307                else LOCO_SUB_CS_ADD_MID_REV
308            )
309            role = f"MID#{position - 1}"
310        cmd = nce_loco_cmd(
311            member_address,
312            long_address=long_address,
313            sub_cmd=sub,
314            data=consist_address,
315        )
316        _log.info(
317            "NCE consist %d: ADD %s addr=%d%s dir=%s (sub=0x%02X)",
318            consist_address, role, member_address,
319            "L" if long_address else "S",
320            "FWD" if direction_normal else "REV",
321            sub,
322        )
323        await self._tc.send_command(cmd)
324
325    async def _send_remove(
326        self, member_address: int, long_address: bool,
327    ) -> None:
328        cmd = nce_loco_cmd(
329            member_address,
330            long_address=long_address,
331            sub_cmd=LOCO_SUB_CS_REMOVE,
332            data=0,
333        )
334        _log.info(
335            "NCE consist: REMOVE addr=%d%s",
336            member_address, "L" if long_address else "S",
337        )
338        await self._tc.send_command(cmd)
339
340    async def _send_kill(
341        self, lead_address: int, lead_long_address: bool,
342    ) -> None:
343        cmd = nce_loco_cmd(
344            lead_address,
345            long_address=lead_long_address,
346            sub_cmd=LOCO_SUB_CS_KILL,
347            data=0,
348        )
349        _log.info(
350            "NCE consist: KILL via LEAD addr=%d%s",
351            lead_address, "L" if lead_long_address else "S",
352        )
353        await self._tc.send_command(cmd)
354
355    def _is_existing_member(self, consist: Consist, address: int) -> bool:
356        return any(loco.address == address for loco in consist.locos)