import can
import sys

# Function to send the motor enable/disable command
def send_motor_enable(bus, enable):
    """
    Sends a command to enable or disable the motor.

    :param bus: The CAN bus
    :param enable: 1 to enable the motor, 0 to disable it
    """
    msg = can.Message()
    msg.arbitration_id = 1  # Message ID
    msg.is_extended_id = False
    msg.dlc = 2  # Message length (flag + 1 byte of data)
    msg.data = [ord('E'), enable]  # 'E' for the command, followed by 0 or 1

    try:
        bus.send(msg)
        state = "enabled" if enable else "disabled"
        print(f"Sent message to {state} motor")
        print(f"Message data: {msg.data}")
    except can.CanError as e:
        print(f"Message failed to send: {e}")
        sys.exit(1)  # Exit the program on failure

def main():
    # CAN interface setup
    bus = None  # Define outside the try block for proper shutdown
    try:
        bus = can.interface.Bus(channel='can0', bustype='socketcan', bitrate=1000000)  # Ensure the bitrate matches the microcontroller settings
        print("CAN bus initialized.")

        # Ensure the state is passed via arguments
        if len(sys.argv) != 2 or sys.argv[1] not in ['0', '1']:
            print("Usage: python3 script_name.py <0|1>")
            print("0 - Disable motor, 1 - Enable motor")
            sys.exit(1)

        enable = int(sys.argv[1])
        send_motor_enable(bus, enable)

    except Exception as e:
        print(f"Error initializing CAN bus: {e}")
        sys.exit(1)

    finally:
        # Ensure the bus is properly shut down
        if bus is not None:
            bus.shutdown()
            print("CAN bus shut down.")

if __name__ == '__main__':
    main()