#!/usr/env/bin python3
import os
import rospy
import rospkg
from PyQt5.QtWidgets import QApplication, QMainWindow, QVBoxLayout, QPushButton, QWidget, QLabel
from PyQt5.QtUiTools import QUiLoader
from rviz_visual_tools import RvizVisualTools

class MainWindow(QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()

        # Load the UI file
        ui_file = os.path.join(rospkg.RosPack().get_path('test_wst'), 'app.ui')
        self.load_ui(ui_file)

        # Connect buttons
        self.button1.clicked.connect(self.handle_button1)
        self.button2.clicked.connect(self.handle_button2)

        # Initialize RViz
        self.rviz = RvizVisualTools("world", "/rviz_visual_tools")

        # Add robot model to RViz
        self.rviz.load_robot()
        
        # Add RViz visualization to the UI
        self.rviz_widget = self.rviz.embed_into("rvizWidget")

    def load_ui(self, ui_file):
        loader = QUiLoader()
        ui_file = QFile(ui_file)
        ui_file.open(QFile.ReadOnly)
        self.window = loader.load(ui_file)
        ui_file.close()
        self.setCentralWidget(self.window.centralwidget)

        # Find UI elements
        self.button1 = self.window.centralwidget.findChild(QPushButton, "button1")
        self.button2 = self.window.centralwidget.findChild(QPushButton, "button2")
        self.rviz_widget = self.window.centralwidget.findChild(QWidget, "rvizWidget")

    def handle_button1(self):
        # Handle button 1 click
        print("Button 1 clicked")

    def handle_button2(self):
        # Handle button 2 click
        print("Button 2 clicked")

if __name__ == '__main__':
    # Initialize the ROS node
    rospy.init_node("rviz_example")

    # Initialize the Qt application
    app = QApplication([])
    window = MainWindow()
    window.show()

    # Start the Qt event loop
    app.exec_()