from pybricks.hubs import PrimeHub
from pybricks.pupdevices import Motor
from pybricks.parameters import Port, Stop
from pybricks.tools import wait, StopWatch, multitask
from umath import pi

class Attachment:
    def __init__(self, port, start_angle=0):
        self.motor = Motor(port)
        self.start_angle = start_angle
    
    def move_attachment(self, degrees, speed):
        self.motor.reset_angle(0)
        target_angle = degrees
        tolerance = 2
        movement_timer = StopWatch()
        movement_timeout = 3000  # 3 seconds timeout
        
        stuck_threshold = 50  # ms without significant movement
        last_angle = 0
        stuck_timer = StopWatch()
        
        while movement_timer.time() < movement_timeout:
            current_angle = self.motor.angle()
            error = target_angle - current_angle
            
            if abs(error) <= tolerance:
                self.motor.stop()
            
            # Check if motor is stuck
            if abs(current_angle - last_angle) < 1:  # Less than 1 degree movement
                if stuck_timer.time() > stuck_threshold:
                    print("Motor appears stuck")
                    self.motor.stop()
            else:
                stuck_timer.reset()
                last_angle = current_angle
            
            if error > 0:
                self.motor.run(speed)
            else:
                self.motor.run(-speed)
            wait(10)  # Consistent timing
        
        self.motor.stop()
        print("Movement timeout reached")

    def reset_attachment(self):
        self.motor.reset_angle(0)
        return self.move_attachment(self.start_angle, 100)

class Robot:
    def __init__(self, left_port, right_port, wheel_diameter):
        # Initialize hub and motors
        self.hub = PrimeHub()
        try:
            self.left_motor = Motor(left_port)
            self.right_motor = Motor(right_port)
        except Exception as e:
            print(f"Failed to initialize motors: {e}")
            
        # Reset and calibrate IMU
        try:
            self.hub.imu.reset_heading(0)
            wait(1000)  # Give IMU time to calibrate
        except Exception as e:
            print(f"Failed to initialize IMU: {e}")
            
        # Robot specs
        self.wheel_diameter = wheel_diameter
        self.wheel_circumference = pi * self.wheel_diameter
        
        # Attachments list
        self.attachments = {}
    
    def add_attachment(self, name, port, start_angle=0):
        try:
            self.attachments[name] = Attachment(port, start_angle)
        except Exception as e:
            print(f"Failed to add attachment '{name}': {e}")
    
    def move_attachment(self, attachment_name, degrees, speed):
        if attachment_name in self.attachments:
            return self.attachments[attachment_name].move_attachment(degrees, speed)
        else:
            print(f"Attachment '{attachment_name}' not found")
    
    def reset_attachment(self, attachment_name):
        if attachment_name in self.attachments:
            return self.attachments[attachment_name].reset_attachment()
        else:
            print(f"Attachment '{attachment_name}' not found")
    
    def straight(self, distance, speed):
        if not (-200 <= speed <= 200):
            print(f"Invalid speed: {speed}. Must be between -200 and 200.")
            
        target_heading = self.hub.imu.heading()
        # Reset distance tracking
        self.left_motor.reset_angle(0)
        self.right_motor.reset_angle(0)
        
        # Calculate target distance in motor degrees
        target_degrees = abs(distance) / self.wheel_circumference * 360
        
        while True:
            # Check current distance traveled (assuming both motors should be positive for forward)
            left_angle = abs(self.left_motor.angle())
            right_angle = abs(self.right_motor.angle())
            average_angle = (left_angle + right_angle) / 2
            
            # Stop if it reached the target
            if average_angle >= target_degrees:
                break
                
            # Get heading error for correction
            current_heading = self.hub.imu.heading()
            heading_error = target_heading - current_heading
            
            # Handle wraparound at 0°/360°
            if heading_error > 180:
                heading_error -= 360
            elif heading_error < -180:
                heading_error += 360
                
            # Calculate motor speeds with correction
            direction = 1 if distance > 0 else -1
            correction = heading_error * 1.5  # Reduced gain for stability
            left_speed = direction * (speed + correction)
            right_speed = direction * (speed - correction)
            
            # Limit speeds to prevent excessive values
            left_speed = max(-200, min(200, left_speed))
            right_speed = max(-200, min(200, right_speed))
            
            # Apply speeds to motors
            self.left_motor.run(left_speed)
            self.right_motor.run(right_speed)
            wait(10)  # Consistent timing
        
        # Stop motors
        self.left_motor.stop()
        self.right_motor.stop()
        wait(50)  # Allow motors to fully stop
    
    def turn(self, theta, speed):
        if not (-200 <= speed <= 200):
            print(f"Invalid speed: {speed}. Must be between -200 and 200.")
            
        start_angle = self.hub.imu.heading()
        target_angle = start_angle + theta
        
        # Normalize target angle to -180 to 180 range
        while target_angle > 180:
            target_angle -= 360
        while target_angle < -180:
            target_angle += 360
            
        while True:
            current_angle = self.hub.imu.heading()
            # Calculate angle error
            angle_error = target_angle - current_angle
            
            # Handle wraparound for shortest path
            if angle_error > 180:
                angle_error -= 360
            elif angle_error < -180:
                angle_error += 360
                
            # Stop if we're close enough
            if abs(angle_error) < 2:
                break     
				
            # Positive angle_error means we need to turn clockwise (right)
            if angle_error > 0:
                # Turn right: left motor forward, right motor backward
                self.left_motor.run(speed)
                self.right_motor.run(-speed)
            else:
                # Turn left: left motor backward, right motor forward
                self.left_motor.run(-speed)
                self.right_motor.run(speed)
            wait(10)  # Consistent timing
        
        # Stop both motors
        self.left_motor.stop()
        self.right_motor.stop()