from PySide6.QtWidgets import QDockWidget, QFormLayout, QHBoxLayout, QWidget, QSlider, QPushButton, QApplication, \
    QMainWindow

from PySide6.QtGui import QColor, QVector3D, QQuaternion
from PySide6.Qt3DExtras import Qt3DExtras
from PySide6.Qt3DRender import Qt3DRender
from PySide6.Qt3DCore import Qt3DCore
from PySide6.QtCore import Qt, QUrl


class Mesh3DViewer(QDockWidget):

    def __init__(self, mesh):
        super().__init__()
        self.mesh_object = None

        # Structure
        self.sliders_body = QFormLayout()

        self.body = QHBoxLayout()
        self.container = QWidget()

        # Components
        self.scene = Qt3DCore.QEntity()

        self.view = Qt3DExtras.Qt3DWindow()
        self.view.defaultFrameGraph().setClearColor(QColor(Qt.gray))

        self.view.camera().lens().setPerspectiveProjection(45.0, 16.0/9.0, 0.1, 1000.0)
        self.view.camera().setPosition(QVector3D(0, 0, 40))
        self.view.camera().setViewCenter(QVector3D(0, 0, 0))
        self.view.setRootEntity(self.scene)

        self.view_widget = QWidget.createWindowContainer(self.view)
        self.view_widget.setMinimumWidth(400)

        self.controller = Qt3DExtras.QOrbitCameraController(self.scene)
        self.controller.setLinearSpeed(80.0)
        self.controller.setLookSpeed(180.0)
        self.controller.setCamera(self.view.camera())

        self.x_sl = QSlider(Qt.Horizontal)
        self.x_sl.setMinimum(0)
        self.x_sl.setMaximum(360)

        self.y_sl = QSlider(Qt.Horizontal)
        self.y_sl.setMinimum(0)
        self.y_sl.setMaximum(360)

        self.z_sl = QSlider(Qt.Horizontal)
        self.z_sl.setMinimum(0)
        self.z_sl.setMaximum(360)

        self.reset_camera_btn = QPushButton('Reset Camera')

        # Assembly
        self.body.addWidget(self.view_widget)

        self.sliders_body.addRow('Roll', self.x_sl)
        self.sliders_body.addRow('Pitch', self.y_sl)
        self.sliders_body.addRow('Yaw', self.z_sl)
        self.sliders_body.addRow(self.reset_camera_btn)
        self.body.addLayout(self.sliders_body)

        self.container.setLayout(self.body)
        self.setWidget(self.container)

        self.add_mesh(mesh)

        # Functionality
        self.x_sl.valueChanged.connect(lambda: self.rotate())
        self.y_sl.valueChanged.connect(lambda: self.rotate())
        self.z_sl.valueChanged.connect(lambda: self.rotate())
        self.reset_camera_btn.clicked.connect(lambda: self.view.camera().setPosition(QVector3D(0, 0, 40)))
        self.reset_camera_btn.clicked.connect(lambda: self.view.camera().setViewCenter(QVector3D(0, 0, 0)))

    def rotate(self):
        roll, pitch, yaw = self.x_sl.value(), self.y_sl.value(), self.z_sl.value()
        self.mesh_object.rotate(roll, pitch, yaw)

    def add_mesh(self, mesh):
        self.mesh_object = MeshObject(self.scene, mesh)


class MeshObject(Qt3DCore.QEntity):

    def __init__(self, root, mesh):
        super().__init__(root)
        self.root = root
        self.mesh = mesh
        self.material = Qt3DExtras.QDiffuseSpecularMaterial()

        self.transform = Qt3DCore.QTransform()
        self.transform.setScale(4)

        self.addComponent(self.mesh)
        self.addComponent(self.transform)

    def rotate(self, roll, pitch, yaw):
        self.transform.setRotation(QQuaternion.fromEulerAngles(QVector3D(roll, pitch, yaw)))


if '__main__' in __name__:
    app = QApplication([])

    window = QMainWindow()
    window.resize(1200, 600)

    monkey_mesh = Qt3DRender.QSceneLoader()
    # monkey_mesh.setSource(QUrl.fromLocalFile('Monkey.obj'))
    monkey_mesh.setSource('Monkey.obj')

    window.addDockWidget(Qt.TopDockWidgetArea, Mesh3DViewer(monkey_mesh))

    window.show()
    app.exec()