from pybricks.hubs import PrimeHub
from pybricks.pupdevices import Motor, ColorSensor, UltrasonicSensor, ForceSensor
from pybricks.parameters import Button, Color, Direction, Port, Side, Stop
from pybricks.robotics import DriveBase
from pybricks.tools import wait, StopWatch
HUB = PrimeHub()
from BatteryDiagnostics import BatteryDiagnostics
from MotorDiagnostics import MotorDiagnostics
from ColorSensorDiagnostics import ColorSensorDiagnostics
from DriveBaseDiagnostics import DriveBaseDiagnostics
from HubDiagnostics import HubDiagnostics
battery = BatteryDiagnostics(HUB)
motor = MotorDiagnostics(HUB, Motor)
colorsensor = ColorSensorDiagnostics(HUB, ColorSensor)
hubdiags = HubDiagnostics(HUB)
drivebase = DriveBaseDiagnostics(HUB, Motor, DriveBase)
CLEARCONFIRM = input("Clear the console before proceeding? Y/N (default: yes): ")
if(CLEARCONFIRM == "Y" or CLEARCONFIRM == "y" or CLEARCONFIRM == "yes" or CLEARCONFIRM == ""):
    print("Clearing console... \n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n")
print("""
 ███████████  █████ █████ ██████   █████   █████████   ██████   ██████ █████   █████████   █████████ 
▒▒███▒▒▒▒▒███▒▒███ ▒▒███ ▒▒██████ ▒▒███   ███▒▒▒▒▒███ ▒▒██████ ██████ ▒▒███   ███▒▒▒▒▒███ ███▒▒▒▒▒███
 ▒███    ▒███ ▒▒███ ███   ▒███▒███ ▒███  ▒███    ▒███  ▒███▒█████▒███  ▒███  ███     ▒▒▒ ▒███    ▒▒▒ 
 ▒██████████   ▒▒█████    ▒███▒▒███▒███  ▒███████████  ▒███▒▒███ ▒███  ▒███ ▒███         ▒▒█████████ 
 ▒███▒▒▒▒▒▒     ▒▒███     ▒███ ▒▒██████  ▒███▒▒▒▒▒███  ▒███ ▒▒▒  ▒███  ▒███ ▒███          ▒▒▒▒▒▒▒▒███
 ▒███            ▒███     ▒███  ▒▒█████  ▒███    ▒███  ▒███      ▒███  ▒███ ▒▒███     ███ ███    ▒███
 █████           █████    █████  ▒▒█████ █████   █████ █████     █████ █████ ▒▒█████████ ▒▒█████████ 
▒▒▒▒▒           ▒▒▒▒▒    ▒▒▒▒▒    ▒▒▒▒▒ ▒▒▒▒▒   ▒▒▒▒▒ ▒▒▒▒▒     ▒▒▒▒▒ ▒▒▒▒▒   ▒▒▒▒▒▒▒▒▒   ▒▒▒▒▒▒▒▒▒  
                                                                                                     
The free and open source diagnostics tool for the LEGO® Education SPIKE™ Prime robots, designed for FIRST Lego League.
Developed by Team 65266, Lego Dynamics.  
Please set your window size to 90% on small screens for best results with the ASCII art.               
          """
          )
while True:
    
    print("\nWhich diagnostic do you want to perform?")
    print("Enter 'b' for battery diagnostics")
    print("Enter 'm' for motor diagnostics")
    print("Enter 'cs' for color sensor diagnostics")
    print("Enter 'h' for hub diagnostics")
    print("Enter 'db' for drive base diagnostics")
    print("Enter 'q' to quit")

    choice = input("Your choice: ").strip().lower()

    if choice == "b":
        print("-----------------------BATTERY DIAGNOSTICS-----------------------")
        print("This test will check the battery voltage and current. It will measure these over a period of 3 seconds and provide average and deviation values. Your voltage should be above 7800 mV for optimal performance.")
        input("Press Enter to begin the battery diagnostics.")
        battery.printAll()
        print("Battery diagnostics completed.")

    elif choice == "m":
        print("------------------------MOTOR DIAGNOSTICS------------------------")
        motor.fullTest()
        print("Motor diagnostics completed.")
    elif choice == "h":
        print("-------------------------HUB DIAGNOSTICS-------------------------")
        hubdiags.printAll()
        print("Hub diagnostics completed.")
    elif choice == "q":
        print("Diagnostics completed successfully. Exiting with code 0. Good luck in the robot game!")
        break
    elif choice == "db":
        print("----------------------DRIVEBASE DIAGNOSTICS----------------------")
        drivebase.printAll()
    elif choice == "cs":
        print("---------------------COLOR SENSOR DIAGNOSTICS---------------------")
        colorsensor.printAll()
        print("Color sensor diagnostics completed.")
        
    else:
        print("Invalid choice. Please enter 'b', 'm', or 'q'.")