#!/usr/bin/env python3
"""UWB Geofence Speed Limiter for SaltyBot (Issue #657).

Subscribes to:
  /saltybot/pose/authoritative  geometry_msgs/PoseWithCovarianceStamped
  /cmd_vel                      geometry_msgs/Twist

Publishes:
  /cmd_vel_limited              geometry_msgs/Twist   -- speed-limited output
  /saltybot/geofence/status     std_msgs/String       -- JSON telemetry
  /saltybot/geofence_violation  std_msgs/Bool         -- triggers e-stop cascade

Zones are defined in a YAML file (zones_file parameter).  Each zone has a
name, polygon vertices in UWB local frame (metres), max_speed (m/s) and
max_angular (rad/s).  The most restrictive of all containing zones wins.

An outer emergency_boundary polygon is mandatory.  If the robot exits it,
all motion is zeroed and /saltybot/geofence_violation is latched True.

Speed transitions are smoothed over `transition_time` seconds using a
linear ramp so that passing through a zone boundary never causes a
velocity spike.
"""

from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import yaml
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist
from geometry_msgs.msg import PoseWithCovarianceStamped
from std_msgs.msg import Bool, String

from saltybot_uwb_geofence.zone_checker import (
    apply_speed_limit,
    min_dist_to_boundary,
    parse_flat_vertices,
    point_in_polygon,
    ramp_toward,
)


class UwbGeofenceNode(Node):
    """ROS2 node enforcing UWB-position-based speed limits inside geofence zones."""

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

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter(
            "zones_file",
            str(Path(__file__).parent.parent.parent / "config" / "uwb_geofence_zones.yaml"),
        )
        self.declare_parameter("transition_time", 0.5)   # seconds
        self.declare_parameter("frequency", 20.0)         # Hz
        self.declare_parameter("default_max_speed", 1.0)  # m/s (outside all named zones)
        self.declare_parameter("default_max_angular", 1.5)  # rad/s
        self.declare_parameter("pose_topic", "/saltybot/pose/authoritative")
        self.declare_parameter("cmd_vel_topic", "/cmd_vel")

        zones_file = self.get_parameter("zones_file").value
        self._transition_time: float = self.get_parameter("transition_time").value
        frequency: float = self.get_parameter("frequency").value
        self._default_max_speed: float = self.get_parameter("default_max_speed").value
        self._default_max_angular: float = self.get_parameter("default_max_angular").value
        pose_topic: str = self.get_parameter("pose_topic").value
        cmd_vel_topic: str = self.get_parameter("cmd_vel_topic").value

        # ── Load zones config ─────────────────────────────────────────────────
        self._zones: List[Dict[str, Any]] = []
        self._emergency_boundary: List[Tuple[float, float]] = []
        self._load_zones(zones_file)

        # ── Runtime state ─────────────────────────────────────────────────────
        self._robot_x: float = 0.0
        self._robot_y: float = 0.0
        self._pose_received: bool = False
        self._last_cmd_vel: Optional[Twist] = None

        # Current smoothed speed limits (ramp toward target)
        self._current_max_speed: float = self._default_max_speed
        self._current_max_angular: float = self._default_max_angular

        # E-stop latch (requires operator reset)
        self._violation_latched: bool = False

        # ── Subscriptions ─────────────────────────────────────────────────────
        self.create_subscription(
            PoseWithCovarianceStamped,
            pose_topic,
            self._on_pose,
            10,
        )
        self.create_subscription(
            Twist,
            cmd_vel_topic,
            self._on_cmd_vel,
            10,
        )

        # ── Publications ──────────────────────────────────────────────────────
        self._pub_limited = self.create_publisher(Twist, "/cmd_vel_limited", 10)
        self._pub_status = self.create_publisher(String, "/saltybot/geofence/status", 10)
        self._pub_violation = self.create_publisher(Bool, "/saltybot/geofence_violation", 1)

        # ── Timer ─────────────────────────────────────────────────────────────
        self._dt: float = 1.0 / frequency
        self.create_timer(self._dt, self._tick)

        self.get_logger().info(
            f"uwb_geofence started: {len(self._zones)} zone(s), "
            f"freq={frequency:.0f}Hz, transition={self._transition_time}s"
        )

    # ── Config loading ────────────────────────────────────────────────────────

    def _load_zones(self, path: str) -> None:
        """Parse YAML zones file and populate self._zones and self._emergency_boundary."""
        p = Path(path)
        if not p.exists():
            self.get_logger().warn(f"zones_file not found: {path}; running with empty zones")
            return

        with open(p) as fh:
            cfg = yaml.safe_load(fh)

        # Emergency boundary (mandatory)
        eb = cfg.get("emergency_boundary", {})
        flat = eb.get("vertices", [])
        if flat:
            try:
                self._emergency_boundary = parse_flat_vertices(flat)
            except ValueError as exc:
                self.get_logger().error(f"emergency_boundary vertices invalid: {exc}")
        else:
            self.get_logger().warn("No emergency_boundary defined — boundary checking disabled")

        # Named speed zones
        for z in cfg.get("zones", []):
            try:
                verts = parse_flat_vertices(z["vertices"])
            except (KeyError, ValueError) as exc:
                self.get_logger().error(f"Skipping zone '{z.get('name', '?')}': {exc}")
                continue
            self._zones.append(
                {
                    "name": z.get("name", "unnamed"),
                    "description": z.get("description", ""),
                    "max_speed": float(z.get("max_speed", self._default_max_speed)),
                    "max_angular": float(z.get("max_angular", self._default_max_angular)),
                    "vertices": verts,
                }
            )
            self.get_logger().info(
                f"  zone '{z['name']}': max_speed={z.get('max_speed')} m/s, "
                f"{len(verts)} vertices"
            )

    # ── Callbacks ─────────────────────────────────────────────────────────────

    def _on_pose(self, msg: PoseWithCovarianceStamped) -> None:
        self._robot_x = msg.pose.pose.position.x
        self._robot_y = msg.pose.pose.position.y
        self._pose_received = True

    def _on_cmd_vel(self, msg: Twist) -> None:
        self._last_cmd_vel = msg

    # ── Main loop ────────────────────────────────────────────────────────────

    def _tick(self) -> None:
        """Called at `frequency` Hz.  Compute limits, ramp, publish."""
        if not self._pose_received or self._last_cmd_vel is None:
            return

        x, y = self._robot_x, self._robot_y

        # ── Emergency boundary check ──────────────────────────────────────────
        if self._emergency_boundary:
            inside_boundary = point_in_polygon(x, y, self._emergency_boundary)
            if not inside_boundary:
                self._trigger_violation(x, y)
                return

        # Clear latch once back inside boundary (manual operator resume)
        if self._violation_latched:
            # Publish safe=False until operator clears via /saltybot/geofence_violation
            # (this node keeps publishing zero until reset)
            self._publish_zero_cmd()
            self._pub_violation.publish(Bool(data=True))
            return

        # ── Zone classification ───────────────────────────────────────────────
        active_zones: List[str] = []
        target_max_speed = self._default_max_speed
        target_max_angular = self._default_max_angular

        for zone in self._zones:
            if point_in_polygon(x, y, zone["vertices"]):
                active_zones.append(zone["name"])
                if zone["max_speed"] < target_max_speed:
                    target_max_speed = zone["max_speed"]
                if zone["max_angular"] < target_max_angular:
                    target_max_angular = zone["max_angular"]

        # ── Smooth ramp toward target limits ──────────────────────────────────
        ramp_rate_speed = (self._default_max_speed / max(self._transition_time, 0.01))
        ramp_rate_angular = (self._default_max_angular / max(self._transition_time, 0.01))

        self._current_max_speed = ramp_toward(
            self._current_max_speed, target_max_speed, ramp_rate_speed * self._dt
        )
        self._current_max_angular = ramp_toward(
            self._current_max_angular, target_max_angular, ramp_rate_angular * self._dt
        )

        # ── Apply limits to cmd_vel ───────────────────────────────────────────
        cmd = self._last_cmd_vel
        lx, ly, lz, ax, ay, wz = apply_speed_limit(
            cmd.linear.x, cmd.linear.y, cmd.angular.z,
            self._current_max_speed, self._current_max_angular,
            linear_z=cmd.linear.z, angular_x=cmd.angular.x, angular_y=cmd.angular.y,
        )
        limited = Twist()
        limited.linear.x = lx
        limited.linear.y = ly
        limited.linear.z = lz
        limited.angular.x = ax
        limited.angular.y = ay
        limited.angular.z = wz
        self._pub_limited.publish(limited)

        # ── Status ───────────────────────────────────────────────────────────
        dist_to_boundary = (
            min_dist_to_boundary(x, y, self._emergency_boundary)
            if self._emergency_boundary
            else -1.0
        )
        status = {
            "robot_xy": [round(x, 3), round(y, 3)],
            "active_zones": active_zones,
            "current_max_speed": round(self._current_max_speed, 3),
            "current_max_angular": round(self._current_max_angular, 3),
            "target_max_speed": round(target_max_speed, 3),
            "dist_to_emergency_boundary": round(dist_to_boundary, 3),
            "violation": False,
        }
        self._pub_status.publish(String(data=json.dumps(status)))
        self._pub_violation.publish(Bool(data=False))

    # ── Helpers ───────────────────────────────────────────────────────────────

    def _trigger_violation(self, x: float, y: float) -> None:
        """Robot has exited emergency boundary.  Zero all motion, latch violation."""
        if not self._violation_latched:
            self.get_logger().error(
                f"GEOFENCE VIOLATION: robot ({x:.3f}, {y:.3f}) is outside emergency boundary!"
            )
            self._violation_latched = True
        self._publish_zero_cmd()
        self._pub_violation.publish(Bool(data=True))
        status = {
            "robot_xy": [round(x, 3), round(y, 3)],
            "active_zones": [],
            "current_max_speed": 0.0,
            "current_max_angular": 0.0,
            "target_max_speed": 0.0,
            "dist_to_emergency_boundary": 0.0,
            "violation": True,
        }
        self._pub_status.publish(String(data=json.dumps(status)))

    def _publish_zero_cmd(self) -> None:
        self._pub_limited.publish(Twist())



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


if __name__ == "__main__":
    main()
