Zordon/src/ui/add_job.py

import copy
import os.path
import pathlib
import socket
import threading

import psutil
from PyQt6.QtCore import QThread, pyqtSignal, Qt, pyqtSlot
from PyQt6.QtWidgets import (
    QApplication, QWidget, QVBoxLayout, QHBoxLayout, QLabel, QLineEdit, QPushButton, QFileDialog, QSpinBox, QComboBox,
    QGroupBox, QCheckBox, QProgressBar, QPlainTextEdit, QDoubleSpinBox, QMessageBox, QListWidget, QListWidgetItem
)
from requests import Response

from src.api.server_proxy import RenderServerProxy
from src.engines.engine_manager import EngineManager
from src.ui.engine_help_viewer import EngineHelpViewer
from src.utilities.zeroconf_server import ZeroconfServer


class NewRenderJobForm(QWidget):
    def __init__(self, project_path=None):
        super().__init__()
        self.project_path = project_path

        # UI
        self.project_group = None
        self.load_file_group = None
        self.current_engine_options = None
        self.file_format_combo = None
        self.renderer_options_layout = None
        self.cameras_list = None
        self.cameras_group = None
        self.renderer_version_combo = None
        self.worker_thread = None
        self.msg_box = None
        self.engine_help_viewer = None
        self.raw_args = None
        self.submit_progress_label = None
        self.submit_progress = None
        self.renderer_type = None
        self.process_label = None
        self.process_progress_bar = None
        self.splitjobs_same_os = None
        self.enable_splitjobs = None
        self.server_input = None
        self.submit_button = None
        self.notes_input = None
        self.priority_input = None
        self.end_frame_input = None
        self.start_frame_input = None
        self.output_path_input = None
        self.scene_file_input = None
        self.scene_file_browse_button = None
        self.job_name_input = None

        # Job / Server Data
        self.server_proxy = RenderServerProxy(socket.gethostname())
        self.renderer_info = None
        self.project_info = None

        # Setup
        self.setWindowTitle("New Job")
        self.setup_ui()
        self.setup_project()

        # get renderer info in bg thread
        t = threading.Thread(target=self.update_renderer_info)
        t.start()

        self.show()

    def setup_ui(self):
        # Main Layout
        main_layout = QVBoxLayout(self)

        # Loading File Group
        self.load_file_group = QGroupBox("Loading")
        load_file_layout = QVBoxLayout(self.load_file_group)
        # progress bar
        progress_layout = QHBoxLayout()
        self.process_progress_bar = QProgressBar()
        self.process_progress_bar.setMinimum(0)
        self.process_progress_bar.setMaximum(0)
        self.process_label = QLabel("Processing")
        progress_layout.addWidget(self.process_label)
        progress_layout.addWidget(self.process_progress_bar)
        load_file_layout.addLayout(progress_layout)
        main_layout.addWidget(self.load_file_group)

        # Project Group
        self.project_group = QGroupBox("Project")
        server_layout = QVBoxLayout(self.project_group)
        # File Path
        scene_file_picker_layout = QHBoxLayout()
        self.scene_file_input = QLineEdit()
        self.scene_file_input.setText(self.project_path)
        self.scene_file_browse_button = QPushButton("Browse...")
        self.scene_file_browse_button.clicked.connect(self.browse_scene_file)
        scene_file_picker_layout.addWidget(QLabel("File:"))
        scene_file_picker_layout.addWidget(self.scene_file_input)
        scene_file_picker_layout.addWidget(self.scene_file_browse_button)
        server_layout.addLayout(scene_file_picker_layout)
        # Server List
        server_list_layout = QHBoxLayout()
        server_list_layout.setSpacing(0)
        self.server_input = QComboBox()
        server_list_layout.addWidget(QLabel("Hostname:"), 1)
        server_list_layout.addWidget(self.server_input, 3)
        server_layout.addLayout(server_list_layout)
        main_layout.addWidget(self.project_group)
        self.update_server_list()
        # Priority
        priority_layout = QHBoxLayout()
        priority_layout.addWidget(QLabel("Priority:"), 1)
        self.priority_input = QComboBox()
        self.priority_input.addItems(["High", "Medium", "Low"])
        self.priority_input.setCurrentIndex(1)
        priority_layout.addWidget(self.priority_input, 3)
        server_layout.addLayout(priority_layout)
        # Splitjobs
        self.enable_splitjobs = QCheckBox("Automatically split render across multiple servers")
        self.enable_splitjobs.setEnabled(True)
        server_layout.addWidget(self.enable_splitjobs)
        self.splitjobs_same_os = QCheckBox("Only render on same OS")
        self.splitjobs_same_os.setEnabled(True)
        server_layout.addWidget(self.splitjobs_same_os)

        # Output Settings Group
        self.output_settings_group = QGroupBox("Output Settings")
        output_settings_layout = QVBoxLayout(self.output_settings_group)
        # output path
        output_path_layout = QHBoxLayout()
        output_path_layout.addWidget(QLabel("Render name:"))
        self.output_path_input = QLineEdit()
        output_path_layout.addWidget(self.output_path_input)
        output_settings_layout.addLayout(output_path_layout)
        # file format
        file_format_layout = QHBoxLayout()
        file_format_layout.addWidget(QLabel("Format:"))
        self.file_format_combo = QComboBox()
        file_format_layout.addWidget(self.file_format_combo)
        output_settings_layout.addLayout(file_format_layout)
        # frame range
        frame_range_layout = QHBoxLayout(self.output_settings_group)
        self.start_frame_input = QSpinBox()
        self.start_frame_input.setRange(1, 99999)
        self.end_frame_input = QSpinBox()
        self.end_frame_input.setRange(1, 99999)
        frame_range_layout.addWidget(QLabel("Frames:"))
        frame_range_layout.addWidget(self.start_frame_input)
        frame_range_layout.addWidget(QLabel("to"))
        frame_range_layout.addWidget(self.end_frame_input)
        output_settings_layout.addLayout(frame_range_layout)
        # resolution
        resolution_layout = QHBoxLayout(self.output_settings_group)
        self.resolution_x_input = QSpinBox()
        self.resolution_x_input.setRange(1, 9999)  # Assuming max resolution width 9999
        self.resolution_x_input.setValue(1920)
        self.resolution_y_input = QSpinBox()
        self.resolution_y_input.setRange(1, 9999)  # Assuming max resolution height 9999
        self.resolution_y_input.setValue(1080)
        self.frame_rate_input = QDoubleSpinBox()
        self.frame_rate_input.setRange(1, 9999)  # Assuming max resolution width 9999
        self.frame_rate_input.setDecimals(3)
        self.frame_rate_input.setValue(23.976)
        resolution_layout.addWidget(QLabel("Resolution:"))
        resolution_layout.addWidget(self.resolution_x_input)
        resolution_layout.addWidget(QLabel("x"))
        resolution_layout.addWidget(self.resolution_y_input)
        resolution_layout.addWidget(QLabel("@"))
        resolution_layout.addWidget(self.frame_rate_input)
        resolution_layout.addWidget(QLabel("fps"))
        output_settings_layout.addLayout(resolution_layout)
        # add group to layout
        main_layout.addWidget(self.output_settings_group)

        # Renderer Group
        self.renderer_group = QGroupBox("Renderer Settings")
        renderer_group_layout = QVBoxLayout(self.renderer_group)
        renderer_layout = QHBoxLayout()
        renderer_layout.addWidget(QLabel("Renderer:"))
        self.renderer_type = QComboBox()
        self.renderer_type.currentIndexChanged.connect(self.renderer_changed)
        renderer_layout.addWidget(self.renderer_type)
        # Version
        renderer_layout.addWidget(QLabel("Version:"))
        self.renderer_version_combo = QComboBox()
        self.renderer_version_combo.addItem('latest')
        renderer_layout.addWidget(self.renderer_version_combo)
        renderer_group_layout.addLayout(renderer_layout)
        # dynamic options
        self.renderer_options_layout = QVBoxLayout()
        renderer_group_layout.addLayout(self.renderer_options_layout)
        # Raw Args
        raw_args_layout = QHBoxLayout(self.renderer_group)
        raw_args_layout.addWidget(QLabel("Raw Args:"))
        self.raw_args = QLineEdit()
        raw_args_layout.addWidget(self.raw_args)
        args_help_button = QPushButton("?")
        args_help_button.clicked.connect(self.args_help_button_clicked)
        raw_args_layout.addWidget(args_help_button)
        renderer_group_layout.addLayout(raw_args_layout)
        main_layout.addWidget(self.renderer_group)

        # Cameras Group
        self.cameras_group = QGroupBox("Cameras")
        cameras_layout = QVBoxLayout(self.cameras_group)
        self.cameras_list = QListWidget()
        self.cameras_group.setHidden(True)
        cameras_layout.addWidget(self.cameras_list)
        main_layout.addWidget(self.cameras_group)

        # Notes Group
        self.notes_group = QGroupBox("Additional Notes")
        notes_layout = QVBoxLayout(self.notes_group)
        self.notes_input = QPlainTextEdit()
        notes_layout.addWidget(self.notes_input)
        main_layout.addWidget(self.notes_group)

        # Submit Button
        self.submit_button = QPushButton("Submit Job")
        self.submit_button.clicked.connect(self.submit_job)
        main_layout.addWidget(self.submit_button)

        self.submit_progress = QProgressBar()
        self.submit_progress.setMinimum(0)
        self.submit_progress.setMaximum(0)
        self.submit_progress.setHidden(True)
        main_layout.addWidget(self.submit_progress)

        self.submit_progress_label = QLabel("Submitting...")
        self.submit_progress_label.setHidden(True)
        main_layout.addWidget(self.submit_progress_label)

        self.toggle_renderer_enablement(False)

    def update_renderer_info(self):
        # get the renderer info and add them all to the ui
        self.renderer_info = self.server_proxy.get_renderer_info()
        self.renderer_type.addItems(self.renderer_info.keys())
        # select the best renderer for the file type
        engine = EngineManager.engine_for_project_path(self.project_path)
        self.renderer_type.setCurrentText(engine.name().lower())
        # refresh ui
        self.renderer_changed()

    def renderer_changed(self):
        # load the version numbers
        current_renderer = self.renderer_type.currentText().lower() or self.renderer_type.itemText(0)
        self.renderer_version_combo.clear()
        self.renderer_version_combo.addItem('latest')
        self.file_format_combo.clear()
        if current_renderer:
            renderer_vers = [version_info['version'] for version_info in self.renderer_info[current_renderer]['versions']]
            self.renderer_version_combo.addItems(renderer_vers)
            self.file_format_combo.addItems(self.renderer_info[current_renderer]['supported_export_formats'])

    def update_server_list(self):
        clients = ZeroconfServer.found_hostnames()
        self.server_input.clear()
        self.server_input.addItems(clients)

    def browse_scene_file(self):
        file_name, _ = QFileDialog.getOpenFileName(self, "Select Scene File")
        if file_name:
            self.scene_file_input.setText(file_name)
            self.setup_project()

    def setup_project(self):
        # UI stuff on main thread
        self.process_progress_bar.setHidden(False)
        self.process_label.setHidden(False)
        self.toggle_renderer_enablement(False)

        output_name, _ = os.path.splitext(os.path.basename(self.scene_file_input.text()))
        output_name = output_name.replace(' ', '_')
        self.output_path_input.setText(output_name)
        file_name = self.scene_file_input.text()

        # setup bg worker
        self.worker_thread = GetProjectInfoWorker(window=self, project_path=file_name)
        self.worker_thread.message_signal.connect(self.post_get_project_info_update)
        self.worker_thread.start()

    def browse_output_path(self):
        directory = QFileDialog.getExistingDirectory(self, "Select Output Directory")
        if directory:
            self.output_path_input.setText(directory)

    def args_help_button_clicked(self):
        url = (f'http://{self.server_proxy.hostname}:{self.server_proxy.port}/api/renderer/'
               f'{self.renderer_type.currentText()}/help')
        self.engine_help_viewer = EngineHelpViewer(url)
        self.engine_help_viewer.show()

# -------- Update --------

    def post_get_project_info_update(self):
        """Called by the GetProjectInfoWorker - Do not call directly."""
        try:
            # Set the best renderer we can find
            input_path = self.scene_file_input.text()
            engine = EngineManager.engine_for_project_path(input_path)

            engine_index = self.renderer_type.findText(engine.name().lower())
            if engine_index >= 0:
                self.renderer_type.setCurrentIndex(engine_index)
            else:
                self.renderer_type.setCurrentIndex(0)  #todo: find out why we don't have renderer info yet
                # not ideal but if we don't have the renderer info we have to pick something

            self.output_path_input.setText(os.path.basename(input_path))

            # cleanup progress UI
            self.load_file_group.setHidden(True)
            self.toggle_renderer_enablement(True)

            # Load scene data
            self.start_frame_input.setValue(self.project_info.get('frame_start'))
            self.end_frame_input.setValue(self.project_info.get('frame_end'))
            self.resolution_x_input.setValue(self.project_info.get('resolution_x'))
            self.resolution_y_input.setValue(self.project_info.get('resolution_y'))
            self.frame_rate_input.setValue(self.project_info.get('fps'))

            # Cameras
            self.cameras_list.clear()
            if self.project_info.get('cameras'):
                self.cameras_group.setHidden(False)
                found_active = False
                for camera in self.project_info['cameras']:
                    # create the list items and make them checkable
                    item = QListWidgetItem(f"{camera['name']} - {camera['lens']}mm")
                    item.setFlags(item.flags() | Qt.ItemFlag.ItemIsUserCheckable)
                    is_checked = camera['is_active'] or len(self.project_info['cameras']) == 1
                    found_active = found_active or is_checked
                    item.setCheckState(Qt.CheckState.Checked if is_checked else Qt.CheckState.Unchecked)
                    self.cameras_list.addItem(item)
                if not found_active:
                    self.cameras_list.item(0).setCheckState(Qt.CheckState.Checked)
            else:
                self.cameras_group.setHidden(True)

            # Dynamic Engine Options
            clear_layout(self.renderer_options_layout)  # clear old options
            # dynamically populate option list
            self.current_engine_options = engine().ui_options()
            for option in self.current_engine_options:
                h_layout = QHBoxLayout()
                label = QLabel(option['name'].capitalize() + ':')
                h_layout.addWidget(label)
                if option.get('options'):
                    combo_box = QComboBox()
                    for opt in option['options']:
                        combo_box.addItem(opt)
                    h_layout.addWidget(combo_box)
                else:
                    text_box = QLineEdit()
                    h_layout.addWidget(text_box)
                self.renderer_options_layout.addLayout(h_layout)
        except AttributeError as e:
            pass

    def toggle_renderer_enablement(self, enabled=False):
        """Toggle on/off all the render settings"""
        self.project_group.setHidden(not enabled)
        self.output_settings_group.setHidden(not enabled)
        self.renderer_group.setHidden(not enabled)
        self.notes_group.setHidden(not enabled)
        if not enabled:
            self.cameras_group.setHidden(True)
        self.submit_button.setEnabled(enabled)

    def after_job_submission(self, result):

        # UI cleanup
        self.submit_progress.setMaximum(0)
        self.submit_button.setHidden(False)
        self.submit_progress.setHidden(True)
        self.submit_progress_label.setHidden(True)
        self.process_progress_bar.setHidden(True)
        self.process_label.setHidden(True)
        self.toggle_renderer_enablement(True)

        self.msg_box = QMessageBox()
        if result.ok:
            self.msg_box.setStandardButtons(QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No)
            self.msg_box.setIcon(QMessageBox.Icon.Information)
            self.msg_box.setText("Job successfully submitted to server. Submit another?")
            self.msg_box.setWindowTitle("Success")
            x = self.msg_box.exec()
            if x == QMessageBox.StandardButton.No:
                self.close()
        else:
            self.msg_box.setStandardButtons(QMessageBox.StandardButton.Ok)
            self.msg_box.setIcon(QMessageBox.Icon.Critical)
            self.msg_box.setText(result.text or "Unknown error")
            self.msg_box.setWindowTitle("Error")
            self.msg_box.exec()

# -------- Submit Job Calls --------

    def submit_job(self):

        # Pre-worker UI
        self.submit_progress.setHidden(False)
        self.submit_progress_label.setHidden(False)
        self.submit_button.setHidden(True)
        self.submit_progress.setMaximum(0)

        # submit job in background thread
        self.worker_thread = SubmitWorker(window=self)
        self.worker_thread.update_ui_signal.connect(self.update_submit_progress)
        self.worker_thread.message_signal.connect(self.after_job_submission)
        self.worker_thread.start()

    @pyqtSlot(str, str)
    def update_submit_progress(self, hostname, percent):
        # Update the UI here. This slot will be executed in the main thread
        self.submit_progress_label.setText(f"Transferring to {hostname} - {percent}%")
        self.submit_progress.setMaximum(100)
        self.submit_progress.setValue(int(percent))


class SubmitWorker(QThread):
    """Worker class called to submit all the jobs to the server and update the UI accordingly"""

    message_signal = pyqtSignal(Response)
    update_ui_signal = pyqtSignal(str, str)

    def __init__(self, window):
        super().__init__()
        self.window = window

    def run(self):
        def create_callback(encoder):
            encoder_len = encoder.len

            def callback(monitor):
                percent = f"{monitor.bytes_read / encoder_len * 100:.0f}"
                self.update_ui_signal.emit(hostname, percent)
            return callback

        hostname = self.window.server_input.currentText()
        job_json = {'owner': psutil.Process().username() + '@' + socket.gethostname(),
                    'renderer': self.window.renderer_type.currentText().lower(),
                    'engine_version': self.window.renderer_version_combo.currentText(),
                    'args': {'raw': self.window.raw_args.text()},
                    'output_path': self.window.output_path_input.text(),
                    'start_frame': self.window.start_frame_input.value(),
                    'end_frame': self.window.end_frame_input.value(),
                    'priority': self.window.priority_input.currentIndex() + 1,
                    'notes': self.window.notes_input.toPlainText(),
                    'enable_split_jobs': self.window.enable_splitjobs.isChecked(),
                    'split_jobs_same_os': self.window.splitjobs_same_os.isChecked()}

        # get the dynamic args
        for i in range(self.window.renderer_options_layout.count()):
            item = self.window.renderer_options_layout.itemAt(i)
            layout = item.layout()  # get the layout
            for x in range(layout.count()):
                z = layout.itemAt(x)
                widget = z.widget()
                if isinstance(widget, QComboBox):
                    job_json['args'][self.window.current_engine_options[i]['name']] = widget.currentText()
                elif isinstance(widget, QLineEdit):
                    job_json['args'][self.window.current_engine_options[i]['name']] = widget.text()

        # determine if any cameras are checked
        selected_cameras = []
        if self.window.cameras_list.count() and not self.window.cameras_group.isHidden():
            for index in range(self.window.cameras_list.count()):
                item = self.window.cameras_list.item(index)
                if item.checkState() == Qt.CheckState.Checked:
                    selected_cameras.append(item.text().rsplit('-', 1)[0].strip())  # cleanup to just camera name

        # process cameras into nested format
        input_path = self.window.scene_file_input.text()
        if selected_cameras:
            job_list = []
            for cam in selected_cameras:
                job_copy = copy.deepcopy(job_json)
                job_copy['args']['camera'] = cam
                job_copy['name'] = pathlib.Path(input_path).stem.replace(' ', '_') + "-" + cam.replace(' ', '')
                job_list.append(job_copy)
        else:
            job_list = [job_json]

        # presubmission tasks
        engine = EngineManager.engine_with_name(self.window.renderer_type.currentText().lower())
        input_path = engine().perform_presubmission_tasks(input_path)
        # submit
        result = self.window.server_proxy.post_job_to_server(file_path=input_path, job_list=job_list,
                                                      callback=create_callback)
        self.message_signal.emit(result)


class GetProjectInfoWorker(QThread):
    """Worker class called to retrieve information about a project file on a background thread and update the UI"""

    message_signal = pyqtSignal()

    def __init__(self, window, project_path):
        super().__init__()
        self.window = window
        self.project_path = project_path

    def run(self):
        engine = EngineManager.engine_for_project_path(self.project_path)
        self.window.project_info = engine().get_project_info(self.project_path)
        self.message_signal.emit()


def clear_layout(layout):
    if layout is not None:
        # Go through the layout's items in reverse order
        for i in reversed(range(layout.count())):
            # Take the item at the current position
            item = layout.takeAt(i)

            # Check if the item is a widget
            if item.widget():
                # Remove the widget and delete it
                widget_to_remove = item.widget()
                widget_to_remove.setParent(None)
                widget_to_remove.deleteLater()
            elif item.layout():
                # If the item is a sub-layout, clear its contents recursively
                clear_layout(item.layout())
                # Then delete the layout
                item.layout().deleteLater()

# Run the application
if __name__ == '__main__':
    app = QApplication([])
    window = NewRenderJobForm()
    app.exec()