#!/usr/bin/env python3
"""Waypoint Follower - Follows relative waypoints without GPS."""

import rclpy
from rclpy.node import Node
from geometry_msgs.msg import PoseStamped, Twist
from nav_msgs.msg import Odometry, Path
from std_msgs.msg import Int32, String, Bool
import numpy as np
from enum import Enum


class WaypointState(Enum):
    IDLE = 0
    NAVIGATING = 1
    REACHED = 2
    PAUSED = 3


class WaypointFollower(Node):
    
    def __init__(self):
        super().__init__('waypoint_follower')
        
        self.declare_parameter('waypoint_radius', 0.5)
        self.declare_parameter('max_velocity', 2.0)
        self.declare_parameter('kp_linear', 0.8)
        self.declare_parameter('kp_angular', 1.0)
        
        self.waypoint_radius = self.get_parameter('waypoint_radius').value
        self.max_velocity = self.get_parameter('max_velocity').value
        self.kp_linear = self.get_parameter('kp_linear').value
        self.kp_angular = self.get_parameter('kp_angular').value
        
        self.waypoints = []
        self.current_idx = 0
        self.state = WaypointState.IDLE
        self.current_pose = None
        
        self.odom_sub = self.create_subscription(
            Odometry, '/uav/local_position/odom', self.odom_callback, 10)
        
        self.waypoint_sub = self.create_subscription(
            PoseStamped, '/waypoint_follower/add_waypoint', self.add_waypoint_callback, 10)
        
        self.path_sub = self.create_subscription(
            Path, '/waypoint_follower/set_path', self.set_path_callback, 10)
        
        self.command_sub = self.create_subscription(
            String, '/waypoint_follower/command', self.command_callback, 10)
        
        self.cmd_vel_pub = self.create_publisher(Twist, '/uav/cmd_vel', 10)
        self.current_waypoint_pub = self.create_publisher(Int32, '/waypoint_follower/current_index', 10)
        self.status_pub = self.create_publisher(String, '/waypoint_follower/status', 10)
        self.reached_pub = self.create_publisher(Bool, '/waypoint_follower/waypoint_reached', 10)
        
        self.timer = self.create_timer(0.05, self.control_loop)
        
        self.get_logger().info('Waypoint Follower Started - Using LOCAL coordinates only')
    
    def odom_callback(self, msg):
        self.current_pose = msg.pose.pose
    
    def add_waypoint_callback(self, msg):
        waypoint = np.array([
            msg.pose.position.x,
            msg.pose.position.y,
            msg.pose.position.z
        ])
        self.waypoints.append(waypoint)
        self.get_logger().info(f'Added waypoint {len(self.waypoints)}: {waypoint}')
    
    def set_path_callback(self, msg):
        self.waypoints = []
        for pose_stamped in msg.poses:
            waypoint = np.array([
                pose_stamped.pose.position.x,
                pose_stamped.pose.position.y,
                pose_stamped.pose.position.z
            ])
            self.waypoints.append(waypoint)
        self.current_idx = 0
        self.get_logger().info(f'Set path with {len(self.waypoints)} waypoints')
    
    def command_callback(self, msg):
        cmd = msg.data.lower()
        
        if cmd == 'start':
            if len(self.waypoints) > 0:
                self.state = WaypointState.NAVIGATING
                self.get_logger().info('Starting waypoint navigation')
        elif cmd == 'pause':
            self.state = WaypointState.PAUSED
            self.stop()
        elif cmd == 'resume':
            self.state = WaypointState.NAVIGATING
        elif cmd == 'stop':
            self.state = WaypointState.IDLE
            self.stop()
        elif cmd == 'clear':
            self.waypoints = []
            self.current_idx = 0
            self.state = WaypointState.IDLE
        elif cmd == 'next':
            if self.current_idx < len(self.waypoints) - 1:
                self.current_idx += 1
        elif cmd == 'prev':
            if self.current_idx > 0:
                self.current_idx -= 1
    
    def control_loop(self):
        self.publish_status()
        
        if self.state != WaypointState.NAVIGATING:
            return
        
        if self.current_pose is None or len(self.waypoints) == 0:
            return
        
        if self.current_idx >= len(self.waypoints):
            self.state = WaypointState.IDLE
            self.stop()
            self.get_logger().info('All waypoints reached!')
            return
        
        current_pos = np.array([
            self.current_pose.position.x,
            self.current_pose.position.y,
            self.current_pose.position.z
        ])
        
        target = self.waypoints[self.current_idx]
        error = target - current_pos
        distance = np.linalg.norm(error)
        
        if distance < self.waypoint_radius:
            reached_msg = Bool()
            reached_msg.data = True
            self.reached_pub.publish(reached_msg)
            
            self.get_logger().info(f'Reached waypoint {self.current_idx + 1}/{len(self.waypoints)}')
            self.current_idx += 1
            return
        
        direction = error / distance
        speed = min(self.kp_linear * distance, self.max_velocity)
        velocity = direction * speed
        
        cmd = Twist()
        cmd.linear.x = float(velocity[0])
        cmd.linear.y = float(velocity[1])
        cmd.linear.z = float(velocity[2])
        
        self.cmd_vel_pub.publish(cmd)
    
    def stop(self):
        cmd = Twist()
        self.cmd_vel_pub.publish(cmd)
    
    def publish_status(self):
        idx_msg = Int32()
        idx_msg.data = self.current_idx
        self.current_waypoint_pub.publish(idx_msg)
        
        status_msg = String()
        status_msg.data = f'{self.state.name}|{self.current_idx}/{len(self.waypoints)}'
        self.status_pub.publish(status_msg)


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


if __name__ == '__main__':
    main()