import tkinter as tk, robomasterpy as rm
from tkinter import ttk

class RobomasterUI:
    def __init__(self, master, cmd):
        self.master = master
        master.title("Robomaster Control")
        self.cmd = cmd
        self.output_text = tk.Text(master, height=10, width=50)
        self.output_text.grid(row=6, column=0, columnspan=4, padx=10, pady=10)

        arm_frame = ttk.LabelFrame(master, text="Arm Control")
        arm_frame.grid(row=0, column=0, padx=10, pady=10, sticky="nw")

        tk.Label(arm_frame, text="X:").grid(row=0, column=0, padx=5, pady=5, sticky="e")
        self.x_entry = tk.Entry(arm_frame, width=10)
        self.x_entry.grid(row=0, column=1, padx=5, pady=5)

        tk.Label(arm_frame, text="Y:").grid(row=1, column=0, padx=5, pady=5, sticky="e")
        self.y_entry = tk.Entry(arm_frame, width=10)
        self.y_entry.grid(row=1, column=1, padx=5, pady=5)

        tk.Button(arm_frame, text="Move", command=self.move_arm).grid(row=0, column=2, padx=5, pady=5)
        tk.Button(arm_frame, text="Move To", command=self.moveto_arm).grid(row=1, column=2, padx=5, pady=5)
        tk.Button(arm_frame, text="Recenter", command=self.recenter_arm).grid(row=2, column=0, columnspan=3, pady=5, sticky="we")
        tk.Button(arm_frame, text="Stop", command=self.stop_arm).grid(row=3, column=0, columnspan=3, pady=5, sticky="we")
        tk.Button(arm_frame, text="Get Position", command=self.get_arm_pos).grid(row=4, column=0, columnspan=3, pady=5, sticky="we")

        gripper_frame = ttk.LabelFrame(master, text="Gripper Control")
        gripper_frame.grid(row=0, column=1, padx=10, pady=10, sticky="nw")

        tk.Label(gripper_frame, text="Force (1-4, optional):").grid(row=0, column=0, padx=5, pady=5, sticky="e")
        self.force_entry = tk.Entry(gripper_frame, width=5)
        self.force_entry.grid(row=0, column=1, padx=5, pady=5)

        tk.Button(gripper_frame, text="Open Gripper", command=self.open_gripper).grid(row=1, column=0, columnspan=2, pady=5, sticky="we")
        tk.Button(gripper_frame, text="Close Gripper", command=self.close_gripper).grid(row=2, column=0, columnspan=2, pady=5, sticky="we")
        tk.Button(gripper_frame, text="Gripper Status", command=self.get_gripper_status).grid(row=3, column=0, columnspan=2, pady=5, sticky="we")

        chassis_frame = ttk.LabelFrame(master, text="Chassis Control")
        chassis_frame.grid(row=1, column=0, columnspan=2, padx=10, pady=10, sticky="nw")

        tk.Label(chassis_frame, text="X Speed (m/s):").grid(row=0, column=0, padx=5, pady=5, sticky="e")
        self.chassis_x_speed_entry = tk.Entry(chassis_frame, width=5)
        self.chassis_x_speed_entry.insert(0, "0.1")
        self.chassis_x_speed_entry.grid(row=0, column=1, padx=5, pady=5)

        tk.Label(chassis_frame, text="Y Speed (m/s):").grid(row=0, column=2, padx=5, pady=5, sticky="e")
        self.chassis_y_speed_entry = tk.Entry(chassis_frame, width=5)
        self.chassis_y_speed_entry.insert(0, "0.1")
        self.chassis_y_speed_entry.grid(row=0, column=3, padx=5, pady=5)

        tk.Label(chassis_frame, text="Z Speed (rad/s):").grid(row=0, column=4, padx=5, pady=5, sticky="e")
        self.chassis_z_speed_entry = tk.Entry(chassis_frame, width=5)
        self.chassis_z_speed_entry.insert(0, "0.5")
        self.chassis_z_speed_entry.grid(row=0, column=5, padx=5, pady=5)

        tk.Button(chassis_frame, text="Move Chassis", command=self.move_chassis_button).grid(row=1, column=0, columnspan=6, pady=5, sticky="we")
        tk.Button(chassis_frame, text="Stop Chassis", command=self.stop_chassis).grid(row=2, column=0, columnspan=6, pady=5, sticky="we")

        master.bind("<Up>", lambda event: self.chassis_move_arrow(y=1))
        master.bind("<Down>", lambda event: self.chassis_move_arrow(y=-1))
        master.bind("<Left>", lambda event: self.chassis_move_arrow(x=-1))
        master.bind("<Right>", lambda event: self.chassis_move_arrow(x=1))
        master.bind("<KeyPress-q>", lambda event: self.chassis_move_arrow(z=1))
        master.bind("<KeyPress-e>", lambda event: self.chassis_move_arrow(z=-1))
        master.bind("<KeyRelease-Up>", lambda event: self.chassis_move_arrow(y=0) if event.keysym == 'Up' else None)
        master.bind("<KeyRelease-Down>", lambda event: self.chassis_move_arrow(y=0) if event.keysym == 'Down' else None)
        master.bind("<KeyRelease-Left>", lambda event: self.chassis_move_arrow(x=0) if event.keysym == 'Left' else None)
        master.bind("<KeyRelease-Right>", lambda event: self.chassis_move_arrow(x=0) if event.keysym == 'Right' else None)
        master.bind("<KeyRelease-q>", lambda event: self.chassis_move_arrow(z=0) if event.keysym == 'q' else None)
        master.bind("<KeyRelease-e>", lambda event: self.chassis_move_arrow(z=0) if event.keysym == 'e' else None)
        master.bind("<space>", lambda event: self.stop_chassis())
        master.bind("<Button-1>", self.deselect_textbox)

    def deselect_textbox(self, event):
        focused_widget = self.master.focus_get()
        if isinstance(focused_widget, tk.Entry) and event.widget != focused_widget:
            self.master.focus_set()

    def _command_executor(self, func, *args, **kwargs):
        try:
            result = func(*args, **kwargs)
            self.master.after(0, self._display_output, str(result))
        except Exception as e:
            self.master.after(0, self._display_output, f"Error: {e}")

    def _display_output(self, output_str):
        self.output_text.insert(tk.END, output_str + "\n")
        self.output_text.see(tk.END)

    def move_arm(self):
        x_str, y_str = self.x_entry.get(), self.y_entry.get()
        try: x, y = float(x_str), float(y_str)
        except ValueError: return self._display_output("Error: Invalid X or Y")
        self._command_executor(self.cmd.arm_move, x=x, y=y)

    def moveto_arm(self):
        x_str, y_str = self.x_entry.get(), self.y_entry.get()
        try: x, y = float(x_str), float(y_str)
        except ValueError: return self._display_output("Error: Invalid X or Y")
        self._command_executor(self.cmd.arm_moveto, x=x, y=y)

    def recenter_arm(self):
        self._command_executor(self.cmd.arm_recenter)

    def stop_arm(self):
        self._command_executor(self.cmd.arm_stop)

    def get_arm_pos(self):
        self._command_executor(self.cmd.arm_pos)

    def open_gripper(self):
        force_str = self.force_entry.get()
        force = None
        if force_str:
            try:
                force = int(force_str)
                if force not in range(1, 5): raise ValueError
            except ValueError: return self._display_output("Error: Force must be 1-4")
        self._command_executor(self.cmd.gripper_open, force=force)

    def close_gripper(self):
        force_str = self.force_entry.get()
        force = None
        if force_str:
            try:
                force = int(force_str)
                if force not in range(1, 5): raise ValueError
            except ValueError: return self._display_output("Error: Force must be 1-4")
        self._command_executor(self.cmd.gripper_close, force=force)

    def get_gripper_status(self):
        self._command_executor(self.cmd.gripper_status)

    def move_chassis_button(self):
        try:
            x = float(self.chassis_x_speed_entry.get())
            y = float(self.chassis_y_speed_entry.get())
            z = float(self.chassis_z_speed_entry.get())
        except ValueError:
            return self._display_output("Error: Invalid speed value")
        self._command_executor(self.cmd.chassis_move, x=x, y=y, z=z)

    def stop_chassis(self):
        self._command_executor(self.cmd.chassis_move, x=0, y=0, z=0)

    def chassis_move_arrow(self, x=None, y=None, z=None):
        try:
            xs = float(self.chassis_x_speed_entry.get()) if x is not None else 0
            ys = float(self.chassis_y_speed_entry.get()) if y is not None else 0
            zs = float(self.chassis_z_speed_entry.get()) if z is not None else 0
            current_x = xs * (x or 0)
            current_y = ys * (y or 0)
            current_z = zs * (z or 0)
            self._command_executor(self.cmd.chassis_move, x=current_x, y=current_y, z=current_z, speed_xy=None, speed_z=None)
        except ValueError:
            self._display_output("Error: Invalid speed value")

def main():
    ip_address = "192.168.2.1"
    try:
        cmd = rm.Commander(ip=ip_address)
        print(f"Connected to Robomaster. Version: {cmd.version()}")
        root = tk.Tk()
        RobomasterUI(root, cmd)
        root.mainloop()
    except Exception as e:
        print(f"Error connecting to Robomaster: {e}")
        print("Make sure the Robomaster is powered on and connected.")

if __name__ == "__main__":
    main()