"""VESC CAN health monitor — pure state machine.

No ROS2, python-can, or hardware dependencies.  Safe to import in unit tests.

State transitions per VESC
--------------------------

    HEALTHY  ──(elapsed > timeout_s = 0.5 s)──► DEGRADED  (send alive, log warn)
                                                    │
                    ◄──(frame received)─────────────┘ (log recovery)
                                                    │
                              (elapsed > estop_s = 2.0 s)──► ESTOP  (trigger e-stop)
                                                                 │
                    ◄──(frame received)──────────────────────────┘ (clear e-stop)

    BUS_OFF is a parallel override: set by node when the CAN interface reports
    bus-off; cleared when the interface recovers.  While active it always
    produces TRIGGER_ESTOP.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum


# ---------------------------------------------------------------------------
# Enumerations
# ---------------------------------------------------------------------------

class VescHealthState(str, Enum):
    HEALTHY  = "healthy"
    DEGRADED = "degraded"   # >timeout_s, alive frames being sent
    ESTOP    = "estop"      # >estop_s unresponsive, e-stop asserted
    BUS_OFF  = "bus_off"    # CAN bus-off error


class Action(str, Enum):
    SEND_ALIVE       = "send_alive"       # send GET_VALUES request
    TRIGGER_ESTOP    = "trigger_estop"    # publish zero cmd_vel / /estop
    CLEAR_ESTOP      = "clear_estop"      # stop asserting e-stop
    LOG_WARN         = "log_warn"         # WARN-level log
    LOG_ERROR        = "log_error"        # ERROR-level log
    LOG_RECOVERY     = "log_recovery"     # INFO-level recovery log


# ---------------------------------------------------------------------------
# Per-VESC state machine
# ---------------------------------------------------------------------------

@dataclass
class VescMonitor:
    """Monitors one VESC and drives the recovery state machine.

    All timestamps are ``time.monotonic()`` seconds.

    Usage::

        mon = VescMonitor(can_id=61)
        actions = mon.tick(now)           # call at ~10 Hz
        actions = mon.on_frame(now)       # call when a CAN frame arrives
    """

    can_id: int
    timeout_s: float = 0.5    # frame-drop detection threshold
    estop_s:   float = 2.0    # escalation to e-stop threshold
    alive_retry_s: float = 0.2  # interval between alive retries in DEGRADED

    # Mutable runtime state (not init args)
    state: VescHealthState = field(default=VescHealthState.HEALTHY, init=False)
    last_frame_ts: float = field(default=0.0, init=False)       # 0 = never
    _first_tick_ts: float = field(default=0.0, init=False)      # set on first tick
    _last_alive_ts: float = field(default=0.0, init=False)

    # ------------------------------------------------------------------

    def on_frame(self, now: float) -> list[Action]:
        """Call whenever a STATUS frame arrives from this VESC."""
        self.last_frame_ts = now

        prev = self.state
        if prev in (VescHealthState.DEGRADED, VescHealthState.ESTOP,
                    VescHealthState.BUS_OFF):
            self.state = VescHealthState.HEALTHY
            actions: list[Action] = [Action.LOG_RECOVERY]
            if prev in (VescHealthState.ESTOP, VescHealthState.BUS_OFF):
                actions.append(Action.CLEAR_ESTOP)
            return actions
        return []

    def on_bus_off(self, now: float) -> list[Action]:
        """Call when the CAN interface enters bus-off state."""
        if self.state != VescHealthState.BUS_OFF:
            self.state = VescHealthState.BUS_OFF
            return [Action.LOG_ERROR, Action.TRIGGER_ESTOP]
        return [Action.TRIGGER_ESTOP]

    def on_bus_ok(self, now: float) -> list[Action]:
        """Call when the CAN interface recovers from bus-off."""
        if self.state == VescHealthState.BUS_OFF:
            self.state = VescHealthState.HEALTHY
            self.last_frame_ts = 0.0   # wait for first fresh frame
            return [Action.LOG_RECOVERY, Action.CLEAR_ESTOP]
        return []

    def tick(self, now: float) -> list[Action]:
        """Periodic update (~10 Hz).  Returns list of actions to take."""
        # BUS_OFF and ESTOP keep asserting until a frame or bus recovery
        if self.state == VescHealthState.BUS_OFF:
            return [Action.TRIGGER_ESTOP]
        if self.state == VescHealthState.ESTOP:
            return [Action.TRIGGER_ESTOP]

        # Track elapsed since startup for the "never received" case so we
        # don't immediately e-stop on node startup before any frames arrive.
        if self._first_tick_ts == 0.0:
            self._first_tick_ts = now

        if self.last_frame_ts > 0:
            elapsed = now - self.last_frame_ts
        else:
            elapsed = now - self._first_tick_ts

        if elapsed >= self.estop_s:
            # Escalate to e-stop
            if self.state != VescHealthState.ESTOP:
                self.state = VescHealthState.ESTOP
                return [Action.TRIGGER_ESTOP, Action.LOG_ERROR]
            return [Action.TRIGGER_ESTOP]

        if elapsed >= self.timeout_s:
            if self.state == VescHealthState.HEALTHY:
                # First timeout — transition to DEGRADED and send alive
                self.state = VescHealthState.DEGRADED
                self._last_alive_ts = now
                return [Action.SEND_ALIVE, Action.LOG_WARN]
            elif self.state == VescHealthState.DEGRADED:
                # Retry alive at interval
                if now - self._last_alive_ts >= self.alive_retry_s:
                    self._last_alive_ts = now
                    return [Action.SEND_ALIVE]
            return []

        # elapsed < timeout_s — connection is healthy
        if self.state == VescHealthState.DEGRADED:
            # Frame arrived (on_frame handles ESTOP/BUS_OFF → HEALTHY)
            # If somehow we reach here still DEGRADED, clear it
            self.state = VescHealthState.HEALTHY
            return [Action.LOG_RECOVERY]

        return []

    # ------------------------------------------------------------------
    # Introspection helpers
    # ------------------------------------------------------------------

    @property
    def is_healthy(self) -> bool:
        return self.state == VescHealthState.HEALTHY

    @property
    def estop_active(self) -> bool:
        return self.state in (VescHealthState.ESTOP, VescHealthState.BUS_OFF)

    def elapsed(self, now: float) -> float:
        """Seconds since last frame (inf if never received)."""
        return (now - self.last_frame_ts) if self.last_frame_ts > 0 else float("inf")


# ---------------------------------------------------------------------------
# Dual-VESC wrapper
# ---------------------------------------------------------------------------

@dataclass
class HealthFsm:
    """Manages two VescMonitor instances and aggregates bus-off state."""

    left_id:  int = 61
    right_id: int = 79
    timeout_s: float = 0.5
    estop_s:   float = 2.0

    def __post_init__(self) -> None:
        self.left  = VescMonitor(self.left_id,  self.timeout_s, self.estop_s)
        self.right = VescMonitor(self.right_id, self.timeout_s, self.estop_s)
        self._bus_off = False

    def monitors(self) -> tuple[VescMonitor, VescMonitor]:
        return self.left, self.right

    def on_frame(self, can_id: int, now: float) -> list[Action]:
        mon = self._get(can_id)
        return mon.on_frame(now) if mon else []

    def on_bus_off(self, now: float) -> list[Action]:
        if not self._bus_off:
            self._bus_off = True
            return self.left.on_bus_off(now) + self.right.on_bus_off(now)
        return [Action.TRIGGER_ESTOP]

    def on_bus_ok(self, now: float) -> list[Action]:
        if self._bus_off:
            self._bus_off = False
            return self.left.on_bus_ok(now) + self.right.on_bus_ok(now)
        return []

    def tick(self, now: float) -> dict[int, list[Action]]:
        """Returns {can_id: [actions]} for each VESC."""
        return {
            self.left_id:  self.left.tick(now),
            self.right_id: self.right.tick(now),
        }

    @property
    def any_estop(self) -> bool:
        return self.left.estop_active or self.right.estop_active

    def _get(self, can_id: int) -> VescMonitor | None:
        if can_id == self.left_id:
            return self.left
        if can_id == self.right_id:
            return self.right
        return None
