main.py

# main.py
import sys
import bosdyn.client
from bosdyn.client.lease import LeaseClient, LeaseKeepAlive
from bosdyn.client.time_sync import TimeSyncClient, TimeSyncEndpoint
from PyQt6.QtCore import QThread, pyqtSignal, QPropertyAnimation, QRect, QEasingCurve, Qt
from PyQt6.QtGui import QFont
from PyQt6.QtWidgets import QApplication, QMainWindow, QPushButton, QVBoxLayout, QWidget, QLabel
from PyQt6.QtCore import QThread, pyqtSignal
from config import ROBOT_IP, ROBOT_USERNAME, ROBOT_PASSWORD
from robot_modules.docking import undock_robot, dock_robot
from robot_modules.movement import stand_robot, sit_robot, velocity_move, velocity_stop
from PyQt6.QtCore import QTimer
# --- Thread for Robot Connection ---
class RobotConnectThread(QThread):
    finished = pyqtSignal(str)   # success message
    failed = pyqtSignal(str)     # error message

    def __init__(self, hostname, username, password):
        super().__init__()
        self.hostname = hostname
        self.username = username
        self.password = password
        self.robot = None
        self.lease_keepalive = None

    def run(self):
        try:
            # Create SDK + robot and authenticate
            sdk = bosdyn.client.create_standard_sdk('RobotUI-App')
            self.robot = sdk.create_robot(self.hostname)
            self.robot.authenticate(self.username, self.password)

            # Time-sync (important for valid command timestamps)
            ts_client = self.robot.ensure_client(TimeSyncClient.default_service_name)
            TimeSyncEndpoint(ts_client).establish_timesync(max_samples=5)

            # Acquire lease and keep it alive so commands aren't rejected
            lease_client = self.robot.ensure_client(LeaseClient.default_service_name)
            self.lease_keepalive = LeaseKeepAlive(
                lease_client, must_acquire=True, return_at_exit=True
            )

            self.finished.emit("Connected & lease acquired.")
        except Exception as e:
            self.failed.emit(f"Failed to connect: {e}")

# --- Generic worker for commands so the UI stays responsive ---
class CommandThread(QThread):
    finished = pyqtSignal(str)  # success message
    failed = pyqtSignal(str)    # error message

    def __init__(self, fn, *args, **kwargs):
        super().__init__()
        self.fn = fn
        self.args = args
        self.kwargs = kwargs

    def run(self):
        try:
            self.fn(*self.args, **self.kwargs)
            self.finished.emit("OK")
        except Exception as e:
            self.failed.emit(str(e))

# --- Main Window ---
class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()

        # ----- Window -----
        self.setWindowTitle("Spot Controller")
        self.setGeometry(100, 100, 520, 360)

        central = QWidget()
        self.setCentralWidget(central)
        self.layout = QVBoxLayout(central)
        self.layout.setSpacing(18)
        self.layout.setContentsMargins(28, 28, 28, 28)
        self.layout.setAlignment(Qt.AlignmentFlag.AlignHCenter | Qt.AlignmentFlag.AlignTop)

        # ----- Title -----
        self.title = QLabel("Spot Controller")
        self.title.setFont(QFont("Arial", 22, QFont.Weight.Bold))
        self.title.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.title.setStyleSheet("color:#2b2b2b; letter-spacing:0.5px;")
        self.layout.addWidget(self.title)

        # ----- Status -----
        self.status_label = QLabel("Status: Disconnected")
        self.status_label.setFont(QFont("Arial", 14))
        self.status_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self._set_status("Disconnected", "neutral")
        self.layout.addWidget(self.status_label)

        # ----- Buttons -----
        self.connect_button = self._create_button("Connect", base="#0A84FF")   # blue
        self.undock_button  = self._create_button("Undock", base="#34C759")    # green
        self.dock_button    = self._create_button("Dock",   base="#FFCC00")    # yellow

        self.undock_button.setEnabled(False)
        self.dock_button.setEnabled(False)

        self.connect_button.clicked.connect(self.start_connection)
        self.undock_button.clicked.connect(self.undock_command)
        self.dock_button.clicked.connect(self.dock_command)

        self.layout.addWidget(self.connect_button)
        self.layout.addWidget(self.undock_button)
        self.layout.addWidget(self.dock_button)

                # Stand button
        self.stand_button = self._create_button("Stand", base="#4CAF50")
        self.stand_button.setEnabled(False)
        self.stand_button.clicked.connect(self.stand_command)
        self.layout.addWidget(self.stand_button)

        # Sit button
        self.sit_button = self._create_button("Sit", base="#E67E22")
        self.sit_button.setEnabled(False)
        self.sit_button.clicked.connect(self.sit_command)
        self.layout.addWidget(self.sit_button)


        # Window fade-in
        self._fade_in()

        self.connection_thread = None
        self.robot = None
        self._last_cmd_thread = None  # keep worker threads alive
        self._keys_held = set()
        self._vel_timer = QTimer(self)
        self._vel_timer.setInterval(100)  # 10 Hz
        self._vel_timer.timeout.connect(self._send_velocity)


    # ---------- UI helpers ----------
    def _create_button(self, text, base="#0A84FF"):
        btn = QPushButton(text)
        btn.setFont(QFont("Arial", 15, QFont.Weight.DemiBold))
        btn.setFixedSize(240, 56)
        btn.setCursor(Qt.CursorShape.PointingHandCursor)
        btn.setStyleSheet(f"""
            QPushButton {{
                background-color: {base};
                color: white;
                border: none;
                border-radius: 14px;
            }}
            QPushButton:hover {{
                filter: brightness(1.05);
            }}
            QPushButton:disabled {{
                background-color: #c8c8c8;
                color: #f0f0f0;
            }}
        """)
        btn.pressed.connect(lambda b=btn: self._press_anim(b))
        return btn

    def _press_anim(self, button):
        r = button.geometry()
        anim = QPropertyAnimation(button, b"geometry", self)
        anim.setDuration(120)
        anim.setEasingCurve(QEasingCurve.Type.InOutQuad)
        anim.setStartValue(r)
        anim.setEndValue(QRect(r.x()+5, r.y()+5, r.width()-10, r.height()-10))
        anim.finished.connect(lambda: button.setGeometry(r))  # restore
        anim.start()

    def _fade_in(self):
        self.setWindowOpacity(0.0)
        anim = QPropertyAnimation(self, b"windowOpacity", self)
        anim.setDuration(650)
        anim.setEasingCurve(QEasingCurve.Type.InOutCubic)
        anim.setStartValue(0.0)
        anim.setEndValue(1.0)
        anim.start()

    def _set_status(self, text, tone="neutral"):
        # tone: "ok" (green), "error" (red), "busy" (blue), "neutral" (gray)
        colors = {
            "ok": "#1f9d55",
            "error": "#d64545",
            "busy": "#0A84FF",
            "neutral": "#5a5a5a",
        }
        self.status_label.setText(f"Status: {text}")
        self.status_label.setStyleSheet(
            f"color:{colors.get(tone, '#5a5a5a')}; font-size:16px;"
        )

    # ---------- Connect flow ----------
    def start_connection(self):
        self._set_status("Connecting…", "busy")
        self.connect_button.setEnabled(False)

        self.connection_thread = RobotConnectThread(ROBOT_IP, ROBOT_USERNAME, ROBOT_PASSWORD)
        self.connection_thread.finished.connect(self.on_connected)
        self.connection_thread.failed.connect(self.on_connect_failed)
        self.connection_thread.start()

    def on_connected(self, message):
        self.robot = self.connection_thread.robot
        self._set_status(message, "ok")
        self.undock_button.setEnabled(True)
        self.dock_button.setEnabled(True)
        self.stand_button.setEnabled(True)  
        self.sit_button.setEnabled(True) 

    def on_connect_failed(self, message):
        self._set_status(message, "error")
        self.connect_button.setEnabled(True)

    # ---------- Undock ----------
    def undock_command(self):
        if not self.robot:
            self._set_status("Not connected to robot.", "error")
            return
        self._set_status("Undocking…", "busy")
        self.undock_button.setEnabled(False)

        th = CommandThread(undock_robot, self.robot)
        th.finished.connect(lambda _: self._undock_done(None))
        th.failed.connect(lambda err: self._undock_done(err))
        th.start()
        self._last_cmd_thread = th

    def _undock_done(self, err):
        if err:
            self._set_status(f"Undock error: {err}", "error")
        else:
            self._set_status("Undocked successfully.", "ok")
        self.undock_button.setEnabled(True)

    # ---------- Dock ----------
    def dock_command(self):
        if not self.robot:
            self._set_status("Not connected to robot.", "error")
            return
        self._set_status("Docking to 520…", "busy")
        self.dock_button.setEnabled(False)

        th = CommandThread(dock_robot, self.robot, 520)
        th.finished.connect(lambda _: self._dock_done(None))
        th.failed.connect(lambda err: self._dock_done(err))
        th.start()
        self._last_cmd_thread = th

    def _dock_done(self, err):
        if err:
            self._set_status(f"Dock error: {err}", "error")
        else:
            self._set_status("Docked successfully.", "ok")
        self.dock_button.setEnabled(True)

         # ---------- Stand ----------
    def stand_command(self):
        if not self.robot:
            self._set_status("Not connected to robot.", "error")
            return
        self._set_status("Standing…", "busy")

        th = CommandThread(stand_robot, self.robot)
        th.finished.connect(lambda _: self._set_status("Stand complete.", "ok"))
        th.failed.connect(lambda err: self._set_status(f"Stand error: {err}", "error"))
        th.start()
        self._last_cmd_thread = th

    # ---------- Sit ----------
    def sit_command(self):
        if not self.robot:
            self._set_status("Not connected to robot.", "error")
            return
        self._set_status("Sitting…", "busy")

        th = CommandThread(sit_robot, self.robot)
        th.finished.connect(lambda _: self._set_status("Sit complete.", "ok"))
        th.failed.connect(lambda err: self._set_status(f"Sit error: {err}", "error"))
        th.start()
        self._last_cmd_thread = th
       # ---------- Keyboard Controls ----------
    def keyPressEvent(self, event):
        if not self.robot:
            return
        key = event.key()
        if key in (Qt.Key.Key_W, Qt.Key.Key_S, Qt.Key.Key_A, Qt.Key.Key_D, Qt.Key.Key_Q, Qt.Key.Key_E):
            if not event.isAutoRepeat():
                self._keys_held.add(key)
                if not self._vel_timer.isActive():
                    self._vel_timer.start()

    def keyReleaseEvent(self, event):
        if not self.robot:
            return
        key = event.key()
        if key in (Qt.Key.Key_W, Qt.Key.Key_S, Qt.Key.Key_A, Qt.Key.Key_D, Qt.Key.Key_Q, Qt.Key.Key_E):
            if not event.isAutoRepeat():
                self._keys_held.discard(key)
                if not self._keys_held:
                    self._vel_timer.stop()
                    velocity_stop(self.robot)
    def _send_velocity(self):
        vx = vy = vrot = 0.0
        if Qt.Key.Key_W in self._keys_held: vx += 0.5     # forward
        if Qt.Key.Key_S in self._keys_held: vx -= 0.5     # backward
        if Qt.Key.Key_A in self._keys_held: vy += 0.4     # strafe left
        if Qt.Key.Key_D in self._keys_held: vy -= 0.4     # strafe right
        if Qt.Key.Key_Q in self._keys_held: vrot += 1.0   # rotate left
        if Qt.Key.Key_E in self._keys_held: vrot -= 1.0   # rotate right

        # Refresh short-duration command every tick for smooth motion
        velocity_move(self.robot, vx, vy, vrot, duration=0.15)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())