# SmedgeRender.py
#
# Smedge integration with Nuke
#
# Smedge
# Copyright Uberware. All Rights Reserved

"""Simple Smedge submit interface for Nuke.

It will use the Smedge Python API, which expects the Smedge "python" folder to be in your PYTHONPATH environment.

To integrate with the Nuke menus, use something like this in menu.py:
    import sys
    sys.path.append("/Applications/Smedge.app/Contents/MacOS/python")
    nuke.menu("Nuke").addCommand("Render/Submit to Smedge", "import smedge.dcc.nuke; win = smedge.dcc.nuke.SmedgeSubmitWindow()")
"""


from pathlib import Path
from typing import List
from uuid import UUID

import nuke
import smedge

try:
    from PySide6 import QtWidgets, QtCore
except ImportError:
    from PySide2 import QtWidgets, QtCore


class SmedgeSubmitWindow(smedge.SmedgeSubmitWindow):
    """A dialog box to override some details about how the job gets submitted."""

    def __init__(self):
        """Initialize a new SmedgeSubmitWindow object."""
        # Ensure the script has been saved to disk at some point
        root_name = nuke.root().name()
        if root_name == "Root":
            raise RuntimeError("You must save the script before it can be rendered on the farm")
        nuke.scriptSave()
        self.script_path = Path(root_name)
        # Set up the window
        self.mode = None
        first_frame = int(nuke.root().knob("first_frame").value())
        last_frame = int(nuke.root().knob("last_frame").value())
        super().__init__(frame_range=(first_frame, last_frame), packet_size=10)

    def add_properties(self, properties: smedge.PropertiesType):
        """Customize to add the Nuke specific properties."""
        # Add the mode to the list of job properties
        self.mode = QtWidgets.QComboBox()
        self.mode.addItems(["A single job", "One job for each Write node"])
        self.mode.setToolTip("Choose one job for the entire script or one job for each Write node")
        properties["Mode"] = self.mode
        # Use the script file name as the default job name
        self.name.setText(self.script_path.name)
        super().add_properties(properties)

    def submit(self, job: smedge.JobType) -> List[UUID]:
        """Submit button callback to do the work."""
        # Build the base job
        job["Type"] = "Nuke"
        job["Scene"] = self.script_path
        # Create the list of jobs to submit
        jobs = []
        if self.mode.currentIndex() == 0:
            # One job for the whole scene
            job["ImageFormat"] = ";".join(node["file"].value() for node in _get_write_nodes())
            jobs.append(job)
        else:
            # One job for each write node
            for node in _get_write_nodes():
                node_job = job.copy()
                node_job["Name"] += " :: " + node.name()
                node_job["WriteNode"] = node.name()
                node_job["ImageFormat"] = node["file"].value()
                jobs.append(node_job)
        # Send to the farm
        return smedge.submit_jobs(jobs, self.paused.isChecked())


def _get_write_nodes() -> List[nuke.Node]:
    """Get all write nodes in the script."""
    return [node for node in nuke.allNodes() if node.Class() == "Write"]
