# utils/robot_control.py - Shared driving and control functions
from pybricks.robotics import DriveBase
from config import ROBOT_CONFIG
import time

# Initialize drive base (done once)
drive_base = DriveBase(
    ROBOT_CONFIG['left_motor'], 
    ROBOT_CONFIG['right_motor'], 
    wheel_diameter=56, 
    axle_track=114
)

def drive_straight(distance, speed=200):
    """Drive straight for a given distance in mm"""
    drive_base.straight(distance, speed=speed)

def turn_angle(angle, speed=100):
    """Turn by a given angle in degrees"""
    drive_base.turn(angle, speed=speed)

def drive_until_line(speed=100, sensor='color_left'):
    """Drive until line is detected"""
    sensor_obj = ROBOT_CONFIG[sensor]
    drive_base.drive(speed, 0)
    
    while sensor_obj.color() != Color.BLACK:
        time.sleep(0.01)
    
    drive_base.stop()

def return_to_base():
    """Return robot to launch area - implement based on your strategy"""
    print("Returning to base...")
    # Add your return-to-base logic here
    pass

def reset_robot():
    """Reset all motors and sensors to default state"""
    drive_base.stop()
    for motor in ROBOT_CONFIG.get('attachment_motors', []):
        motor.stop()
    print("Robot reset completed")