"""
docking_node.py — Autonomous docking station return (Issue #158).

Overview
────────
  Orchestrates dock detection, Nav2 corridor approach, visual-servo final
  alignment, and charge monitoring.  Interrupts the active Nav2 mission when
  battery drops to 20 % and resumes when charged to 80 %.

Pipeline (20 Hz)
────────────────
  1. Camera callback → ArucoDetector → DockPose (if marker visible)
  2. IR callback     → IRBeaconDetector → presence flag (fallback)
  3. Battery callback→ ChargeMonitor → ChargeEvents
  4. 20 Hz timer     → DockingStateMachine.tick()
                     → publish cmd_vel  (VISUAL_SERVO + UNDOCKING)
                     → Nav2 action calls (NAV2_APPROACH)
                     → publish DockingStatus

Subscribes
──────────
  <image_topic>          sensor_msgs/Image          — dock camera
  <camera_info_topic>    sensor_msgs/CameraInfo     — intrinsics for PnP
  <ir_topic>             std_msgs/Float32           — IR beacon ADC sample
  <battery_topic>        sensor_msgs/BatteryState   — SOC + current

Publishes
─────────
  /saltybot/docking_status   saltybot_docking_msgs/DockingStatus
  /saltybot/docking_cmd_vel  geometry_msgs/Twist   — servo + undock commands
  /saltybot/resume_mission   std_msgs/Bool          — trigger mission resume

Services
────────
  /saltybot/dock    saltybot_docking_msgs/Dock
  /saltybot/undock  saltybot_docking_msgs/Undock

Action client
─────────────
  /navigate_to_pose   nav2_msgs/action/NavigateToPose  (optional)

Parameters
──────────
  control_rate           20.0     Hz
  image_topic            /camera/image_raw
  camera_info_topic      /camera/camera_info
  ir_topic               /saltybot/ir_beacon
  battery_topic          /saltybot/battery_state
  aruco_marker_id        42
  aruco_marker_size_m    0.10
  ir_threshold           0.50
  battery_low_pct        20.0
  battery_high_pct       80.0
  servo_range_m          1.00     m   (switch to IBVS within this distance)
  k_linear               0.30
  k_angular              0.80
  lateral_tol_m          0.005    m   (±5 mm alignment tolerance)
  contact_distance_m     0.05     m
  max_linear_ms          0.10     m/s
  max_angular_rads       0.30     rad/s
  undock_speed_ms       -0.20     m/s
  undock_ticks_max       50       (2.5 s at 20 Hz)
  lost_ticks_max         40       (2.0 s at 20 Hz)
  contact_timeout_ticks  20       (1.0 s at 20 Hz)
  dock_pose              [0,0,0,0,0,0,1]   [x,y,z,qx,qy,qz,qw] map frame
  pre_dock_pose          [0,0,0,0,0,0,1]   1 m in front of dock
"""

import json
import math

import rclpy
from rclpy.action import ActionClient
from rclpy.node import Node
from rclpy.qos import HistoryPolicy, QoSProfile, ReliabilityPolicy

from geometry_msgs.msg import PoseStamped, Twist
from sensor_msgs.msg import BatteryState, CameraInfo, Image
from std_msgs.msg import Bool, Float32, String

from saltybot_docking.charge_monitor import ChargeEvent, ChargeMonitor
from saltybot_docking.dock_detector import ArucoDetector, DockPose, IRBeaconDetector
from saltybot_docking.docking_state_machine import (
    DockingInputs,
    DockingState,
    DockingStateMachine,
)
from saltybot_docking.visual_servo import VisualServo

try:
    from saltybot_docking_msgs.msg import DockingStatus
    from saltybot_docking_msgs.srv import Dock, Undock
    _MSGS_OK = True
except ImportError:
    _MSGS_OK = False

try:
    from nav2_msgs.action import NavigateToPose
    _NAV2_OK = True
except ImportError:
    _NAV2_OK = False

try:
    from saltybot_social_msgs.msg import Mood
    _SOCIAL_OK = True
except ImportError:
    _SOCIAL_OK = False

try:
    import cv2
    import numpy as np
    _CV2_OK = True
except ImportError:
    _CV2_OK = False


# ── Helper ────────────────────────────────────────────────────────────────────

def _pose_param_to_stamped(vals: list, frame_id: str, stamp) -> PoseStamped:
    """Convert [x,y,z,qx,qy,qz,qw] list to PoseStamped."""
    ps = PoseStamped()
    ps.header.frame_id = frame_id
    ps.header.stamp    = stamp
    if len(vals) >= 7:
        ps.pose.position.x    = float(vals[0])
        ps.pose.position.y    = float(vals[1])
        ps.pose.position.z    = float(vals[2])
        ps.pose.orientation.x = float(vals[3])
        ps.pose.orientation.y = float(vals[4])
        ps.pose.orientation.z = float(vals[5])
        ps.pose.orientation.w = float(vals[6])
    return ps


# ── Node ──────────────────────────────────────────────────────────────────────

class DockingNode(Node):

    def __init__(self):
        super().__init__("docking")

        # ── Parameters ───────────────────────────────────────────────────────
        self._declare_params()
        p = self._load_params()

        # ── Sub-systems ──────────────────────────────────────────────────────
        self._aruco = ArucoDetector(
            marker_id=p["aruco_marker_id"],
            marker_size_m=p["aruco_marker_size_m"],
        )
        self._ir = IRBeaconDetector(
            threshold=p["ir_threshold"],
        )
        self._servo = VisualServo(
            k_linear=p["k_linear"],
            k_angular=p["k_angular"],
            lateral_tol_m=p["lateral_tol_m"],
            contact_distance_m=p["contact_distance_m"],
            max_linear_ms=p["max_linear_ms"],
            max_angular_rads=p["max_angular_rads"],
        )
        self._charge = ChargeMonitor(
            low_threshold_pct=p["battery_low_pct"],
            high_threshold_pct=p["battery_high_pct"],
        )
        self._fsm = DockingStateMachine(
            battery_low_pct=p["battery_low_pct"],
            battery_high_pct=p["battery_high_pct"],
            servo_range_m=p["servo_range_m"],
            undock_speed_ms=p["undock_speed_ms"],
            undock_ticks_max=p["undock_ticks_max"],
            lost_ticks_max=p["lost_ticks_max"],
            contact_timeout_ticks=p["contact_timeout_ticks"],
        )

        # ── Mutable state ────────────────────────────────────────────────────
        self._latest_dock_pose: DockPose | None = None
        self._latest_servo_aligned: bool = False
        self._nav2_arrived:  bool = False
        self._nav2_failed:   bool = False
        self._dock_requested:  bool = False
        self._undock_requested: bool = False
        self._camera_matrix  = None
        self._dist_coeffs    = None
        self._nav2_goal_handle = None

        # ── QoS ──────────────────────────────────────────────────────────────
        reliable = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            history=HistoryPolicy.KEEP_LAST,
            depth=10,
        )
        best_effort = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
        )

        # ── Subscriptions ────────────────────────────────────────────────────
        self.create_subscription(Image,       p["image_topic"],       self._image_cb,    best_effort)
        self.create_subscription(CameraInfo,  p["camera_info_topic"], self._cam_info_cb, reliable)
        self.create_subscription(Float32,     p["ir_topic"],          self._ir_cb,       best_effort)
        self.create_subscription(BatteryState, p["battery_topic"],    self._battery_cb,  reliable)

        # ── Publishers ───────────────────────────────────────────────────────
        self._cmd_vel_pub = self.create_publisher(
            Twist, "/saltybot/docking_cmd_vel", reliable
        )
        self._resume_pub = self.create_publisher(
            Bool, "/saltybot/resume_mission", reliable
        )
        self._docking_state_pub = self.create_publisher(
            String, "/saltybot/docking_state", reliable
        )
        self._status_pub = None
        if _MSGS_OK:
            self._status_pub = self.create_publisher(
                DockingStatus, "/saltybot/docking_status", reliable
            )
        self._mood_pub = None
        if _SOCIAL_OK:
            self._mood_pub = self.create_publisher(
                Mood, "/saltybot/mood", reliable
            )

        # ── Services ─────────────────────────────────────────────────────────
        if _MSGS_OK:
            self.create_service(Dock,   "/saltybot/dock",   self._dock_srv)
            self.create_service(Undock, "/saltybot/undock", self._undock_srv)

        # ── Nav2 action client (optional) ────────────────────────────────────
        self._nav2_client = None
        if _NAV2_OK:
            self._nav2_client = ActionClient(self, NavigateToPose, "navigate_to_pose")

        # ── Timer ────────────────────────────────────────────────────────────
        rate = p["control_rate"]
        self._timer = self.create_timer(1.0 / rate, self._control_cb)

        self.get_logger().info(
            f"DockingNode ready  rate={rate}Hz  "
            f"nav2={'yes' if _NAV2_OK else 'no'}  "
            f"aruco={'yes' if _CV2_OK else 'no'}"
        )

    # ── Parameters ────────────────────────────────────────────────────────────

    def _declare_params(self) -> None:
        self.declare_parameter("control_rate",          20.0)
        self.declare_parameter("image_topic",           "/camera/image_raw")
        self.declare_parameter("camera_info_topic",     "/camera/camera_info")
        self.declare_parameter("ir_topic",              "/saltybot/ir_beacon")
        self.declare_parameter("battery_topic",         "/saltybot/battery_state")
        self.declare_parameter("aruco_marker_id",       42)
        self.declare_parameter("aruco_marker_size_m",   0.10)
        self.declare_parameter("ir_threshold",          0.50)
        self.declare_parameter("battery_low_pct",       20.0)
        self.declare_parameter("battery_high_pct",      80.0)
        self.declare_parameter("servo_range_m",         1.00)
        self.declare_parameter("k_linear",              0.30)
        self.declare_parameter("k_angular",             0.80)
        self.declare_parameter("lateral_tol_m",         0.005)
        self.declare_parameter("contact_distance_m",    0.05)
        self.declare_parameter("max_linear_ms",         0.10)
        self.declare_parameter("max_angular_rads",      0.30)
        self.declare_parameter("undock_speed_ms",      -0.20)
        self.declare_parameter("undock_ticks_max",      50)
        self.declare_parameter("lost_ticks_max",        40)
        self.declare_parameter("contact_timeout_ticks", 20)
        self.declare_parameter("dock_pose",     [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0])
        self.declare_parameter("pre_dock_pose", [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0])

    def _load_params(self) -> dict:
        g = self.get_parameter
        keys = [
            "control_rate",
            "image_topic", "camera_info_topic", "ir_topic", "battery_topic",
            "aruco_marker_id", "aruco_marker_size_m", "ir_threshold",
            "battery_low_pct", "battery_high_pct",
            "servo_range_m", "k_linear", "k_angular",
            "lateral_tol_m", "contact_distance_m", "max_linear_ms", "max_angular_rads",
            "undock_speed_ms", "undock_ticks_max", "lost_ticks_max", "contact_timeout_ticks",
            "dock_pose", "pre_dock_pose",
        ]
        return {k: g(k).value for k in keys}

    # ── Subscriptions ─────────────────────────────────────────────────────────

    def _cam_info_cb(self, msg: CameraInfo) -> None:
        if _CV2_OK and self._camera_matrix is None:
            self._camera_matrix = np.array(msg.k, dtype=np.float64).reshape(3, 3)
            self._dist_coeffs   = np.array(msg.d, dtype=np.float64)

    def _image_cb(self, msg: Image) -> None:
        if not _CV2_OK or self._camera_matrix is None:
            return
        try:
            img = self._ros_image_to_bgr(msg)
        except Exception:
            return
        pose = self._aruco.detect(img, self._camera_matrix, self._dist_coeffs)
        if pose is not None:
            self._latest_dock_pose = pose
            # Update visual servo for aligned flag
            servo_result = self._servo.update(pose)
            self._latest_servo_aligned = servo_result.aligned

    def _ir_cb(self, msg: Float32) -> None:
        detected = self._ir.update(msg.data)
        # IR gives presence only — supply a synthetic far DockPose if no ArUco
        if detected and self._latest_dock_pose is None:
            # Beacon says dock is ahead, distance unknown; use sentinel 5.0 m
            self._latest_dock_pose = DockPose(
                distance_m=5.0, yaw_rad=0.0, lateral_m=0.0, source="ir"
            )

    def _battery_cb(self, msg: BatteryState) -> None:
        pct = float(msg.percentage) * 100.0 if msg.percentage <= 1.0 else float(msg.percentage)
        current = float(msg.current)
        events = self._charge.update(pct, current)
        for ev in events:
            if ev == ChargeEvent.THRESHOLD_LOW:
                self.get_logger().warn(
                    f"Battery low ({pct:.1f}%) — initiating auto-dock"
                )
            elif ev == ChargeEvent.CHARGING_STARTED:
                self.get_logger().info("Charging started")
            elif ev == ChargeEvent.THRESHOLD_HIGH:
                self.get_logger().info(
                    f"Battery charged ({pct:.1f}%) — ready to undock"
                )

    # ── 20 Hz control loop ────────────────────────────────────────────────────

    def _control_cb(self) -> None:
        inp = DockingInputs(
            battery_pct=self._charge.percentage,
            is_charging=self._charge.is_charging,
            dock_pose=self._latest_dock_pose,
            servo_aligned=self._latest_servo_aligned,
            dock_requested=self._dock_requested,
            undock_requested=self._undock_requested,
            nav2_arrived=self._nav2_arrived,
            nav2_failed=self._nav2_failed,
        )
        # Consume one-shot flags
        self._dock_requested   = False
        self._undock_requested = False
        self._nav2_arrived     = False
        self._nav2_failed      = False
        # Clear dock pose so stale detections don't persist past one tick
        self._latest_dock_pose     = None
        self._latest_servo_aligned = False

        out = self._fsm.tick(inp)

        # ── State-entry side effects ──────────────────────────────────────────
        if out.state_changed:
            self.get_logger().info(f"Docking FSM → {out.state.value}")
            # Publish charging mood when docking state is reached
            if out.state == DockingState.CHARGING and self._mood_pub is not None:
                mood_msg = Mood()
                mood_msg.mood = "happy"
                mood_msg.intensity = 1.0
                self._mood_pub.publish(mood_msg)

        if out.request_nav2:
            self._send_nav2_goal()
        if out.cancel_nav2:
            self._cancel_nav2_goal()

        if out.mission_interrupted:
            self.get_logger().warn("Mission interrupted for auto-dock (battery low)")

        if out.mission_resume:
            msg = Bool()
            msg.data = True
            self._resume_pub.publish(msg)
            self.get_logger().info("Auto-dock complete — resuming mission")

        # ── Velocity commands ─────────────────────────────────────────────────
        if out.state in (DockingState.VISUAL_SERVO, DockingState.UNDOCKING):
            twist = Twist()
            twist.linear.x  = out.cmd_linear
            twist.angular.z = out.cmd_angular
            self._cmd_vel_pub.publish(twist)

        # ── Status ────────────────────────────────────────────────────────────
        # Publish docking state (Issue #489)
        state_msg = String()
        state_msg.data = out.state.value
        self._docking_state_pub.publish(state_msg)

        if self._status_pub is not None:
            self._publish_status(inp, out)

    # ── Nav2 integration ──────────────────────────────────────────────────────

    def _send_nav2_goal(self) -> None:
        if not _NAV2_OK or self._nav2_client is None:
            self.get_logger().warn("Nav2 unavailable — skipping pre-dock approach")
            self._nav2_arrived = True   # skip straight to VISUAL_SERVO next tick
            return
        if not self._nav2_client.wait_for_server(timeout_sec=1.0):
            self.get_logger().warn("Nav2 action server not ready — skipping approach")
            self._nav2_arrived = True
            return
        p = self._load_params()
        goal = NavigateToPose.Goal()
        goal.pose = _pose_param_to_stamped(
            p["pre_dock_pose"], "map", self.get_clock().now().to_msg()
        )
        send_future = self._nav2_client.send_goal_async(goal)
        send_future.add_done_callback(self._nav2_goal_response_cb)

    def _nav2_goal_response_cb(self, future) -> None:
        goal_handle = future.result()
        if not goal_handle or not goal_handle.accepted:
            self.get_logger().warn("Nav2 goal rejected")
            self._nav2_failed = True
            return
        self._nav2_goal_handle = goal_handle
        result_future = goal_handle.get_result_async()
        result_future.add_done_callback(self._nav2_result_cb)

    def _nav2_result_cb(self, future) -> None:
        result = future.result()
        if result and result.status == 4:   # SUCCEEDED
            self._nav2_arrived = True
        else:
            self._nav2_failed = True

    def _cancel_nav2_goal(self) -> None:
        if self._nav2_goal_handle is not None:
            self._nav2_goal_handle.cancel_goal_async()
            self._nav2_goal_handle = None

    # ── Services ──────────────────────────────────────────────────────────────

    def _dock_srv(self, request, response):
        if self._fsm.state not in (DockingState.IDLE,) and not request.force:
            response.success = False
            response.message = f"Already in state {self._fsm.state.value}; use force=True to override"
            return response
        self._dock_requested = True
        response.success = True
        response.message = "Docking initiated"
        return response

    def _undock_srv(self, request, response):
        if self._fsm.state == DockingState.IDLE:
            response.success = False
            response.message = "Already idle"
            return response
        self._undock_requested = True
        response.success = True
        response.message = "Undocking initiated"
        return response

    # ── Publishing ────────────────────────────────────────────────────────────

    def _publish_status(self, inp: DockingInputs, out) -> None:
        msg = DockingStatus()
        msg.stamp         = self.get_clock().now().to_msg()
        msg.state         = out.state.value
        msg.dock_detected = inp.dock_pose is not None
        msg.distance_m    = float(inp.dock_pose.distance_m) if inp.dock_pose else float("nan")
        msg.lateral_error_m = float(inp.dock_pose.lateral_m) if inp.dock_pose else float("nan")
        msg.battery_pct   = float(self._charge.percentage)
        msg.charging      = self._charge.is_charging
        msg.aligned       = inp.servo_aligned
        self._status_pub.publish(msg)

    # ── Utilities ─────────────────────────────────────────────────────────────

    @staticmethod
    def _ros_image_to_bgr(msg: Image):
        """Convert sensor_msgs/Image to cv2 BGR numpy array."""
        import numpy as np
        data = np.frombuffer(msg.data, dtype=np.uint8)
        if msg.encoding in ("bgr8", "rgb8"):
            img = data.reshape(msg.height, msg.width, 3)
            if msg.encoding == "rgb8":
                img = img[:, :, ::-1].copy()
        elif msg.encoding == "mono8":
            img = data.reshape(msg.height, msg.width)
        else:
            img = data.reshape(msg.height, msg.width, -1)
        return img


# ── Entry point ───────────────────────────────────────────────────────────────

def main(args=None):
    rclpy.init(args=args)
    node = DockingNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.try_shutdown()


if __name__ == "__main__":
    main()
