#!/usr/bin/env python3

import time
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist
from turtlesim.srv import Spawn, Kill

class MultiTurtleController(Node):
    def __init__(self):
        super().__init__('multi_turtle_controller')
        
        self.spawn_client = self.create_client(Spawn, '/spawn')
        self.kill_client = self.create_client(Kill, '/kill')
        
        # Wait for service availability
        while not self.spawn_client.wait_for_service(timeout_sec=1.0):
            self.get_logger().info('Waiting for /spawn service...')

    def kill_default_turtle(self):
        """Removes default 'turtle1' to prepare a clean canvas."""
        if self.kill_client.wait_for_service(timeout_sec=1.0):
            req = Kill.Request()
            req.name = 'turtle1'
            future = self.kill_client.call_async(req)
            rclpy.spin_until_future_complete(self, future)

    def spawn_turtle(self, name, x, y, theta):
        """Spawns a turtle at specific (x, y, theta) coordinates."""
        req = Spawn.Request()
        req.x = float(x)
        req.y = float(y)
        req.theta = float(theta)
        req.name = name
        future = self.spawn_client.call_async(req)
        rclpy.spin_until_future_complete(self, future)
        self.get_logger().info(f'Spawned {name} at ({x}, {y})')

    def publish_velocities(self, publishers, linear_x, angular_z, duration):
        """
        Publishes velocity commands to all registered publishers simultaneously
        while spinning the node to process incoming ROS 2 events.
        """
        msg = Twist()
        msg.linear.x = float(linear_x)
        msg.angular.z = float(angular_z)
        
        start_time = time.time()
        while (time.time() - start_time) < duration and rclpy.ok():
            for pub in publishers:
                pub.publish(msg)
            rclpy.spin_once(self, timeout_sec=0.05)
            
        # Send zero-velocity stop pulse to halt movement smoothly
        stop_msg = Twist()
        for pub in publishers:
            pub.publish(stop_msg)
        rclpy.spin_once(self, timeout_sec=0.05)


def run_square_formation(node):
    """
    Spawns 4 turtles in a square pattern and drives them concurrently along a square path.
    """
    node.get_logger().info("=== STARTING 4-TURTLE SQUARE FORMATION ===")
    
    turtle_configs = [
        ('sq_turtle1', 3.0, 3.0, 0.0),
        ('sq_turtle2', 7.0, 3.0, 0.0),
        ('sq_turtle3', 7.0, 7.0, 0.0),
        ('sq_turtle4', 3.0, 7.0, 0.0),
    ]
    
    publishers = []
    for name, x, y, theta in turtle_configs:
        node.spawn_turtle(name, x, y, theta)
        pub = node.create_publisher(Twist, f'/{name}/cmd_vel', 10)
        publishers.append(pub)
        
    # Allow publisher topics to establish connections across the ROS graph
    time_start = time.time()
    while time.time() - time_start < 1.0:
        rclpy.spin_once(node, timeout_sec=0.1)

    # Drive 4 sides of the square
    side_length_duration = 2.0  # Forward movement duration (seconds)
    turn_90_deg_duration = 1.57 # ~1.57 rad/s angular speed for 1 sec = 90-degree turn
    
    for side in range(4):
        node.get_logger().info(f"Square - Side {side + 1}: Moving Forward")
        node.publish_velocities(publishers, linear_x=1.5, angular_z=0.0, duration=side_length_duration)
        
        node.get_logger().info(f"Square - Corner {side + 1}: Turning Right")
        node.publish_velocities(publishers, linear_x=0.0, angular_z=-1.57, duration=1.0)


def run_zigzag_formation(node):
    """
    Spawns 3 turtles in a column and drives them in a synchronized zigzag path.
    """
    node.get_logger().info("=== STARTING 3-TURTLE ZIGZAG FORMATION ===")
    
    turtle_configs = [
        ('zig_turtle1', 1.5, 2.5, 0.0),
        ('zig_turtle2', 1.5, 5.5, 0.0),
        ('zig_turtle3', 1.5, 8.5, 0.0),
    ]
    
    publishers = []
    for name, x, y, theta in turtle_configs:
        node.spawn_turtle(name, x, y, theta)
        pub = node.create_publisher(Twist, f'/{name}/cmd_vel', 10)
        publishers.append(pub)
        
    time_start = time.time()
    while time.time() - time_start < 1.0:
        rclpy.spin_once(node, timeout_sec=0.1)

    num_legs = 3
    
    for leg in range(num_legs):
        # Phase A: Turn Up-Right (Zig) & Drive Forward
        node.get_logger().info(f"Zigzag Leg {leg + 1}a: Steering Left (Zig)")
        node.publish_velocities(publishers, linear_x=0.0, angular_z=0.785, duration=1.0)
        node.publish_velocities(publishers, linear_x=1.5, angular_z=0.0, duration=1.5)
        
        # Phase B: Turn Down-Right (Zag) & Drive Forward
        node.get_logger().info(f"Zigzag Leg {leg + 1}b: Steering Right (Zag)")
        node.publish_velocities(publishers, linear_x=0.0, angular_z=-1.57, duration=1.0)
        node.publish_velocities(publishers, linear_x=1.5, angular_z=0.0, duration=1.5)
        
        # Re-orient forward to align for next zigzag iteration
        node.publish_velocities(publishers, linear_x=0.0, angular_z=0.785, duration=1.0)


def main():
    rclpy.init()
    node = MultiTurtleController()
    
    # Remove default turtle1
    node.kill_default_turtle()
    
    try:
        # Phase 1: 4 Turtles - Square Formation
        run_square_formation(node)
        
        # Delay between formations
        time_start = time.time()
        while time.time() - time_start < 2.0:
            rclpy.spin_once(node, timeout_sec=0.1)
        
        # Phase 2: 3 Turtles - Zigzag Formation
        run_zigzag_formation(node)
        
        node.get_logger().info("=== ALL FORMATION MOVEMENTS COMPLETED SUCCESSFULLY ===")
    except KeyboardInterrupt:
        node.get_logger().info("Simulation stopped by user.")
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()