from pybricks.parameters import Direction, Port, Side, Stop
from pybricks.robotics import DriveBase
from pybricks.tools import wait, StopWatch

from usys import stdin
from uselect import poll

class DriveBaseDiagnostics:
    def __init__(self, hub, motorclass, dbclass):
        self.hub = hub
        self.motorclass = motorclass
        self.dbclass = dbclass
        self.drive_base = None
        self.PORT_MAP = {
            "A": Port.A,
            "B": Port.B,
            "C": Port.C,
            "D": Port.D,
            "E": Port.E,
            "F": Port.F,
            }
    def initializeDriveBase(self):

        print("DriveBase setup:")
        print("1 = Load from savefile (paste JSON)")
        print("2 = Use defaults")
        print("3 = Enter values manually")

        choice = input("Choose an option: ")

        # Default values
        WHEEL_DIAMETER = 68.8
        AXLE_TRACK = 180
        DRIVE_SETTINGS = (600, 2000, 300, 2000)

        # Motor ports (None until set)
        leftmotorport = Port.A
        rightmotorport = Port.B

        # -----------------------------
        # OPTION 1: LOAD SAVEFILE
        # -----------------------------
        if choice == "1":
            print("Paste JSON:")
            raw = input("> ")

            # --- wheel ---
            if "\"wheel\"" in raw:
                part = raw.split("\"wheel\"")[1]
                part = part.split(":")[1]
                part = part.split(",")[0]
                WHEEL_DIAMETER = float(part)

            # --- axle ---
            if "\"axle\"" in raw:
                part = raw.split("\"axle\"")[1]
                part = part.split(":")[1]
                part = part.split(",")[0]
                AXLE_TRACK = float(part)

            # --- settings ---
            if "\"settings\"" in raw:
                part = raw.split("\"settings\"")[1]
                part = part.split("[")[1]
                part = part.split("]")[0]
                nums = part.split(",")
                DRIVE_SETTINGS = (int(nums[0]), int(nums[1]), int(nums[2]), int(nums[3]))

                            # --- left motor port ---
                if "\"left_port\"" in raw:
                    part = raw.split("\"left_port\"")[1]
                    part = part.split("\"")[1]   # first quoted value
                    leftmotorport = part

                # --- right motor port ---
                if "\"right_port\"" in raw:
                    part = raw.split("\"right_port\"")[1]
                    part = part.split("\"")[1]
                    rightmotorport = part


            print("Loaded config.")

        # -----------------------------
        # OPTION 3: MANUAL ENTRY
        # -----------------------------
        elif choice == "3":
            WHEEL_DIAMETER = float(input("Wheel diameter: "))
            AXLE_TRACK = float(input("Axle track: "))

            print("Enter drive settings:")
            s1 = int(input("Straight speed: "))
            s2 = int(input("Straight accel: "))
            s3 = int(input("Turn rate: "))
            s4 = int(input("Turn accel: "))
            DRIVE_SETTINGS = (s1, s2, s3, s4)

            # Ask for motor ports HERE (manual only)
            leftmotorport = input("Left motor port: ")
            rightmotorport = input("Right motor port: ")

        # -----------------------------
        # OPTION 2: DEFAULTS
        # -----------------------------
        else:
            pass

        # -----------------------------
        # CREATE MOTORS
        # -----------------------------
        left_motor = self.motorclass(leftmotorport, Direction.COUNTERCLOCKWISE)
        right_motor = self.motorclass(rightmotorport, Direction.CLOCKWISE)

        # -----------------------------
        # CREATE DRIVEBASE
        # -----------------------------
        self.drive_base = self.dbclass(left_motor, right_motor, WHEEL_DIAMETER, AXLE_TRACK)
        self.drive_base.settings(*DRIVE_SETTINGS)
        self.drive_base.use_gyro(True)

        print("DriveBase initialized.")
        return DRIVE_SETTINGS


    
    def printAll(self):
        self.driveRobot()
    def driveRobot(self):
        drivesettings = self.initializeDriveBase()
        print(drivesettings)
        keyboard = poll()
        keyboard.register(stdin)

        last_key_time = StopWatch()
        last_key_time.reset()

        while True:
            key = None

            # Check for keypress immediately
            if keyboard.poll(0):
                key = stdin.read(1).upper()
                last_key_time.reset()

            # Process key
            if key:
                if key == "W":
                    self.drive_base.drive(drivesettings[0], 0)
                elif key == "A":
                    self.drive_base.drive(drivesettings[0], -180)
                elif key == "S":
                    self.drive_base.drive(-drivesettings[0], 0)
                elif key == "D":
                    self.drive_base.drive(drivesettings[0], 180)
                elif key == "X":
                    break

            # Auto-stop if no key for 120 ms
            if last_key_time.time() > 120:
                self.drive_base.stop()

            # Tiny manual delay to avoid 100% CPU
            # (much faster than wait())
            for _ in range(200):
                pass