#!/usr/bin/env python3
"""gimbal_node.py — ROS2 gimbal control node for SaltyBot pan/tilt camera head (Issue #548).

<<<<<<< HEAD
Controls pan/tilt gimbal via JLINK binary protocol over serial to ESP32 BALANCE.
=======
Controls pan/tilt gimbal via JLINK binary protocol over serial to ESP32-S3.
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
Implements smooth trapezoidal motion profiles with configurable axis limits.

Subscribed topics:
  /saltybot/gimbal/cmd       (geometry_msgs/Vector3)  x=pan_deg, y=tilt_deg, z=speed_deg_s
  /saltybot/gimbal/look_at_target (geometry_msgs/PointStamped) 3D look-at target (camera frame)

Published topics:
  /saltybot/gimbal/state     (std_msgs/String)    JSON: pan_deg, tilt_deg, moving, fault, …
  /saltybot/gimbal/cmd_echo  (geometry_msgs/Vector3)  current target cmd (monitoring)

Services:
  /saltybot/gimbal/home      (std_srvs/Trigger)   command gimbal to home position
  /saltybot/gimbal/look_at   (std_srvs/Trigger)   execute look-at to last look_at_target

Parameters (config/gimbal_params.yaml):
  serial_port          /dev/ttyTHS1   JLINK serial device
  baud_rate            921600
  pan_limit_deg        150.0          soft limit ± degrees
  tilt_limit_deg        45.0          soft limit ± degrees
  home_pan_deg           0.0          home pan position
  home_tilt_deg          0.0          home tilt position
  max_speed_deg_s       90.0          default motion speed
  accel_deg_s2         180.0          trapezoidal acceleration
  update_rate_hz        20.0          motion profile update rate
  state_publish_hz      10.0          /saltybot/gimbal/state publish rate
  reconnect_delay_s      2.0          serial reconnect interval
  camera_focal_length_px 600.0        D435i focal length (for look_at projection)
  image_width_px         848
  image_height_px        480
"""

from __future__ import annotations

import json
import math
import threading
import time
from typing import Optional

import rclpy
from rclpy.node import Node
from rclpy.qos import HistoryPolicy, QoSProfile, ReliabilityPolicy
import serial

from geometry_msgs.msg import PointStamped, Vector3
from std_msgs.msg import String
from std_srvs.srv import Trigger

from .jlink_gimbal import (
    encode_gimbal_pos,
    encode_gimbal_home,
    decode_gimbal_state,
    TLM_GIMBAL_STATE,
    GimbalStateFrame,
    STX, ETX,
    _crc16_ccitt,
)


def _clamp(v: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, v))


class MotionAxis:
    """Single-axis trapezoidal velocity profile.

    Tracks current position and velocity; each tick() call advances the
    axis toward the target using bounded acceleration and velocity.
    """

    def __init__(self, initial_pos: float, max_speed: float, accel: float) -> None:
        self.pos = initial_pos
        self.vel = 0.0
        self.target = initial_pos
        self.max_speed = max_speed
        self.accel = accel

    def set_target(self, target: float, speed: Optional[float] = None) -> None:
        self.target = target
        if speed is not None:
            self.max_speed = max(1.0, speed)

    def tick(self, dt: float) -> float:
        """Advance one timestep. Returns new position."""
        error = self.target - self.pos
        if abs(error) < 0.01:
            self.pos = self.target
            self.vel = 0.0
            return self.pos

        # Deceleration distance at current speed: v²/(2a)
        decel_dist = (self.vel ** 2) / (2.0 * self.accel) if self.accel > 0 else 0.0
        direction = 1.0 if error > 0 else -1.0

        if abs(error) <= decel_dist + 0.01:
            # Deceleration phase
            desired_vel = direction * math.sqrt(max(0.0, 2.0 * self.accel * abs(error)))
        else:
            # Acceleration/cruise phase
            desired_vel = direction * self.max_speed

        # Apply acceleration limit
        dv = desired_vel - self.vel
        dv = _clamp(dv, -self.accel * dt, self.accel * dt)
        self.vel += dv
        self.vel = _clamp(self.vel, -self.max_speed, self.max_speed)

        self.pos += self.vel * dt
        # Don't overshoot
        if direction > 0 and self.pos > self.target:
            self.pos = self.target
            self.vel = 0.0
        elif direction < 0 and self.pos < self.target:
            self.pos = self.target
            self.vel = 0.0

        return self.pos

    @property
    def is_moving(self) -> bool:
        return abs(self.target - self.pos) > 0.05 or abs(self.vel) > 0.05


class JLinkGimbalSerial:
    """Serial connection manager for JLINK gimbal frames.

    Handles connect/disconnect/reconnect and provides send() and
    read_state() methods. Thread-safe via internal lock.
    """

    _PARSER_WAIT_STX  = 0
    _PARSER_WAIT_CMD  = 1
    _PARSER_WAIT_LEN  = 2
    _PARSER_PAYLOAD   = 3
    _PARSER_CRC_HI    = 4
    _PARSER_CRC_LO    = 5
    _PARSER_WAIT_ETX  = 6

    def __init__(self, port: str, baud: int) -> None:
        self._port = port
        self._baud = baud
        self._ser: Optional[serial.Serial] = None
        self._lock = threading.Lock()
        self._parser_state = self._PARSER_WAIT_STX
        self._parser_cmd = 0
        self._parser_len = 0
        self._parser_payload = bytearray()
        self._parser_crc_rcvd = 0

    def connect(self) -> bool:
        try:
            self._ser = serial.Serial(
                port=self._port,
                baudrate=self._baud,
                timeout=0.05,
                bytesize=8,
                parity=serial.PARITY_NONE,
                stopbits=1,
            )
            self._reset_parser()
            return True
        except Exception as e:
            print(f"[gimbal] serial connect failed ({self._port}): {e}")
            self._ser = None
            return False

    def disconnect(self) -> None:
        with self._lock:
            if self._ser:
                self._ser.close()
                self._ser = None

    def is_connected(self) -> bool:
        return self._ser is not None and self._ser.is_open

    def send(self, frame: bytes) -> bool:
        with self._lock:
            if not self.is_connected():
                return False
            try:
                self._ser.write(frame)
                return True
            except Exception as e:
                print(f"[gimbal] serial write error: {e}")
                self._ser = None
                return False

    def read_pending(self) -> Optional[GimbalStateFrame]:
        """Drain RX buffer; return GimbalStateFrame if a complete telemetry
        frame was received, or None if no complete frame is available."""
        with self._lock:
            if not self.is_connected():
                return None
            try:
                raw = self._ser.read(self._ser.in_waiting or 1)
            except Exception:
                self._ser = None
                return None

        result = None
        for b in raw:
            frame = self._feed(b)
            if frame is not None:
                result = frame  # keep last complete frame
        return result

    def _reset_parser(self) -> None:
        self._parser_state = self._PARSER_WAIT_STX
        self._parser_payload = bytearray()

    def _feed(self, byte: int) -> Optional[GimbalStateFrame]:
        s = self._parser_state

        if s == self._PARSER_WAIT_STX:
            if byte == STX:
                self._parser_state = self._PARSER_WAIT_CMD
            return None

        if s == self._PARSER_WAIT_CMD:
            self._parser_cmd = byte
            self._parser_state = self._PARSER_WAIT_LEN
            return None

        if s == self._PARSER_WAIT_LEN:
            self._parser_len = byte
            self._parser_payload = bytearray()
            if byte > 64:
                self._reset_parser()
                return None
            self._parser_state = self._PARSER_PAYLOAD if byte > 0 else self._PARSER_CRC_HI
            return None

        if s == self._PARSER_PAYLOAD:
            self._parser_payload.append(byte)
            if len(self._parser_payload) == self._parser_len:
                self._parser_state = self._PARSER_CRC_HI
            return None

        if s == self._PARSER_CRC_HI:
            self._parser_crc_rcvd = byte << 8
            self._parser_state = self._PARSER_CRC_LO
            return None

        if s == self._PARSER_CRC_LO:
            self._parser_crc_rcvd |= byte
            self._parser_state = self._PARSER_WAIT_ETX
            return None

        if s == self._PARSER_WAIT_ETX:
            self._reset_parser()
            if byte != ETX:
                return None
            crc_data = bytes([self._parser_cmd, self._parser_len]) + self._parser_payload
            if _crc16_ccitt(crc_data) != self._parser_crc_rcvd:
                return None
            if self._parser_cmd == TLM_GIMBAL_STATE:
                return decode_gimbal_state(bytes(self._parser_payload))
            return None

        self._reset_parser()
        return None


class GimbalNode(Node):
    """ROS2 gimbal control node: smooth motion profiles + JLINK serial bridge."""

    def __init__(self) -> None:
        super().__init__("gimbal_node")

        # ── Parameters ─────────────────────────────────────────────────────
        self.declare_parameter("serial_port",           "/dev/ttyTHS1")
        self.declare_parameter("baud_rate",             921600)
        self.declare_parameter("pan_limit_deg",         150.0)
        self.declare_parameter("tilt_limit_deg",         45.0)
        self.declare_parameter("home_pan_deg",            0.0)
        self.declare_parameter("home_tilt_deg",           0.0)
        self.declare_parameter("max_speed_deg_s",        90.0)
        self.declare_parameter("accel_deg_s2",          180.0)
        self.declare_parameter("update_rate_hz",         20.0)
        self.declare_parameter("state_publish_hz",       10.0)
        self.declare_parameter("reconnect_delay_s",       2.0)
        self.declare_parameter("camera_focal_length_px", 600.0)
        self.declare_parameter("image_width_px",          848)
        self.declare_parameter("image_height_px",         480)

        self._port           = self.get_parameter("serial_port").value
        self._baud           = self.get_parameter("baud_rate").value
        self._pan_limit      = self.get_parameter("pan_limit_deg").value
        self._tilt_limit     = self.get_parameter("tilt_limit_deg").value
        self._home_pan       = self.get_parameter("home_pan_deg").value
        self._home_tilt      = self.get_parameter("home_tilt_deg").value
        self._max_speed      = self.get_parameter("max_speed_deg_s").value
        self._accel          = self.get_parameter("accel_deg_s2").value
        update_hz            = self.get_parameter("update_rate_hz").value
        state_hz             = self.get_parameter("state_publish_hz").value
        self._reconnect_del  = self.get_parameter("reconnect_delay_s").value
        self._focal_px       = self.get_parameter("camera_focal_length_px").value
        self._img_w          = self.get_parameter("image_width_px").value
        self._img_h          = self.get_parameter("image_height_px").value

        # ── State ──────────────────────────────────────────────────────────
        self._pan_axis  = MotionAxis(self._home_pan,  self._max_speed, self._accel)
        self._tilt_axis = MotionAxis(self._home_tilt, self._max_speed, self._accel)

        self._hw_state:  Optional[GimbalStateFrame] = None
        self._look_at_target: Optional[PointStamped] = None
        self._last_cmd_time = time.monotonic()
        self._reconnect_ts  = 0.0
        self._state_lock    = threading.Lock()

        # ── Serial ─────────────────────────────────────────────────────────
        self._serial = JLinkGimbalSerial(self._port, self._baud)
        if not self._serial.connect():
            self.get_logger().warn(
                f"[gimbal] serial not available on {self._port} — will retry"
            )

        # ── Subscribers ────────────────────────────────────────────────────
        qos = QoSProfile(
            history=HistoryPolicy.KEEP_LAST,
            depth=5,
            reliability=ReliabilityPolicy.BEST_EFFORT,
        )
        self.create_subscription(Vector3, "/saltybot/gimbal/cmd",
                                 self._on_cmd, qos)
        self.create_subscription(PointStamped, "/saltybot/gimbal/look_at_target",
                                 self._on_look_at_target, 10)

        # ── Publishers ─────────────────────────────────────────────────────
        self._pub_state    = self.create_publisher(String,  "/saltybot/gimbal/state",    10)
        self._pub_cmd_echo = self.create_publisher(Vector3, "/saltybot/gimbal/cmd_echo", 10)

        # ── Services ───────────────────────────────────────────────────────
        self.create_service(Trigger, "/saltybot/gimbal/home",    self._svc_home)
        self.create_service(Trigger, "/saltybot/gimbal/look_at", self._svc_look_at)

        # ── Timers ─────────────────────────────────────────────────────────
        self.create_timer(1.0 / update_hz, self._motion_tick)
        self.create_timer(1.0 / state_hz,  self._publish_state)

        self.get_logger().info(
            f"[gimbal] node ready — port={self._port} "
            f"pan=±{self._pan_limit}° tilt=±{self._tilt_limit}° "
            f"speed={self._max_speed}°/s accel={self._accel}°/s²"
        )

    # ── Subscriber callbacks ───────────────────────────────────────────────

    def _on_cmd(self, msg: Vector3) -> None:
        """Handle /saltybot/gimbal/cmd (x=pan, y=tilt, z=speed_deg_s)."""
        pan   = _clamp(msg.x, -self._pan_limit,  self._pan_limit)
        tilt  = _clamp(msg.y, -self._tilt_limit, self._tilt_limit)
        speed = msg.z if msg.z > 0.0 else self._max_speed
        speed = _clamp(speed, 1.0, self._max_speed)
        with self._state_lock:
            self._pan_axis.set_target(pan,  speed)
            self._tilt_axis.set_target(tilt, speed)
        self._last_cmd_time = time.monotonic()

    def _on_look_at_target(self, msg: PointStamped) -> None:
        """Cache look-at 3D target for use by /look_at service."""
        self._look_at_target = msg

    # ── Service callbacks ──────────────────────────────────────────────────

    def _svc_home(self, _req: Trigger.Request, resp: Trigger.Response) -> Trigger.Response:
        """Return gimbal to home position."""
        with self._state_lock:
            self._pan_axis.set_target(self._home_pan,  self._max_speed)
            self._tilt_axis.set_target(self._home_tilt, self._max_speed)
        self._serial.send(encode_gimbal_home())
        resp.success = True
        resp.message = f"Homing to pan={self._home_pan}° tilt={self._home_tilt}°"
        self.get_logger().info("[gimbal] home command")
        return resp

    def _svc_look_at(self, _req: Trigger.Request, resp: Trigger.Response) -> Trigger.Response:
        """Convert last look_at_target point to pan/tilt and command motion."""
        tgt = self._look_at_target
        if tgt is None:
            resp.success = False
            resp.message = "No look_at_target received on /saltybot/gimbal/look_at_target"
            return resp

        # Project 3D point (camera frame) → pan/tilt angles.
        # Assumes point is in camera_depth_optical_frame (Z=forward, X=right, Y=down).
        x, y, z = tgt.point.x, tgt.point.y, tgt.point.z
        if z <= 0.0:
            resp.success = False
            resp.message = f"Invalid look_at_target: z={z:.3f} (must be > 0)"
            return resp

        pan_deg  =  math.degrees(math.atan2(x, z))
        tilt_deg = -math.degrees(math.atan2(y, z))  # negative: down in image = tilt up

        pan_deg  = _clamp(pan_deg,  -self._pan_limit,  self._pan_limit)
        tilt_deg = _clamp(tilt_deg, -self._tilt_limit, self._tilt_limit)

        with self._state_lock:
            self._pan_axis.set_target(pan_deg,  self._max_speed)
            self._tilt_axis.set_target(tilt_deg, self._max_speed)

        resp.success = True
        resp.message = (
            f"Looking at ({x:.2f}, {y:.2f}, {z:.2f}) → "
            f"pan={pan_deg:.1f}° tilt={tilt_deg:.1f}°"
        )
        self.get_logger().info(f"[gimbal] look_at: {resp.message}")
        return resp

    # ── Motion tick ────────────────────────────────────────────────────────

    def _motion_tick(self) -> None:
        """Advance motion profiles and send JLINK frame at update_rate_hz."""
        # Reconnect if needed
        if not self._serial.is_connected():
            now = time.monotonic()
            if now - self._reconnect_ts >= self._reconnect_del:
                self._reconnect_ts = now
                if self._serial.connect():
                    self.get_logger().info(f"[gimbal] serial reconnected on {self._port}")

        dt = 1.0 / self.get_parameter("update_rate_hz").value

        with self._state_lock:
            pan  = self._pan_axis.tick(dt)
            tilt = self._tilt_axis.tick(dt)
            moving = self._pan_axis.is_moving or self._tilt_axis.is_moving
            speed = max(self._pan_axis.max_speed, self._tilt_axis.max_speed)

        # Send JLINK frame
        frame = encode_gimbal_pos(pan, tilt, speed)
        self._serial.send(frame)

        # Read telemetry
        hw = self._serial.read_pending()
        if hw is not None:
            with self._state_lock:
                self._hw_state = hw

        # Echo command
        echo = Vector3(x=pan, y=tilt, z=speed)
        self._pub_cmd_echo.publish(echo)

    # ── State publisher ────────────────────────────────────────────────────

    def _publish_state(self) -> None:
        """Publish gimbal state as JSON string."""
        with self._state_lock:
            pan    = self._pan_axis.pos
            tilt   = self._tilt_axis.pos
            pan_t  = self._pan_axis.target
            tilt_t = self._tilt_axis.target
            moving = self._pan_axis.is_moving or self._tilt_axis.is_moving
            hw     = self._hw_state

        state: dict = {
            "pan_deg":       round(pan,  2),
            "tilt_deg":      round(tilt, 2),
            "pan_target_deg":  round(pan_t,  2),
            "tilt_target_deg": round(tilt_t, 2),
            "moving":        moving,
            "serial_ok":     self._serial.is_connected(),
        }
        if hw is not None:
            state["hw_pan_deg"]        = round(hw.pan_deg, 1)
            state["hw_tilt_deg"]       = round(hw.tilt_deg, 1)
            state["hw_pan_speed_raw"]  = hw.pan_speed_raw
            state["hw_tilt_speed_raw"] = hw.tilt_speed_raw
            state["hw_torque_en"]      = hw.torque_en
            state["hw_rx_err_pct"]     = hw.rx_err_pct

        self._pub_state.publish(String(data=json.dumps(state)))

    def destroy_node(self) -> None:
        self._serial.disconnect()
        super().destroy_node()


def main(args=None) -> None:
    rclpy.init(args=args)
    node = GimbalNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
