kota-technician-tool/kota_tech/serialio/manager.py

"""Serial connection manager (GUI-independent).

One background thread owns the port: it opens, reads, frames and parses
lines, executes queued commands (reset, write) and re-finds the board if
the port vanishes (power switch, USB re-enumeration). Callbacks are
invoked from that thread; the Qt layer marshals them to the UI thread.
"""

from __future__ import annotations

import enum
import queue
import sys
import threading
import time
from collections.abc import Callable
from dataclasses import dataclass

import serial

from ..core.logline import LineFramer, LogLine, host_event, parse_line
from .ports import PortInfo, find_port

READ_TIMEOUT_S = 0.1
REOPEN_POLL_S = 0.5


class State(str, enum.Enum):
    DISCONNECTED = "Disconnected"
    CONNECTED = "Connected"
    WAITING = "Waiting for board"   # port lost; polling for it to return


@dataclass
class _Target:
    device: str
    serial_number: str | None


class SerialManager:
    def __init__(
        self,
        on_line: Callable[[LogLine], None],
        on_state: Callable[[State, str], None],
        baudrate: int = 115200,  # ignored by native USB CDC, used for UART
    ):
        self._on_line = on_line
        self._on_state = on_state
        self._baud = baudrate
        self._cmds: queue.Queue[tuple[str, object]] = queue.Queue()
        self._thread: threading.Thread | None = None
        self._stop = threading.Event()
        self._target: _Target | None = None
        self.state = State.DISCONNECTED

    # ---- public API (any thread) -------------------------------------
    def connect(self, port: PortInfo) -> None:
        self.disconnect()
        self._target = _Target(port.device, port.serial_number)
        self._stop.clear()
        self._thread = threading.Thread(target=self._run, name="serial-reader", daemon=True)
        self._thread.start()

    def disconnect(self) -> None:
        if self._thread and self._thread.is_alive():
            self._stop.set()
            self._thread.join(timeout=2.0)
        self._thread = None

    def reset_board(self) -> None:
        """Hard-reset the chip into the application (not download mode)."""
        self._cmds.put(("reset", None))

    def write(self, data: bytes) -> None:
        self._cmds.put(("write", data))

    @property
    def is_running(self) -> bool:
        return bool(self._thread and self._thread.is_alive())

    # ---- reader thread -----------------------------------------------
    def _set_state(self, state: State, detail: str = "") -> None:
        if state != self.state:
            self.state = state
            self._on_state(state, detail)

    def _open(self) -> serial.Serial | None:
        assert self._target
        info = find_port(self._target.serial_number, self._target.device)
        if info is None:
            return None
        # serial_for_url also accepts test URLs such as "loop://".
        ser = serial.serial_for_url(
            info.device, baudrate=self._baud, timeout=READ_TIMEOUT_S, do_not_open=True
        )
        # On USB-Serial-JTAG, RTS=1/DTR=0 holds the chip in reset and DTR=1
        # during reset selects download mode.
        # Windows: pyserial applies both lines atomically at open, so
        # deassert both beforehand (it would otherwise assert both).
        # macOS/Linux: the OS asserts both lines on open and pyserial then
        # writes DTR before RTS, passing through RTS=1/DTR=0 and resetting
        # the box on every connect. Drop RTS first, then DTR.
        posix = sys.platform != "win32"
        if not posix:
            ser.dtr = False
        ser.rts = False
        try:
            ser.open()
            if posix:
                ser.dtr = False
        except (serial.SerialException, OSError):
            return None
        self._target.device = info.device  # COM number / node may change
        return ser

    @staticmethod
    def _set_rts(ser: serial.Serial, state: bool) -> None:
        ser.rts = state
        # Windows usbser.sys only sends SET_CONTROL_LINE_STATE when DTR is
        # written, so an RTS-only change never reaches the chip. Re-write
        # DTR with its current value to push the new RTS (esptool does the
        # same). Harmless on macOS.
        ser.dtr = ser.dtr

    @classmethod
    def _hard_reset(cls, ser: serial.Serial) -> None:
        # Same sequence as esptool HardReset for USB-Serial-JTAG: pulse
        # EN via RTS with DTR (GPIO0) released, long enough for the USB
        # peripheral to settle.
        ser.dtr = False
        cls._set_rts(ser, True)
        time.sleep(0.2)
        cls._set_rts(ser, False)
        time.sleep(0.2)

    def _run(self) -> None:
        framer = LineFramer()
        ser: serial.Serial | None = None
        announced_wait = False
        try:
            while not self._stop.is_set():
                if ser is None:
                    ser = self._open()
                    if ser is None:
                        if not announced_wait:
                            self._set_state(State.WAITING, "Plug in / power on the sensorbox")
                            announced_wait = True
                        self._drain_cmds(None)
                        time.sleep(REOPEN_POLL_S)
                        continue
                    announced_wait = False
                    framer = LineFramer(drop_first_partial=True)
                    self._set_state(State.CONNECTED, ser.port)
                    self._on_line(host_event(f"connected to {ser.port}"))

                try:
                    if self._drain_cmds(ser):
                        framer.cancel_drop()  # output after a reset is always live
                    data = ser.read(ser.in_waiting or 1)
                except (serial.SerialException, OSError) as exc:
                    for raw in framer.flush():
                        self._on_line(parse_line(raw))
                    self._on_line(host_event(f"port lost ({exc.__class__.__name__}); waiting for board"))
                    self._close(ser)
                    ser = None
                    continue

                for raw in framer.feed(data) + framer.poll_idle():
                    self._on_line(parse_line(raw))
        finally:
            for raw in framer.flush():
                self._on_line(parse_line(raw))
            if ser is not None:
                self._close(ser)
            self._on_line(host_event("disconnected"))
            self._set_state(State.DISCONNECTED)

    def _drain_cmds(self, ser: serial.Serial | None) -> bool:
        """Execute queued commands; returns True if the board was reset."""
        did_reset = False
        while True:
            try:
                cmd, arg = self._cmds.get_nowait()
            except queue.Empty:
                return did_reset
            if ser is None:
                self._on_line(host_event(f"'{cmd}' ignored: board not connected"))
                continue
            if cmd == "reset":
                self._on_line(host_event("reset issued by tool"))
                self._hard_reset(ser)
                did_reset = True
            elif cmd == "write":
                ser.write(arg)  # type: ignore[arg-type]

    @staticmethod
    def _close(ser: serial.Serial) -> None:
        try:
            ser.close()
        except (serial.SerialException, OSError):
            pass