# Localization Helper Functions
# These methods allow you to Localize Clips on Bins, Shots, Tracks and Sequences
import hiero.core
from PySide6 import (QtCore, QtGui, QtWidgets)

# This is just a convenience method for returning QActions with a title, triggered method and icon.
def createMenuAction(title, method, icon = None ):
    action = QtGui.QAction(title,None)
    action.setIcon(QtGui.QIcon(icon))
    action.triggered.connect( method )
    return action

def setlocalizationPolicyOnTrack(track, policy = hiero.core.Clip.kOnLocalize ):
    """Localizes all TrackItems in a Track
    setlocalizationPolicyOnTrack( track, policy = hiero.core.Clip.kOnLocalize ):
    @param: track - a list of hiero.core.Clip objects
    @param: policy - the localization policy (kOnLocalize,kAutoLocalize,kOffLocalize)"""
    if not isinstance(track,(hiero.core.VideoTrack,hiero.core.AudioTrack)):
        hiero.core.log.info('First argument must be a hiero.core.VideoTrack/hiero.core.AudioTrack')
        return
    else:
        for ti in track:
            ti.source().setLocalizationPolicy( policy )


def setlocalizationPolicyOnTrackItem(trackItem, policy = hiero.core.Clip.kOnLocalize ):
    """Localizes all TrackItems in a Track
    setlocalizationPolicyOnTrackItem( track, policy = hiero.core.Clip.kOnLocalize ):
    @param: track - a list of hiero.core.Clip objects
    @param: policy - the localization policy (kOnLocalize,kAutoLocalize,kOffLocalize)"""
    if not isinstance(trackItem,(hiero.core.TrackItem)):
        hiero.core.log.info('First argument must be a hiero.core.TrackItem')
        return
    else:
        trackItem.source().setLocalizationPolicy( policy )

def setlocalizationPolicyOnSequence(sequence, policy = hiero.core.Clip.kOnLocalize ):
    """Localizes all Clips used in a Sequence
    @param: clipList - a list of hiero.core.Clip objects
    @param: policy - the localization policy (kOnLocalize,kAutoLocalize,kOffLocalize)"""
    if not isinstance(sequence,(hiero.core.Sequence)):
        hiero.core.log.info('First argument must be a hiero.core.Sequence')
        return
    else:
        for track in sequence.items():
            for ti in track:
                if isinstance(ti,(hiero.core.TrackItem)):
                    clip = ti.source()
                    if clip.localizationPolicy != policy:
                        clip.setLocalizationPolicy(policy)

def setlocalizationPolicyOnBin(bin, policy = hiero.core.Clip.kOnLocalize, recursive = True ):
    """Localizes all Clips, recursively found in all Sub-Bins
    @param: clipList - a list of hiero.core.Clip objects
    @param: policy - the localization policy (kOnLocalize,kAutoLocalize,kOffLocalize)"""

    clips =[]

    if not isinstance(bin,(hiero.core.Bin)):
        hiero.core.log.info('First argument must be a hiero.core.Bin')
        return
    else:
        if recursive:
            # Here we're going to walk the entire Bin view, for all Clips contained in this Bin
            clips += hiero.core.findItemsInBin(bin, filter = hiero.core.Clip)
        else:
            # Here we just get Clips at the 1st level of the Bin
            clips = bin.clips()

    for clip in clips:
        if clip.localizationPolicy != policy:
            clip.setLocalizationPolicy(policy)


def setlocalizationPolicyOnClipWithExtension(clip, extensionList = [None], policy = hiero.core.Clip.kOnLocalize ):
    """Localizes Clips in clipList with a given file extension in extensionList
    localizeClipsWithExtension(clipList, extensionList = None, policy = hiero.core.Clip.kOnLocalize ):
    @param: clipList - a list of hiero.core.Clip objects
    @param: extensionList - a list of file extensions, e.g. ['dpx','exr','mov']
    @param: policy - the localization policy (kOnLocalize,kAutoLocalize,kOnDemandLocalize,kOffLocalize)"""

    if type(extensionList) != list:
        print("Please supply a list of file extensions, e.g. extensionList=['dpx','exr','mov']")
        return
    else:
        clipPath = clip.mediaSource().filename()
        hiero.core.log.info('clipPath is'+str(clipPath))
        if clipPath.lower().endswith( tuple(extensionList) ):
            clip.setLocalizationPolicy( policy )

class CustomLocalizationDialog(QtWidgets.QDialog):

    def __init__(self,itemSelection=None,parent=None):
        super(CustomLocalizationDialog, self).__init__(parent)
        self.setWindowTitle("Select File Extensions to Localize")
        self.setWindowIcon(QtGui.QIcon("icons:BlendingDisabled.png"))
        self.setSizePolicy( QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Fixed )

        layout = QtWidgets.QFormLayout()
        self._itemSelection = itemSelection

        self._formatFilterEdit = QtWidgets.QLineEdit()
        self._formatFilterEdit.setToolTip('Enter a comma separated list of file extensions to localize, e.g. dpx, exr')
        self._formatFilterEdit.setText('dpx')
        layout.addRow("Extension: ",self._formatFilterEdit)

        # Standard buttons for Add/Cancel
        self._buttonbox = QtWidgets.QDialogButtonBox(QtWidgets.QDialogButtonBox.Ok | QtWidgets.QDialogButtonBox.Cancel)
        self._buttonbox.button(QtWidgets.QDialogButtonBox.Ok).setText("Add")
        self._buttonbox.accepted.connect(self.accept)
        self._buttonbox.rejected.connect(self.reject)
        layout.addRow("",self._buttonbox)

        self.setLayout(layout)

    def convertFormatStringToFormatList(self):
        formatString = self._formatFilterEdit.text()
        formats = formatString.split(',')
        return formats


    # Set Tag dialog with Note and Selection info. Adds Tag to entire object, not over a range of frames
    def showDialogAndLocalizeSelection(self):

        if len(self._itemSelection)<1:
            hiero.core.log.info('No valid selected Items can be localized')
            return

        clips = [item for item in self._itemSelection if isinstance(item, hiero.core.Clip)]

        if len(clips)==0:
            return

        if self.exec_():
            exts = self.convertFormatStringToFormatList()
            hiero.core.log.info('Got this extension list:'+str(exts))
            for clip in clips:
                setlocalizationPolicyOnClipWithExtension(clip, extensionList = exts)
            return



# Menu which adds a Tags Menu to the Viewer, Project Bin and Timeline/Spreadsheet
class CustomLocalizeMenu:

    def __init__(self):
        self._localizeMenu = None
        self._customLocalizeDialog = None

        # These actions in the viewer are a special case, because they drill down in to what is currrenly being
        self._localizeBinContents = createMenuAction("This Bin", self.localizeBinSelection,icon="icons:Bin.png")
        self._localizeThisBinSequence = createMenuAction("This Sequence", self.localizeBinSelection,icon="icons:TimelineStroked.png")
        self._localizeBinContentsCustom = createMenuAction("This Bin Custom...", self.localizeBinSelectionCustom,icon="icons:BlendingDisabled.png")
        self._localizeThisSequence  = createMenuAction("This Sequence", self.localizeTimelineSpreadsheetSequence,icon="icons:TimelineStroked.png")
        self._localizeThisTrack  = createMenuAction("This Track", self.localizeTimelineSpreadsheetSelection,icon="icons:Tracks.png")
        self._localizeThisShot  = createMenuAction("This Shot", self.localizeTimelineSpreadsheetSelection,icon="icons:Mask.png")

        hiero.core.events.registerInterest("kShowContextMenu/kBin", self.binViewEventHandler)
        hiero.core.events.registerInterest("kShowContextMenu/kTimeline", self.timelineSpreadsheetEventHandler)
        hiero.core.events.registerInterest("kShowContextMenu/kSpreadsheet", self.timelineSpreadsheetEventHandler)

    def createLocalizeMenuForView(self, viewName):
        localizeMenu = QtWidgets.QMenu("Custom Localize")
        if viewName == 'kBin':
            localizeMenu.addAction(self._localizeBinContents)
            localizeMenu.addAction(self._localizeBinContentsCustom)
            localizeMenu.addAction(self._localizeThisBinSequence)

        elif viewName == 'kTimeline/kSpreadsheet':
            localizeMenu.addAction(self._localizeThisSequence)
            localizeMenu.addAction(self._localizeThisTrack)
            localizeMenu.addAction(self._localizeThisShot)

        return localizeMenu

    # Called from the Bin View, to a BinItem selection.
    # Note: A mixed selection of Bins AND Sequences is not currently supported.
    def localizeBinSelection(self):
        if self._selection == None:
            return
        elif isinstance(self._selection[0],hiero.core.Bin):
            for bin in self._selection:
                setlocalizationPolicyOnBin(bin)

        elif isinstance(self._selection[0],hiero.core.Sequence):
            for seq in self._selection:
                setlocalizationPolicyOnSequence(seq)

    # Called from the Bin View, to localize all Clips with a custom Extension
    def localizeBinSelectionCustom(self):
        if self._selection == None:
            return
        elif isinstance(self._selection[0],hiero.core.Bin):
            clips = []
            for bin in self._selection:
                clips += hiero.core.findItemsInBin(bin, filter = hiero.core.Clip)
            # Present Custom Localize Dialog:
            if len(clips)>=1:
                CustomLocalizationDialog(itemSelection=clips).showDialogAndLocalizeSelection()

    # Called from the Spreadsheet or Timeline View, to localize all Clips in the Shot or Track
    def localizeTimelineSpreadsheetSelection(self):
        if self._selection == None:
            return

        elif isinstance(self._selection[0],hiero.core.TrackItem):
            for ti in self._selection:
                setlocalizationPolicyOnTrackItem(ti)

        elif isinstance(self._selection[0],(hiero.core.VideoTrack,hiero.core.AudioTrack)):
            for track in self._selection:
                setlocalizationPolicyOnTrack(track)

    # Called from the Spreadsheet or Timeline View, to localize all Clips in the current Sequence
    def localizeTimelineSpreadsheetSequence(self):
        if self._selection == None:
            return

        # We should only have one Sequence to localize. Take just the first item
        sequenceItem = self._selection[0]

        if isinstance(sequenceItem,(hiero.core.VideoTrack,hiero.core.AudioTrack)):
            sequence = sequenceItem.parent()
        elif isinstance(sequenceItem,hiero.core.TrackItem):
            sequence = sequenceItem.parent().parent()

        if isinstance(sequence,hiero.core.Sequence):
            setlocalizationPolicyOnSequence(sequence)

    # This handles events from the Project Bin View
    def binViewEventHandler(self,event):

        if not hasattr(event.sender, 'selection'):
            # Something has gone wrong, we should only be here if raised
            # by the Bin view which gives a selection.
            return
        s = event.sender.selection()

        # Return if there's no Selection. We won't add the Localize Menu.
        if s == None:
            return

        # Filter the selection to only act on Bins
        binSelection = [item for item in s if isinstance(item, (hiero.core.Bin))]
        binItemSelection = [item for item in s if isinstance(item,hiero.core.BinItem)]
        sequenceSelection = [item.activeItem() for item in binItemSelection if isinstance(item.activeItem(),hiero.core.Sequence)]
        self._selection = binSelection+sequenceSelection
        hiero.core.log.info('selection is'+ str(self._selection))

        binsSelected = len(binSelection)>=1
        sequencesSelected = len(sequenceSelection)>=1

        if binsSelected and sequencesSelected:
            hiero.core.log.debug('Selection must either be a Bin selection or Sequence Selection')
            return

        # Only add the Menu if Bins or Sequences are selected (this ensures menu isn't added in the Tags Pane)
        if len(self._selection) > 0:
            self._localizeMenu = self.createLocalizeMenuForView('kBin')
            self._localizeBinContents.setEnabled(binsSelected)
            self._localizeBinContentsCustom.setEnabled(binsSelected)
            self._localizeThisBinSequence.setEnabled(sequencesSelected)

            # Insert the Tags menu with the Localization Menu
            hiero.ui.insertMenuAction(self._localizeMenu.menuAction(), event.menu)
        return


    # This handles events from the Project Timeline View
    def timelineSpreadsheetEventHandler(self,event):

        if not hasattr(event.sender, 'selection'):
            # Something has gone wrong, we should only be here if raised
            # by the timeline view which gives a selection.
            return
        s = event.sender.selection()

        # Return if there's no Selection. We won't add the Tags Menu.
        if s == None:
            return

        # Filter the selection to only act on TrackItems, not Transitions etc.
        shotSelection = [item for item in s if isinstance(item, (hiero.core.TrackItem))]
        trackSelection = [item for item in s if isinstance(item, (hiero.core.VideoTrack,hiero.core.AudioTrack))]

        if len(shotSelection)>=1 and len(trackSelection)>=1:
            hiero.core.log.debug('Selection must either be a Shot selection or a Track Selection')
            return

        # We don't currently get a mixture of TrackItem and Tracks but combine the two lists to make a selection
        self._selection = shotSelection+trackSelection

        if len(self._selection) > 0:
            self._localizeMenu = self.createLocalizeMenuForView('kTimeline/kSpreadsheet')
            self._localizeThisShot.setEnabled(len(shotSelection)>=1)
            self._localizeThisTrack.setEnabled(len(trackSelection)>=1)

            # Insert the Tags menu with the Clear Tags Option
            hiero.ui.insertMenuAction(self._localizeMenu.menuAction(),  event.menu)
        return

# Instantiate the Menu to get it to register itself.
customLocalizeMenu = CustomLocalizeMenu()