"""
docking_state_machine.py — FSM for autonomous docking / undocking (Issue #158).

States
──────
  IDLE          : normal operation; watching battery.
  DETECTING     : scanning for dock marker or IR beacon.
  NAV2_APPROACH : Nav2 navigating to pre-dock waypoint (~1 m from dock face).
  VISUAL_SERVO  : IBVS final approach (<= servo_range_m to dock).
  CONTACT       : dock pins touched; waiting for charge current.
  CHARGING      : charging in progress.
  UNDOCKING     : backing away from dock at fixed reverse speed.

Transitions
───────────
  IDLE          → DETECTING     : battery_pct < battery_low_pct  OR  dock_requested
  DETECTING     → NAV2_APPROACH : dock_pose found AND distance > servo_range_m
  DETECTING     → VISUAL_SERVO  : dock_pose found AND distance ≤ servo_range_m
  NAV2_APPROACH → VISUAL_SERVO  : dock_pose found AND distance ≤ servo_range_m
                                   OR  nav2_arrived
  NAV2_APPROACH → DETECTING     : nav2_failed (retry detection)
  VISUAL_SERVO  → CONTACT       : servo_aligned
  VISUAL_SERVO  → DETECTING     : dock lost for > lost_ticks_max ticks
  CONTACT       → CHARGING      : is_charging
  CONTACT       → VISUAL_SERVO  : contact_timeout_ticks exceeded (retry)
  CHARGING      → UNDOCKING     : battery_pct >= battery_high_pct
                                   OR  undock_requested
  UNDOCKING     → IDLE          : undock_ticks >= undock_ticks_max

Any state    → IDLE             : undock_requested  (except CHARGING → UNDOCKING)
DETECTING    → IDLE             : undock_requested  (abort scan)

Output signals
──────────────
  cmd_linear, cmd_angular : velocity during VISUAL_SERVO and UNDOCKING
  request_nav2            : send Nav2 goal (emitted once on ENTRY to NAV2_APPROACH)
  cancel_nav2             : cancel Nav2 goal
  mission_interrupted     : battery forced dock (for mission planner)
  mission_resume          : auto-resume after charge cycle completes
  state_changed           : FSM state changed this tick

Pure module — no ROS2 dependency.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional

from saltybot_docking.dock_detector import DockPose


# ── States ────────────────────────────────────────────────────────────────────

class DockingState(Enum):
    IDLE          = "IDLE"
    DETECTING     = "DETECTING"
    NAV2_APPROACH = "NAV2_APPROACH"
    VISUAL_SERVO  = "VISUAL_SERVO"
    CONTACT       = "CONTACT"
    CHARGING      = "CHARGING"
    UNDOCKING     = "UNDOCKING"


# ── I/O dataclasses ───────────────────────────────────────────────────────────

@dataclass
class DockingInputs:
    """Sensor snapshot fed to the FSM each tick."""
    battery_pct:     float              = 100.0
    is_charging:     bool               = False
    dock_pose:       Optional[DockPose] = None   # None = dock not visible
    servo_aligned:   bool               = False  # visual servo declared aligned
    dock_requested:  bool               = False  # explicit /saltybot/dock call
    undock_requested: bool              = False  # explicit /saltybot/undock call
    nav2_arrived:    bool               = False  # Nav2 action succeeded
    nav2_failed:     bool               = False  # Nav2 action failed/cancelled


@dataclass
class DockingOutputs:
    """Commands and signals produced by the FSM each tick."""
    state:               DockingState = DockingState.IDLE
    cmd_linear:          float        = 0.0     # m/s (VISUAL_SERVO + UNDOCKING)
    cmd_angular:         float        = 0.0     # rad/s (VISUAL_SERVO)
    request_nav2:        bool         = False   # send Nav2 goal
    cancel_nav2:         bool         = False   # cancel Nav2 goal
    mission_interrupted: bool         = False   # battery dock forced
    mission_resume:      bool         = False   # charge cycle complete, resume
    state_changed:       bool         = False   # FSM state changed this tick


# ── DockingStateMachine ───────────────────────────────────────────────────────

class DockingStateMachine:
    """
    Deterministic docking FSM.

    Parameters
    ----------
    battery_low_pct     : SOC that triggers auto-dock (default 15 %)
    battery_high_pct    : SOC at which charging completes (default 80 %)
    servo_range_m       : switch from Nav2 to visual servo (default 1.0 m)
    undock_speed_ms     : reverse speed during undocking (m/s, must be ≤ 0)
    undock_ticks_max    : ticks at undock_speed before returning to IDLE
    lost_ticks_max      : allowed VISUAL_SERVO ticks with no detection
    contact_timeout_ticks : CONTACT ticks before retry alignment
    """

    def __init__(
        self,
        battery_low_pct:       float = 15.0,
        battery_high_pct:      float = 80.0,
        servo_range_m:         float = 1.0,
        undock_speed_ms:       float = -0.20,
        undock_ticks_max:      int   = 50,    # 2.5 s at 20 Hz
        lost_ticks_max:        int   = 40,    # 2.0 s at 20 Hz
        contact_timeout_ticks: int   = 20,    # 1.0 s at 20 Hz
    ):
        self._bat_low       = float(battery_low_pct)
        self._bat_high      = float(battery_high_pct)
        self._servo_range   = float(servo_range_m)
        self._undock_speed  = min(0.0, float(undock_speed_ms))
        self._undock_max    = max(1, int(undock_ticks_max))
        self._lost_max      = max(1, int(lost_ticks_max))
        self._contact_max   = max(1, int(contact_timeout_ticks))

        self._state           = DockingState.IDLE
        self._undock_ticks    = 0
        self._lost_ticks      = 0
        self._contact_ticks   = 0
        self._mission_interrupted = False
        self._resume_on_undock    = False

    # ── Public API ────────────────────────────────────────────────────────────

    @property
    def state(self) -> DockingState:
        return self._state

    def reset(self) -> None:
        """Return to IDLE and clear all counters (use for unit tests)."""
        self._state               = DockingState.IDLE
        self._undock_ticks        = 0
        self._lost_ticks          = 0
        self._contact_ticks       = 0
        self._mission_interrupted = False
        self._resume_on_undock    = False

    def tick(self, inputs: DockingInputs) -> DockingOutputs:
        """Advance the FSM by one tick and return output commands."""
        prev_state = self._state
        out = self._step(inputs)
        out.state         = self._state
        out.state_changed = (self._state != prev_state)
        if out.state_changed:
            self._on_enter(self._state)
        return out

    # ── Internal step ─────────────────────────────────────────────────────────

    def _step(self, inp: DockingInputs) -> DockingOutputs:
        out = DockingOutputs(state=self._state)

        # ── IDLE ──────────────────────────────────────────────────────────────
        if self._state == DockingState.IDLE:
            if inp.dock_requested:
                self._state = DockingState.DETECTING
            elif inp.battery_pct < self._bat_low:
                self._state = DockingState.DETECTING
                out.mission_interrupted   = True
                self._mission_interrupted = True

        # ── DETECTING ─────────────────────────────────────────────────────────
        elif self._state == DockingState.DETECTING:
            if inp.undock_requested:
                self._state = DockingState.IDLE
                self._mission_interrupted = False
            elif inp.dock_pose is not None:
                if inp.dock_pose.distance_m <= self._servo_range:
                    self._state = DockingState.VISUAL_SERVO
                else:
                    self._state    = DockingState.NAV2_APPROACH
                    out.request_nav2 = True

        # ── NAV2_APPROACH ─────────────────────────────────────────────────────
        elif self._state == DockingState.NAV2_APPROACH:
            if inp.undock_requested:
                self._state      = DockingState.IDLE
                out.cancel_nav2  = True
                self._mission_interrupted = False
            elif inp.dock_pose is not None and inp.dock_pose.distance_m <= self._servo_range:
                self._state     = DockingState.VISUAL_SERVO
                out.cancel_nav2 = True
            elif inp.nav2_arrived:
                self._state = DockingState.VISUAL_SERVO
            elif inp.nav2_failed:
                self._state = DockingState.DETECTING   # retry scan

        # ── VISUAL_SERVO ──────────────────────────────────────────────────────
        elif self._state == DockingState.VISUAL_SERVO:
            if inp.undock_requested:
                self._state = DockingState.IDLE
            elif inp.dock_pose is None:
                self._lost_ticks += 1
                if self._lost_ticks >= self._lost_max:
                    self._state = DockingState.DETECTING
            else:
                self._lost_ticks = 0
                if inp.servo_aligned:
                    self._state = DockingState.CONTACT

        # ── CONTACT ───────────────────────────────────────────────────────────
        elif self._state == DockingState.CONTACT:
            if inp.undock_requested:
                self._state = DockingState.UNDOCKING
                self._resume_on_undock = False
            elif inp.is_charging:
                self._state = DockingState.CHARGING
            else:
                self._contact_ticks += 1
                if self._contact_ticks >= self._contact_max:
                    self._state = DockingState.VISUAL_SERVO   # realign

        # ── CHARGING ──────────────────────────────────────────────────────────
        elif self._state == DockingState.CHARGING:
            if inp.undock_requested:
                self._state            = DockingState.UNDOCKING
                self._resume_on_undock = False
            elif inp.battery_pct >= self._bat_high:
                self._state            = DockingState.UNDOCKING
                self._resume_on_undock = self._mission_interrupted

        # ── UNDOCKING ─────────────────────────────────────────────────────────
        elif self._state == DockingState.UNDOCKING:
            out.cmd_linear = self._undock_speed
            self._undock_ticks += 1
            if self._undock_ticks >= self._undock_max:
                out.mission_resume        = self._resume_on_undock
                self._mission_interrupted = False
                self._resume_on_undock    = False
                self._state               = DockingState.IDLE

        return out

    # ── Entry actions ─────────────────────────────────────────────────────────

    def _on_enter(self, state: DockingState) -> None:
        """Reset per-state counters when entering a new state."""
        if state == DockingState.VISUAL_SERVO:
            self._lost_ticks    = 0
        elif state == DockingState.CONTACT:
            self._contact_ticks = 0
        elif state == DockingState.UNDOCKING:
            self._undock_ticks  = 0
