Zordon/src/ui/main_window.py

''' app/ui/main_window.py '''
import datetime
import io
import logging
import os
import subprocess
import sys
import threading
import time

import PIL
from PIL import Image
from PyQt6.QtCore import Qt, QByteArray, QBuffer, QIODevice, QThread
from PyQt6.QtGui import QPixmap, QImage, QFont, QIcon
from PyQt6.QtWidgets import QMainWindow, QWidget, QHBoxLayout, QListWidget, QTableWidget, QAbstractItemView, \
    QTableWidgetItem, QLabel, QVBoxLayout, QHeaderView, QMessageBox, QGroupBox, QPushButton, QListWidgetItem, \
    QFileDialog

from src.render_queue import RenderQueue
from src.utilities.misc_helper import get_time_elapsed, resources_dir, is_localhost
from src.utilities.status_utils import RenderStatus
from src.utilities.zeroconf_server import ZeroconfServer
from .add_job import NewRenderJobForm
from .console import ConsoleWindow
from .engine_browser import EngineBrowserWindow
from .log_viewer import LogViewer
from .widgets.menubar import MenuBar
from .widgets.proportional_image_label import ProportionalImageLabel
from .widgets.statusbar import StatusBar
from .widgets.toolbar import ToolBar
from src.api.serverproxy_manager import ServerProxyManager

logger = logging.getLogger()


class MainWindow(QMainWindow):
    """
    MainWindow

    Args:
        QMainWindow (QMainWindow): Inheritance
    """

    def __init__(self) -> None:
        """
        Initialize the Main-Window.
        """
        super().__init__()

        # Load the queue
        self.job_list_view = None
        self.server_info_ram = None
        self.server_info_cpu = None
        self.server_info_os = None
        self.server_info_hostname = None
        self.engine_browser_window = None
        self.server_info_group = None
        self.current_hostname = None
        self.subprocess_runner = None

        # To pass to console
        self.buffer_handler = None

        # Window-Settings
        self.setWindowTitle("Zordon")
        self.setGeometry(100, 100, 900, 800)
        central_widget = QWidget(self)
        self.setCentralWidget(central_widget)

        main_layout = QHBoxLayout(central_widget)

        # Create a QLabel widget to display the image
        self.image_label = ProportionalImageLabel()
        self.image_label.setMaximumSize(700, 500)
        self.image_label.setFixedHeight(500)
        self.image_label.setAlignment(Qt.AlignmentFlag.AlignTop | Qt.AlignmentFlag.AlignHCenter)
        self.load_image_path(os.path.join(resources_dir(), 'Rectangle.png'))

        # Server list
        self.server_list_view = QListWidget()
        self.server_list_view.itemClicked.connect(self.server_picked)
        list_font = QFont()
        list_font.setPointSize(16)
        self.server_list_view.setFont(list_font)
        self.added_hostnames = []

        self.setup_ui(main_layout)

        self.create_toolbars()

        # Add Widgets to Window
        self.setMenuBar(MenuBar(self))
        self.setStatusBar(StatusBar(self))

        # start background update
        self.bg_update_thread = QThread()
        self.bg_update_thread.run = self.__background_update
        self.bg_update_thread.start()

        # Setup other windows
        self.new_job_window = None
        self.console_window = None
        self.log_viewer_window = None

        # Pick default job
        self.job_picked()

    def setup_ui(self, main_layout):

        # Servers
        server_list_group = QGroupBox("Available Servers")
        list_layout = QVBoxLayout()
        list_layout.addWidget(self.server_list_view)
        list_layout.setContentsMargins(0, 0, 0, 0)
        server_list_group.setLayout(list_layout)
        server_info_group = QGroupBox("Server Info")

        # Server Info Group
        self.server_info_hostname = QLabel()
        self.server_info_os = QLabel()
        self.server_info_cpu = QLabel()
        self.server_info_ram = QLabel()
        server_info_engines_button = QPushButton("Render Engines")
        server_info_engines_button.clicked.connect(self.engine_browser)
        server_info_layout = QVBoxLayout()
        server_info_layout.addWidget(self.server_info_hostname)
        server_info_layout.addWidget(self.server_info_os)
        server_info_layout.addWidget(self.server_info_cpu)
        server_info_layout.addWidget(self.server_info_ram)
        server_info_layout.addWidget(server_info_engines_button)
        server_info_group.setLayout(server_info_layout)

        # Server Button Layout
        server_button_layout = QHBoxLayout()
        add_server_button = QPushButton(text="+")
        remove_server_button = QPushButton(text="-")
        server_button_layout.addWidget(add_server_button)
        server_button_layout.addWidget(remove_server_button)

        # Layouts
        info_layout = QVBoxLayout()
        info_layout.addWidget(server_list_group, stretch=True)
        info_layout.addWidget(server_info_group)
        info_layout.setContentsMargins(0, 0, 0, 0)
        server_list_group.setFixedWidth(260)
        self.server_picked()

        # Job list
        self.job_list_view = QTableWidget()
        self.job_list_view.setRowCount(0)
        self.job_list_view.setColumnCount(8)
        self.job_list_view.setSelectionBehavior(QAbstractItemView.SelectionBehavior.SelectRows)
        self.job_list_view.verticalHeader().setVisible(False)
        self.job_list_view.itemSelectionChanged.connect(self.job_picked)
        self.job_list_view.setEditTriggers(QAbstractItemView.EditTrigger.NoEditTriggers)
        self.refresh_job_headers()

        # Image Layout
        image_group = QGroupBox("Job Preview")
        image_layout = QVBoxLayout(image_group)
        image_layout.setContentsMargins(0, 0, 0, 0)
        image_center_layout = QHBoxLayout()
        image_center_layout.addWidget(self.image_label)
        image_layout.addWidget(self.image_label)
        # image_layout.addLayout(image_center_layout)

        # Job Layout
        job_list_group = QGroupBox("Render Jobs")
        job_list_layout = QVBoxLayout(job_list_group)
        job_list_layout.setContentsMargins(0, 0, 0, 0)
        image_layout.addWidget(self.job_list_view, stretch=True)
        image_layout.addLayout(job_list_layout)

        # Add them all to the window
        main_layout.addLayout(info_layout)

        right_layout = QVBoxLayout()
        right_layout.setContentsMargins(0, 0, 0, 0)
        right_layout.addWidget(image_group)
        # right_layout.addWidget(job_list_group)
        main_layout.addLayout(right_layout)

    def __background_update(self):
        while True:
            self.update_servers()
            self.fetch_jobs()
            time.sleep(0.5)

    def closeEvent(self, event):
        running_jobs = len(RenderQueue.running_jobs())
        if running_jobs:
            reply = QMessageBox.question(self, "Running Jobs",
                                         f"You have {running_jobs} jobs running.\n"
                                         f"Quitting will cancel these renders. Continue?",
                                         QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.Cancel,
                                         QMessageBox.StandardButton.Cancel)
            if reply == QMessageBox.StandardButton.Yes:
                event.accept()
            else:
                event.ignore()

# -- Server Code -- #

    @property
    def current_server_proxy(self):
        return ServerProxyManager.get_proxy_for_hostname(self.current_hostname)

    def server_picked(self):
        """Update the UI elements relevant to the server selection."""
        try:
            # Retrieve the new hostname selected by the user
            new_hostname = self.server_list_view.currentItem().text()

            # Check if the hostname has changed to avoid unnecessary updates
            if new_hostname != self.current_hostname:
                # Update the current hostname and clear the job list
                self.current_hostname = new_hostname
                self.job_list_view.setRowCount(0)
                self.fetch_jobs(clear_table=True)

                # Select the first row if there are jobs listed
                if self.job_list_view.rowCount():
                    self.job_list_view.selectRow(0)

                # Update server information display
                self.update_server_info_display(new_hostname)

        except AttributeError:
            # Handle cases where the server list view might not be properly initialized
            pass

    def update_server_info_display(self, hostname):
        """Updates the server information section of the UI."""
        self.server_info_hostname.setText(hostname or "unknown")
        server_info = ZeroconfServer.get_hostname_properties(hostname)

        # Use the get method with defaults to avoid KeyError
        os_info = f"OS: {server_info.get('system_os', 'Unknown')} {server_info.get('system_os_version', '')}"
        cpu_info = f"CPU: {server_info.get('system_cpu', 'Unknown')} - {server_info.get('system_cpu_cores', 'Unknown')} cores"

        self.server_info_os.setText(os_info.strip())
        self.server_info_cpu.setText(cpu_info)

    def fetch_jobs(self, clear_table=False):

        if not self.current_server_proxy:
            return

        if clear_table:
            self.job_list_view.clear()
            self.refresh_job_headers()

        job_fetch = self.current_server_proxy.get_all_jobs(ignore_token=clear_table)
        if job_fetch:
            num_jobs = len(job_fetch)
            self.job_list_view.setRowCount(num_jobs)

            for row, job in enumerate(job_fetch):

                display_status = job['status'] if job['status'] != RenderStatus.RUNNING.value else \
                    ('%.0f%%' % (job['percent_complete'] * 100))  # if running, show percent, otherwise just show status
                tags = (job['status'],)
                start_time = datetime.datetime.fromisoformat(job['start_time']) if job['start_time'] else None
                end_time = datetime.datetime.fromisoformat(job['end_time']) if job['end_time'] else None

                time_elapsed = "" if (job['status'] != RenderStatus.RUNNING.value and not end_time) else \
                    get_time_elapsed(start_time, end_time)

                name = job.get('name') or os.path.basename(job.get('input_path', ''))
                renderer = f"{job.get('renderer', '')}-{job.get('renderer_version')}"
                priority = str(job.get('priority', ''))
                total_frames = str(job.get('total_frames', ''))

                items = [QTableWidgetItem(job['id']), QTableWidgetItem(name), QTableWidgetItem(renderer),
                         QTableWidgetItem(priority), QTableWidgetItem(display_status), QTableWidgetItem(time_elapsed),
                         QTableWidgetItem(total_frames), QTableWidgetItem(job['date_created'])]

                for col, item in enumerate(items):
                    self.job_list_view.setItem(row, col, item)

# -- Job Code -- #
    def job_picked(self):

        def fetch_preview(job_id):
            try:
                default_image_path = "error.png"
                before_fetch_hostname = self.current_server_proxy.hostname

                response = self.current_server_proxy.request(f'job/{job_id}/thumbnail?size=big')
                if response.ok:
                    try:
                        with io.BytesIO(response.content) as image_data_stream:
                            image = Image.open(image_data_stream)
                            if self.current_server_proxy.hostname == before_fetch_hostname and job_id == \
                                    self.selected_job_ids()[0]:
                                self.load_image_data(image)
                                return
                    except PIL.UnidentifiedImageError:
                        default_image_path = response.text
                else:
                    default_image_path = default_image_path or response.text

                self.load_image_path(os.path.join(resources_dir(), default_image_path))

            except ConnectionError as e:
                logger.error(f"Connection error fetching image: {e}")
            except Exception as e:
                logger.exception(f"Error fetching image: {e}")

        job_id = self.selected_job_ids()[0] if self.selected_job_ids() else None
        local_server = is_localhost(self.current_hostname)

        if job_id:
            fetch_thread = threading.Thread(target=fetch_preview, args=(job_id,))
            fetch_thread.daemon = True
            fetch_thread.start()

            selected_row = self.job_list_view.selectionModel().selectedRows()[0]
            current_status = self.job_list_view.item(selected_row.row(), 4).text()

            # show / hide the stop button
            show_stop_button = current_status.lower() == 'running'
            self.topbar.actions_call['Stop Job'].setEnabled(show_stop_button)
            self.topbar.actions_call['Stop Job'].setVisible(show_stop_button)
            self.topbar.actions_call['Delete Job'].setEnabled(not show_stop_button)
            self.topbar.actions_call['Delete Job'].setVisible(not show_stop_button)

            self.topbar.actions_call['Render Log'].setEnabled(True)
            self.topbar.actions_call['Download'].setEnabled(not local_server)
            self.topbar.actions_call['Download'].setVisible(not local_server)
            self.topbar.actions_call['Open Files'].setEnabled(local_server)
            self.topbar.actions_call['Open Files'].setVisible(local_server)
        else:
            # load default
            default_image_path = os.path.join(resources_dir(), 'Rectangle.png')
            self.load_image_path(default_image_path)
            self.topbar.actions_call['Stop Job'].setVisible(False)
            self.topbar.actions_call['Stop Job'].setEnabled(False)
            self.topbar.actions_call['Delete Job'].setEnabled(False)
            self.topbar.actions_call['Render Log'].setEnabled(False)
            self.topbar.actions_call['Download'].setEnabled(False)
            self.topbar.actions_call['Download'].setVisible(True)
            self.topbar.actions_call['Open Files'].setEnabled(False)
            self.topbar.actions_call['Open Files'].setVisible(False)

    def selected_job_ids(self):
        try:
            selected_rows = self.job_list_view.selectionModel().selectedRows()
            job_ids = []
            for selected_row in selected_rows:
                id_item = self.job_list_view.item(selected_row.row(), 0)
                job_ids.append(id_item.text())
            return job_ids
        except AttributeError:
            return []

    def refresh_job_headers(self):
        self.job_list_view.setHorizontalHeaderLabels(["ID", "Name", "Renderer", "Priority", "Status",
                                                      "Time Elapsed", "Frames", "Date Created"])
        self.job_list_view.setColumnHidden(0, True)

        self.job_list_view.horizontalHeader().setSectionResizeMode(1, QHeaderView.ResizeMode.Stretch)
        self.job_list_view.horizontalHeader().setSectionResizeMode(2, QHeaderView.ResizeMode.ResizeToContents)
        self.job_list_view.horizontalHeader().setSectionResizeMode(3, QHeaderView.ResizeMode.ResizeToContents)
        self.job_list_view.horizontalHeader().setSectionResizeMode(4, QHeaderView.ResizeMode.ResizeToContents)
        self.job_list_view.horizontalHeader().setSectionResizeMode(5, QHeaderView.ResizeMode.ResizeToContents)
        self.job_list_view.horizontalHeader().setSectionResizeMode(6, QHeaderView.ResizeMode.ResizeToContents)
        self.job_list_view.horizontalHeader().setSectionResizeMode(7, QHeaderView.ResizeMode.ResizeToContents)

# -- Image Code -- #

    def load_image_path(self, image_path):
        # Load and set the image using QPixmap
        pixmap = QPixmap(image_path)
        if not pixmap:
            logger.error("Error loading image")
            return
        self.image_label.setPixmap(pixmap)

    def load_image_data(self, pillow_image):
        # Convert the Pillow Image to a QByteArray (byte buffer)
        byte_array = QByteArray()
        buffer = QBuffer(byte_array)
        buffer.open(QIODevice.OpenModeFlag.WriteOnly)
        pillow_image.save(buffer, "PNG")
        buffer.close()

        # Create a QImage from the QByteArray
        image = QImage.fromData(byte_array)

        # Create a QPixmap from the QImage
        pixmap = QPixmap.fromImage(image)

        if not pixmap:
            logger.error("Error loading image")
            return
        self.image_label.setPixmap(pixmap)

    def update_servers(self):
        found_servers = list(set(ZeroconfServer.found_hostnames() + self.added_hostnames))
        # Always make sure local hostname is first
        if found_servers and not is_localhost(found_servers[0]):
            for hostname in found_servers:
                if is_localhost(hostname):
                    found_servers.remove(hostname)
                    found_servers.insert(0, hostname)
                    break

        old_count = self.server_list_view.count()

        # Update proxys
        for hostname in found_servers:
            ServerProxyManager.get_proxy_for_hostname(hostname)  # setup background updates

        # Add in all the missing servers
        current_server_list = []
        for i in range(self.server_list_view.count()):
            current_server_list.append(self.server_list_view.item(i).text())
        for hostname in found_servers:
            if hostname not in current_server_list:
                properties = ZeroconfServer.get_hostname_properties(hostname)
                image_path = os.path.join(resources_dir(), f"{properties.get('system_os', 'Monitor')}.png")
                list_widget = QListWidgetItem(QIcon(image_path), hostname)
                self.server_list_view.addItem(list_widget)

        # find any servers that shouldn't be shown any longer
        servers_to_remove = []
        for i in range(self.server_list_view.count()):
            name = self.server_list_view.item(i).text()
            if name not in found_servers:
                servers_to_remove.append(name)

        # remove any servers that shouldn't be shown any longer
        for server in servers_to_remove:
            # Find and remove the item with the specified text
            for i in range(self.server_list_view.count()):
                item = self.server_list_view.item(i)
                if item is not None and item.text() == server:
                    self.server_list_view.takeItem(i)
                    break  # Stop searching after the first match is found

        if not old_count and self.server_list_view.count():
            self.server_list_view.setCurrentRow(0)
            self.server_picked()

    def create_toolbars(self) -> None:
        """
        Creates and adds the top and right toolbars to the main window.
        """
        # Top Toolbar [PyQt6.QtWidgets.QToolBar]
        self.topbar = ToolBar(self, orientation=Qt.Orientation.Horizontal,
                              style=Qt.ToolButtonStyle.ToolButtonTextUnderIcon, icon_size=(24, 24))
        self.topbar.setMovable(False)

        resources_directory = resources_dir()

        # Top Toolbar Buttons
        self.topbar.add_button(
            "Console", f"{resources_directory}/Console.png", self.open_console_window)
        self.topbar.add_button(
            "Engines", f"{resources_directory}/SoftwareInstaller.png", self.engine_browser)
        self.topbar.add_separator()
        self.topbar.add_button(
            "Stop Job", f"{resources_directory}/StopSign.png", self.stop_job)
        self.topbar.add_button(
            "Delete Job", f"{resources_directory}/Trash.png", self.delete_job)
        self.topbar.add_button(
            "Render Log", f"{resources_directory}/Document.png", self.job_logs)
        self.topbar.add_button(
            "Download", f"{resources_directory}/Download.png", self.download_files)
        self.topbar.add_button(
            "Open Files", f"{resources_directory}/SearchFolder.png", self.open_files)
        self.topbar.add_button(
            "New Job", f"{resources_directory}/AddProduct.png", self.new_job)
        self.addToolBar(Qt.ToolBarArea.TopToolBarArea, self.topbar)

# -- Toolbar Buttons -- #

    def open_console_window(self) -> None:
        """
        Event handler for the "Open Console" button
        """
        self.console_window = ConsoleWindow(self.buffer_handler)
        self.console_window.buffer_handler = self.buffer_handler
        self.console_window.show()

    def engine_browser(self):
        self.engine_browser_window = EngineBrowserWindow(hostname=self.current_hostname)
        self.engine_browser_window.show()

    def job_logs(self) -> None:
        """
        Event handler for the "Logs" button.
        """
        selected_job_ids = self.selected_job_ids()
        if selected_job_ids:
            url = f'http://{self.current_server_proxy.hostname}:{self.current_server_proxy.port}/api/job/{selected_job_ids[0]}/logs'
            self.log_viewer_window = LogViewer(url)
            self.log_viewer_window.show()

    def stop_job(self, event):
        """
        Event handler for the "Exit" button. Closes the application.
        """
        job_ids = self.selected_job_ids()
        if not job_ids:
            return

        if len(job_ids) == 1:
            job = next((job for job in self.current_server_proxy.get_all_jobs() if job.get('id') == job_ids[0]), None)
            if job:
                display_name = job.get('name', os.path.basename(job.get('input_path', '')))
                message = f"Are you sure you want to delete the job:\n{display_name}?"
            else:
                return  # Job not found, handle this case as needed
        else:
            message = f"Are you sure you want to delete these {len(job_ids)} jobs?"

        # Display the message box and check the response in one go
        msg_box = QMessageBox(QMessageBox.Icon.Warning, "Delete Job", message,
                              QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, self)

        if msg_box.exec() == QMessageBox.StandardButton.Yes:
            for job_id in job_ids:
                self.current_server_proxy.cancel_job(job_id, confirm=True)
            self.fetch_jobs(clear_table=True)

    def delete_job(self, event):
        """
        Event handler for the Delete Job button
        """
        job_ids = self.selected_job_ids()
        if not job_ids:
            return

        if len(job_ids) == 1:
            job = next((job for job in self.current_server_proxy.get_all_jobs() if job.get('id') == job_ids[0]), None)
            if job:
                display_name = job.get('name', os.path.basename(job.get('input_path', '')))
                message = f"Are you sure you want to delete the job:\n{display_name}?"
            else:
                return  # Job not found, handle this case as needed
        else:
            message = f"Are you sure you want to delete these {len(job_ids)} jobs?"

        # Display the message box and check the response in one go
        msg_box = QMessageBox(QMessageBox.Icon.Warning, "Delete Job", message,
                              QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, self)

        if msg_box.exec() == QMessageBox.StandardButton.Yes:
            for job_id in job_ids:
                self.current_server_proxy.delete_job(job_id, confirm=True)
            self.fetch_jobs(clear_table=True)

    def download_files(self, event):
        pass

    def open_files(self, event):
        job_ids = self.selected_job_ids()
        if not job_ids:
            return

        for job_id in job_ids:
            job_info = self.current_server_proxy.get_job_info(job_id)
            path = os.path.dirname(job_info['output_path'])

            if sys.platform.startswith('darwin'):
                subprocess.run(['open', path])
            elif sys.platform.startswith('win32'):
                os.startfile(path)
            elif sys.platform.startswith('linux'):
                subprocess.run(['xdg-open', path])
            else:
                raise OSError("Unsupported operating system")

    def new_job(self) -> None:

        file_name, _ = QFileDialog.getOpenFileName(self, "Select Scene File")
        if file_name:
            self.new_job_window = NewRenderJobForm(file_name)
            self.new_job_window.show()