"""
social_nav_node.py -- Social navigation node for saltybot.

Orchestrates person following with multiple modes, voice commands,
waypoint teaching/replay, and A* obstacle avoidance.

Follow modes:
  shadow  -- stay directly behind at follow_distance
  lead    -- move ahead of person
  side    -- stay beside (default right)
  orbit   -- circle around person
  loose   -- relaxed follow with deadband
  tight   -- close follow

Waypoint teaching:
  Voice command "teach route <name>" -> record mode ON
  Voice command "stop teaching" -> save route
  Voice command "replay route <name>" -> playback

Voice commands:
  "follow me" / "shadow" -> SHADOW mode
  "lead me" / "go ahead" -> LEAD mode
  "stay beside" -> SIDE mode
  "orbit" -> ORBIT mode
  "give me space" -> LOOSE mode
  "stay close" -> TIGHT mode
  "stop" / "halt" -> STOP
  "go home" -> navigate to home waypoint
  "teach route <name>" -> start recording
  "stop teaching" -> finish recording
  "replay route <name>" -> playback recorded route
"""

import math
import time
import re
from collections import deque

import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
from geometry_msgs.msg import Twist, PoseStamped
from nav_msgs.msg import OccupancyGrid, Odometry
from std_msgs.msg import String, Int32

from .follow_modes import (
    FollowMode, MODE_VOICE_COMMANDS,
    compute_shadow_target, compute_lead_target, compute_side_target,
    compute_orbit_target, compute_loose_target, compute_tight_target,
)
from .astar import astar, inflate_obstacles
from .waypoint_teacher import WaypointRoute, WaypointReplayer

# Try importing social msgs; fallback gracefully
try:
    from saltybot_social_msgs.msg import PersonStateArray
    _HAS_SOCIAL_MSGS = True
except ImportError:
    _HAS_SOCIAL_MSGS = False

# Proportional controller gains
_K_ANG = 2.0   # angular gain
_K_LIN = 0.8   # linear gain
_HIGH_SPEED_THRESHOLD = 3.0  # m/s person velocity triggers fast mode
_PREDICT_AHEAD_S = 0.3       # seconds to extrapolate position
_TEACH_MIN_DIST = 0.5        # meters between recorded waypoints


class SocialNavNode(Node):
    """Main social navigation orchestrator."""

    def __init__(self):
        super().__init__('social_nav')

        # -- Parameters --
        self.declare_parameter('follow_mode', 'shadow')
        self.declare_parameter('follow_distance', 1.2)
        self.declare_parameter('lead_distance', 2.0)
        self.declare_parameter('orbit_radius', 1.5)
        self.declare_parameter('max_linear_speed', 1.0)
        self.declare_parameter('max_linear_speed_fast', 5.5)
        self.declare_parameter('max_angular_speed', 1.0)
        self.declare_parameter('goal_tolerance', 0.3)
        self.declare_parameter('routes_dir', '/mnt/nvme/saltybot/routes')
        self.declare_parameter('home_x', 0.0)
        self.declare_parameter('home_y', 0.0)
        self.declare_parameter('map_resolution', 0.05)
        self.declare_parameter('obstacle_inflation_cells', 3)

        self._follow_mode = FollowMode(
            self.get_parameter('follow_mode').value)
        self._follow_distance = self.get_parameter('follow_distance').value
        self._lead_distance = self.get_parameter('lead_distance').value
        self._orbit_radius = self.get_parameter('orbit_radius').value
        self._max_lin = self.get_parameter('max_linear_speed').value
        self._max_lin_fast = self.get_parameter('max_linear_speed_fast').value
        self._max_ang = self.get_parameter('max_angular_speed').value
        self._goal_tol = self.get_parameter('goal_tolerance').value
        self._routes_dir = self.get_parameter('routes_dir').value
        self._home_x = self.get_parameter('home_x').value
        self._home_y = self.get_parameter('home_y').value
        self._map_resolution = self.get_parameter('map_resolution').value
        self._inflation_cells = self.get_parameter(
            'obstacle_inflation_cells').value

        # -- State --
        self._robot_x = 0.0
        self._robot_y = 0.0
        self._robot_yaw = 0.0
        self._target_person_pos = None   # (x, y, z)
        self._target_person_bearing = 0.0
        self._target_person_id = -1
        self._person_history = deque(maxlen=5)  # for velocity estimation
        self._stopped = False
        self._go_home = False

        # Occupancy grid for A*
        self._occ_grid = None
        self._occ_origin = (0.0, 0.0)
        self._occ_resolution = 0.05
        self._astar_path = None

        # Orbit state
        self._orbit_angle = 0.0

        # Waypoint teaching / replay
        self._teaching = False
        self._current_route = None
        self._last_teach_x = None
        self._last_teach_y = None
        self._replayer = None

        # -- QoS profiles --
        best_effort_qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST, depth=1)
        reliable_qos = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            history=HistoryPolicy.KEEP_LAST, depth=1)

        # -- Subscriptions --
        if _HAS_SOCIAL_MSGS:
            self.create_subscription(
                PersonStateArray, '/social/persons',
                self._on_persons, best_effort_qos)
        else:
            self.get_logger().warn(
                'saltybot_social_msgs not found; '
                'using /person/target fallback')

        self.create_subscription(
            PoseStamped, '/person/target',
            self._on_person_target, best_effort_qos)
        self.create_subscription(
            String, '/social/speech/command',
            self._on_voice_command, 10)
        self.create_subscription(
            String, '/social/speech/transcript',
            self._on_transcript, 10)
        self.create_subscription(
            OccupancyGrid, '/map',
            self._on_map, reliable_qos)
        self.create_subscription(
            Odometry, '/odom',
            self._on_odom, best_effort_qos)
        self.create_subscription(
            Int32, '/social/attention/target_id',
            self._on_target_id, 10)

        # -- Publishers --
        self._cmd_vel_pub = self.create_publisher(
            Twist, '/cmd_vel', best_effort_qos)
        self._mode_pub = self.create_publisher(
            String, '/social/nav/mode', reliable_qos)
        self._target_pub = self.create_publisher(
            PoseStamped, '/social/nav/target_pos', 10)
        self._status_pub = self.create_publisher(
            String, '/social/nav/status', best_effort_qos)

        # -- Main loop timer (20 Hz) --
        self._timer = self.create_timer(0.05, self._control_loop)

        self.get_logger().info(
            f'Social nav started: mode={self._follow_mode.value}, '
            f'dist={self._follow_distance}m')

    # ----------------------------------------------------------------
    # Subscriptions
    # ----------------------------------------------------------------

    def _on_persons(self, msg):
        """Handle PersonStateArray from social perception."""
        target_id = msg.primary_attention_id
        if self._target_person_id >= 0:
            target_id = self._target_person_id

        for p in msg.persons:
            if p.person_id == target_id or (
                    target_id < 0 and len(msg.persons) > 0):
                pos = (p.position.x, p.position.y, p.position.z)
                self._update_person_position(pos, p.bearing_deg)
                break

    def _on_person_target(self, msg: PoseStamped):
        """Fallback: single person target pose."""
        pos = (msg.pose.position.x, msg.pose.position.y,
               msg.pose.position.z)
        # Estimate bearing from quaternion yaw
        q = msg.pose.orientation
        yaw = math.atan2(2.0 * (q.w * q.z + q.x * q.y),
                         1.0 - 2.0 * (q.y * q.y + q.z * q.z))
        self._update_person_position(pos, math.degrees(yaw))

    def _update_person_position(self, pos, bearing_deg):
        """Update person tracking state and record history."""
        now = time.time()
        self._target_person_pos = pos
        self._target_person_bearing = bearing_deg
        self._person_history.append((now, pos[0], pos[1]))

    def _on_odom(self, msg: Odometry):
        """Update robot pose from odometry."""
        self._robot_x = msg.pose.pose.position.x
        self._robot_y = msg.pose.pose.position.y
        q = msg.pose.pose.orientation
        self._robot_yaw = math.atan2(
            2.0 * (q.w * q.z + q.x * q.y),
            1.0 - 2.0 * (q.y * q.y + q.z * q.z))

    def _on_map(self, msg: OccupancyGrid):
        """Cache occupancy grid for A* planning."""
        w, h = msg.info.width, msg.info.height
        data = np.array(msg.data, dtype=np.int8).reshape((h, w))
        # Convert -1 (unknown) to free (0) for planning
        data[data < 0] = 0
        self._occ_grid = data.astype(np.int32)
        self._occ_origin = (msg.info.origin.position.x,
                            msg.info.origin.position.y)
        self._occ_resolution = msg.info.resolution

    def _on_target_id(self, msg: Int32):
        """Switch target person."""
        self._target_person_id = msg.data
        self.get_logger().info(f'Target person ID set to {msg.data}')

    def _on_voice_command(self, msg: String):
        """Handle discrete voice commands for mode switching."""
        cmd = msg.data.strip().lower()

        if cmd in ('stop', 'halt'):
            self._stopped = True
            self._replayer = None
            self._publish_status('STOPPED')
            return

        if cmd in ('resume', 'go', 'start'):
            self._stopped = False
            self._publish_status('RESUMED')
            return

        matched = MODE_VOICE_COMMANDS.get(cmd)
        if matched:
            self._follow_mode = matched
            self._stopped = False
            mode_msg = String()
            mode_msg.data = self._follow_mode.value
            self._mode_pub.publish(mode_msg)
            self._publish_status(f'MODE: {self._follow_mode.value}')

    def _on_transcript(self, msg: String):
        """Handle free-form voice transcripts for route teaching."""
        text = msg.data.strip().lower()

        # "teach route <name>"
        m = re.match(r'teach\s+route\s+(\w+)', text)
        if m:
            name = m.group(1)
            self._teaching = True
            self._current_route = WaypointRoute(name)
            self._last_teach_x = self._robot_x
            self._last_teach_y = self._robot_y
            self._publish_status(f'TEACHING: {name}')
            self.get_logger().info(f'Recording route: {name}')
            return

        # "stop teaching"
        if 'stop teaching' in text:
            if self._teaching and self._current_route:
                self._current_route.save(self._routes_dir)
                self._publish_status(
                    f'SAVED: {self._current_route.name} '
                    f'({len(self._current_route.waypoints)} pts)')
                self.get_logger().info(
                    f'Route saved: {self._current_route.name}')
            self._teaching = False
            self._current_route = None
            return

        # "replay route <name>"
        m = re.match(r'replay\s+route\s+(\w+)', text)
        if m:
            name = m.group(1)
            try:
                route = WaypointRoute.load(self._routes_dir, name)
                self._replayer = WaypointReplayer(route)
                self._stopped = False
                self._publish_status(f'REPLAY: {name}')
                self.get_logger().info(f'Replaying route: {name}')
            except FileNotFoundError:
                self._publish_status(f'ROUTE NOT FOUND: {name}')
            return

        # "go home"
        if 'go home' in text:
            self._go_home = True
            self._stopped = False
            self._publish_status('GO HOME')
            return

        # Also try mode commands from transcript
        for phrase, mode in MODE_VOICE_COMMANDS.items():
            if phrase in text:
                self._follow_mode = mode
                self._stopped = False
                mode_msg = String()
                mode_msg.data = self._follow_mode.value
                self._mode_pub.publish(mode_msg)
                self._publish_status(f'MODE: {self._follow_mode.value}')
                return

    # ----------------------------------------------------------------
    # Control loop
    # ----------------------------------------------------------------

    def _control_loop(self):
        """Main 20Hz control loop."""
        # Record waypoint if teaching
        if self._teaching and self._current_route:
            self._maybe_record_waypoint()

        # Publish zero velocity if stopped
        if self._stopped:
            self._publish_cmd_vel(0.0, 0.0)
            return

        # Determine navigation target
        target = self._get_nav_target()
        if target is None:
            self._publish_cmd_vel(0.0, 0.0)
            return

        tx, ty, tz = target

        # Publish debug target
        self._publish_target_pose(tx, ty, tz)

        # Check if arrived
        dist_to_target = math.hypot(tx - self._robot_x, ty - self._robot_y)
        if dist_to_target < self._goal_tol:
            self._publish_cmd_vel(0.0, 0.0)
            if self._go_home:
                self._go_home = False
                self._publish_status('HOME REACHED')
            return

        # Try A* path if map available
        if self._occ_grid is not None:
            path_target = self._plan_astar(tx, ty)
            if path_target:
                tx, ty = path_target

        # Determine speed limit
        max_lin = self._max_lin
        person_vel = self._estimate_person_velocity()
        if person_vel > _HIGH_SPEED_THRESHOLD:
            max_lin = self._max_lin_fast

        # Compute and publish cmd_vel
        lin, ang = self._compute_cmd_vel(
            self._robot_x, self._robot_y, self._robot_yaw,
            tx, ty, max_lin)
        self._publish_cmd_vel(lin, ang)

    def _get_nav_target(self):
        """Determine current navigation target based on mode/state."""
        # Route replay takes priority
        if self._replayer and not self._replayer.is_done:
            self._replayer.check_arrived(self._robot_x, self._robot_y)
            wp = self._replayer.current_waypoint()
            if wp:
                return (wp.x, wp.y, wp.z)
            else:
                self._replayer = None
                self._publish_status('REPLAY DONE')
                return None

        # Go home
        if self._go_home:
            return (self._home_x, self._home_y, 0.0)

        # Person following
        if self._target_person_pos is None:
            return None

        # Predict person position ahead for high-speed tracking
        px, py, pz = self._predict_person_position()
        bearing = self._target_person_bearing
        robot_pos = (self._robot_x, self._robot_y, 0.0)

        if self._follow_mode == FollowMode.SHADOW:
            return compute_shadow_target(
                (px, py, pz), bearing, self._follow_distance)
        elif self._follow_mode == FollowMode.LEAD:
            return compute_lead_target(
                (px, py, pz), bearing, self._lead_distance)
        elif self._follow_mode == FollowMode.SIDE:
            return compute_side_target(
                (px, py, pz), bearing, self._follow_distance)
        elif self._follow_mode == FollowMode.ORBIT:
            self._orbit_angle = (self._orbit_angle + 1.0) % 360.0
            return compute_orbit_target(
                (px, py, pz), self._orbit_angle, self._orbit_radius)
        elif self._follow_mode == FollowMode.LOOSE:
            return compute_loose_target(
                (px, py, pz), robot_pos, self._follow_distance)
        elif self._follow_mode == FollowMode.TIGHT:
            return compute_tight_target(
                (px, py, pz), self._follow_distance)

        return (px, py, pz)

    def _predict_person_position(self):
        """Extrapolate person position using velocity from recent samples."""
        if self._target_person_pos is None:
            return (0.0, 0.0, 0.0)

        px, py, pz = self._target_person_pos

        if len(self._person_history) >= 3:
            # Use last 3 samples for velocity estimation
            t0, x0, y0 = self._person_history[-3]
            t1, x1, y1 = self._person_history[-1]
            dt = t1 - t0
            if dt > 0.01:
                vx = (x1 - x0) / dt
                vy = (y1 - y0) / dt
                speed = math.hypot(vx, vy)
                if speed > _HIGH_SPEED_THRESHOLD:
                    px += vx * _PREDICT_AHEAD_S
                    py += vy * _PREDICT_AHEAD_S

        return (px, py, pz)

    def _estimate_person_velocity(self) -> float:
        """Estimate person speed from recent position history."""
        if len(self._person_history) < 2:
            return 0.0
        t0, x0, y0 = self._person_history[-2]
        t1, x1, y1 = self._person_history[-1]
        dt = t1 - t0
        if dt < 0.01:
            return 0.0
        return math.hypot(x1 - x0, y1 - y0) / dt

    def _plan_astar(self, target_x, target_y):
        """Run A* on occupancy grid, return next waypoint in world coords."""
        grid = self._occ_grid
        res = self._occ_resolution
        ox, oy = self._occ_origin

        # World to grid
        def w2g(wx, wy):
            return (int((wy - oy) / res), int((wx - ox) / res))

        start = w2g(self._robot_x, self._robot_y)
        goal = w2g(target_x, target_y)

        rows, cols = grid.shape
        if not (0 <= start[0] < rows and 0 <= start[1] < cols):
            return None
        if not (0 <= goal[0] < rows and 0 <= goal[1] < cols):
            return None

        inflated = inflate_obstacles(grid, self._inflation_cells)
        path = astar(inflated, start, goal)

        if path and len(path) > 1:
            # Follow a lookahead point (3 steps ahead or end)
            lookahead_idx = min(3, len(path) - 1)
            r, c = path[lookahead_idx]
            wx = ox + c * res + res / 2.0
            wy = oy + r * res + res / 2.0
            self._astar_path = path
            return (wx, wy)

        return None

    def _compute_cmd_vel(self, rx, ry, ryaw, tx, ty, max_lin):
        """Proportional controller: compute linear and angular velocity."""
        dx = tx - rx
        dy = ty - ry
        dist = math.hypot(dx, dy)
        angle_to_target = math.atan2(dy, dx)
        angle_error = angle_to_target - ryaw

        # Normalize angle error to [-pi, pi]
        while angle_error > math.pi:
            angle_error -= 2.0 * math.pi
        while angle_error < -math.pi:
            angle_error += 2.0 * math.pi

        angular_vel = _K_ANG * angle_error
        angular_vel = max(-self._max_ang,
                          min(self._max_ang, angular_vel))

        # Reduce linear speed when turning hard
        angle_factor = max(0.0, 1.0 - abs(angle_error) / (math.pi / 2.0))
        linear_vel = _K_LIN * dist * angle_factor
        linear_vel = max(0.0, min(max_lin, linear_vel))

        return (linear_vel, angular_vel)

    # ----------------------------------------------------------------
    # Waypoint teaching
    # ----------------------------------------------------------------

    def _maybe_record_waypoint(self):
        """Record waypoint if robot moved > _TEACH_MIN_DIST."""
        if self._last_teach_x is None:
            self._last_teach_x = self._robot_x
            self._last_teach_y = self._robot_y

        dist = math.hypot(
            self._robot_x - self._last_teach_x,
            self._robot_y - self._last_teach_y)

        if dist >= _TEACH_MIN_DIST:
            yaw_deg = math.degrees(self._robot_yaw)
            self._current_route.add(
                self._robot_x, self._robot_y, 0.0, yaw_deg)
            self._last_teach_x = self._robot_x
            self._last_teach_y = self._robot_y

    # ----------------------------------------------------------------
    # Publishers
    # ----------------------------------------------------------------

    def _publish_cmd_vel(self, linear: float, angular: float):
        twist = Twist()
        twist.linear.x = linear
        twist.angular.z = angular
        self._cmd_vel_pub.publish(twist)

    def _publish_target_pose(self, x, y, z):
        msg = PoseStamped()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = 'map'
        msg.pose.position.x = x
        msg.pose.position.y = y
        msg.pose.position.z = z
        self._target_pub.publish(msg)

    def _publish_status(self, status: str):
        msg = String()
        msg.data = status
        self._status_pub.publish(msg)
        self.get_logger().info(f'Nav status: {status}')


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


if __name__ == '__main__':
    main()
