vse_toolbox/operators/operators.py

# SPDX-License-Identifier: GPL-2.0-or-later

import bpy
import importlib
import json
import re
import vse_toolbox

from bpy_extras.io_utils import ImportHelper 
from bpy.props import (
    CollectionProperty,
    BoolProperty,
    EnumProperty,
    IntProperty,
    StringProperty,
)
from bpy.types import (
    Operator,
)
from pathlib import Path
from vse_toolbox.constants import (
    ASSET_PREVIEWS,
    CASTING_BUFFER,
    CONFIG_DIR,
    EDITS,
    EDIT_SUFFIXES,
    MOVIES,
    MOVIE_SUFFIXES,
    PREVIEWS_DIR,
    SOUNDS,
    SOUND_SUFFIXES,
)
from vse_toolbox.sequencer_utils import (
    clean_sequencer,
    get_strips,
    set_active_strip,
    import_edit,
    import_movie,
    import_sound,
    rename_strips,
    render_strips,
    set_channels,
    get_channel_index,
    new_text_strip
)
from vse_toolbox.bl_utils import get_addon_prefs, get_scene_settings, get_strip_settings
from vse_toolbox.file_utils import install_module, norm_name, norm_str


class VSETB_OT_tracker_connect(Operator):
    bl_idname = "vse_toolbox.tracker_connect"
    bl_label = "Connect to Tracker"
    bl_description = "Connect to Tracker"

    @classmethod
    def poll(cls, context):
        prefs = get_addon_prefs()
        if prefs.tracker:
        
            return True

    def execute(self, context):
        prefs = get_addon_prefs()
        settings = get_scene_settings()
        try:
            prefs.tracker.connect()
            self.report({'INFO'}, f'Sucessfully login to {settings.tracker_name.title()}')
            return {"FINISHED"}
        except Exception as e:
            print('e: ', e)
            self.report({'ERROR'}, f'Cannot connect to tracker.')
            return {"CANCELLED"}


class VSETB_OT_export_csv(Operator):
    bl_idname = "vse_toolbox.export_csv"
    bl_label = "Set Scene"
    bl_description = "Set Scene for Breakdown"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        return True

    def execute(self, context):
        self.report({'ERROR'}, f'Export not implemented yet.')

        return {"CANCELLED"}


class VSETB_OT_auto_select_files(Operator):
    bl_idname = "import.auto_select_files"
    bl_label = "Auto Select"
    bl_description = "Auto Select Files"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        return True

    def get_items(self, items=[]):
        if not items:
            return [('NONE', 'None', '', 0)]
        return [(e, e, '', i) for i, e in enumerate(sorted(items))]

    def execute(self, context):
        params = context.space_data.params
        directory = Path(params.directory.decode())
        
        EDITS.clear()
        MOVIES.clear()
        SOUNDS.clear()
        
        edits = []
        movies = []
        sounds = []

        for file_entry in directory.glob('*'):
            if file_entry.is_dir():
                continue

            if file_entry.suffix in EDIT_SUFFIXES:
                edits.append(file_entry.name)
            elif file_entry.suffix in MOVIE_SUFFIXES:
                movies.append(file_entry.name)
            elif file_entry.suffix in SOUND_SUFFIXES:
                sounds.append(file_entry.name)

        EDITS.extend(self.get_items(items=edits))
        MOVIES.extend(self.get_items(items=movies))
        SOUNDS.extend(self.get_items(items=sounds))

        return {'FINISHED'}


class VSETB_OT_import_files(Operator):
    bl_idname = "vse_toolbox.import_files"
    bl_label = "Import"
    bl_description = "Import Edit"
    bl_options = {"REGISTER", "UNDO"}
    
    directory : StringProperty(subtype='DIR_PATH')
    filepath: StringProperty(
        name="File Path",
        description="Filepath used for importing the file",
        maxlen=1024,
        subtype='FILE_PATH',
    )
    files : CollectionProperty(type=bpy.types.OperatorFileListElement)
    clean_sequencer : BoolProperty(
        name="Clean Sequencer",
        default=False,
        description="Clean all existing strips in sequencer",
    )
    
    import_edit : BoolProperty(name='', default=True)
    edit: EnumProperty(name='', items=lambda s, c: EDITS)

    import_movie : BoolProperty(name='', default=False)
    movie: EnumProperty(name='', items=lambda s, c: MOVIES)

    import_sound : BoolProperty(name='', default=False)
    sound: EnumProperty(name='', items=lambda s, c: SOUNDS)

    @classmethod
    def poll(cls, context):
        return True

    def draw(self, context):
        scn = context.scene
        settings = get_scene_settings()

        layout = self.layout
        layout.use_property_split = True 
        layout.use_property_decorate = False

        col = layout.column(align=True)
        col.operator('import.auto_select_files', text='Auto Select')

        row = self.layout.row(heading="Import Edit", align=True)
        row.prop(self, 'import_edit')
        sub = row.row(align=True)
        sub.active = self.import_edit
        sub.prop(self, 'edit')

        row = self.layout.row(heading="Import Movie", align=True)
        row.prop(self, 'import_movie')
        sub = row.row()
        sub.active = self.import_movie
        sub.prop(self, 'movie')

        row = self.layout.row(heading="Import Sound", align=True)
        row.prop(self, 'import_sound')
        sub = row.row()
        sub.active = self.import_sound
        sub.prop(self, 'sound')

        col = layout.column()
        col.separator()
        col.prop(self, 'clean_sequencer')

    def invoke(self, context, event):
        context.window_manager.fileselect_add(self)
        return {'RUNNING_MODAL'}

    def execute(self, context):
        otio = install_module('opentimelineio')

        edit_filepath = Path(self.directory, self.edit)
        if not edit_filepath.exists():
            self.import_edit = False

        movie_filepath = Path(self.directory, self.movie)
        if not movie_filepath.exists():
            self.import_movie = False
        
        sound_filepath = Path(self.directory, self.sound)
        if not sound_filepath.exists():
            self.import_sound = False
        
        if self.clean_sequencer:
            clean_sequencer(
                edit=self.import_edit,
                movie=self.import_movie,
                sound=self.import_sound,
            )
        
        if self.import_edit:
            print(f'[>.] Loading Edit from: {str(edit_filepath)}')
            import_edit(edit_filepath, adapter="cmx_3600")
        
        if self.import_movie:
            print(f'[>.] Loading Movie from: {str(movie_filepath)}')
            import_movie(movie_filepath)

        if self.import_sound:
            print(f'[>.] Loading Audio from: {str(sound_filepath)}')
            import_sound(sound_filepath)
        elif not self.import_sound and self.import_movie:
            print(f'[>.] Loading Audio from Movie: {str(movie_filepath)}')
            import_sound(movie_filepath)
        
        context.scene.sequence_editor.sequences.update()

        return {"FINISHED"}


class VSETB_OT_load_assets(Operator):
    bl_idname = "vse_toolbox.load_assets"
    bl_label = "Load Assets for current projects"
    bl_description = "Load Assets"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        settings = get_scene_settings()
        if settings.active_project:
            return True

    def execute(self, context):
        settings = get_scene_settings()
        prefs = get_addon_prefs()
        tracker = prefs.tracker
        
        tracker.connect()
        project = settings.active_project
        project.assets.clear()
        
        assets = tracker.get_assets(project['id'])
        if not assets:
            self.report({'ERROR'}, f'No Assets found for {project.name}.')

        for asset_data in assets:
            asset = project.assets.add()

            asset.name = asset_data['id']
            asset.norm_name = norm_name(asset_data['name'], separator=' ', format=str.lower)

            asset.tracker_name = asset_data['name']
            asset.id = asset_data['id']
            asset.asset_type = asset_data['asset_type']
            
            datas = asset_data.get('data', {})
            for key, values in datas.items():
                asset[key] = values

            preview_id = asset_data.get('preview_file_id')
            if preview_id:
                asset.preview = preview_id
                preview_path = Path(PREVIEWS_DIR / project.id / preview_id).with_suffix('.png')
                tracker.download_preview(preview_id, preview_path)
                
                if preview_id not in ASSET_PREVIEWS:
                    ASSET_PREVIEWS.load(preview_id, preview_path.as_posix(), 'IMAGE', True)


        self.report({'INFO'}, f'Assets for {project.name} successfully loaded')
        return {"FINISHED"}


class VSETB_OT_load_projects(Operator):
    bl_idname = "vse_toolbox.load_projects"
    bl_label = "Load Projects"
    bl_description = "Load Projects"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        return True

    def execute(self, context):
        settings = get_scene_settings()
        prefs = get_addon_prefs()
        tracker = prefs.tracker
        
        settings.projects.clear()
        tracker.connect()

        for project_data in tracker.get_projects():


            project = settings.projects.add()
            project.type = project_data['production_type'].upper().replace(' ', '')
            project.name = project_data['name']
            project.id = project_data['id']

            if project.type == 'TVSHOW':
                for episode_data in tracker.get_episodes(project_data):
                    episode = project.episodes.add()
                    episode.name = episode_data['name']
                    episode.id = episode_data['id']

            for metadata_data in tracker.get_shots_metadata(project_data):
                #print(metadata_data)
                metadata_type = project.metadata_types.add()
                metadata_type.name = metadata_data['field_name']
                metadata_type['choices'] = metadata_data['choices']

            for asset_type_data in tracker.get_asset_types(project_data):
               asset_type = project.asset_types.add()
               asset_type.name = asset_type_data['name']
        
        if settings.active_project:
            settings.active_project.set_strip_metadata()
        
        bpy.ops.vse_toolbox.load_assets()

        return {'FINISHED'}


class VSETB_OT_new_episode(Operator):
    bl_idname = "vse_toolbox.new_episode"
    bl_label = "New Epispde"
    bl_description = "Add new Episode to Project"
    bl_options = {"REGISTER", "UNDO"}

    episode_name : StringProperty(name="Episode Name", default="")

    @classmethod
    def poll(cls, context):
        return True

    def invoke(self, context, event):
        scn = context.scene
        settings = get_scene_settings()
        
        return context.window_manager.invoke_props_dialog(self)

    def execute(self, context):
        settings = get_scene_settings()
        prefs = get_addon_prefs()
        tracker = prefs.tracker

        episode_name = settings.episode_template.format(index=int(self.episode_name))
        
        print(self.episode_name)
        print('episode_name: ', episode_name)

        episode = tracker.get_episode(episode_name)
        if episode:
            self.report({'ERROR'}, f'Episode {episode_name} already exists')
            return {"CANCELLED"}
        
        tracker.new_episode(episode_name)
        # tracker.get_episodes
        tracker.update_project()
        self.report({'INFO'}, f'Episode {episode_name} successfully created')
        return {'FINISHED'}  
        
        
class VSETB_OT_reload_addon(Operator):
    bl_idname = "vse_toolbox.reload_addon"
    bl_options = {"UNDO"}
    bl_label = 'Reload VSE ToolBox Addon'
    bl_description = 'Reload The VSE ToolBox Addon'

    def execute(self, context):

        print('Execute reload')

        vse_toolbox.unregister()
        importlib.reload(vse_toolbox)
        vse_toolbox.register()

        return {'FINISHED'}


class VSETB_OT_rename(Operator):
    bl_idname = "vse_toolbox.strips_rename"
    bl_label = "Rename Strips"
    bl_description = "Rename Strips"
    bl_options = {"REGISTER", "UNDO"}

    template : StringProperty(name="Strip Template Name", default="")
    increment : IntProperty(name="Name Increment", default=0)
    channel_name : StringProperty(name="Channel Name", default="")
    #selected_only : BoolProperty(name="Selected Only", default=False)
    start_number : IntProperty(name="Start Number", default=0, min=0)
    by_sequence : BoolProperty(
        name="Reset By Sequence",
        description="Reset Start Number for each sequence",
        default=False
    )

    @classmethod
    def poll(cls, context):
        return True

    def invoke(self, context, event):
        scn = context.scene
        settings = get_scene_settings()
        
        return context.window_manager.invoke_props_dialog(self)

    def draw(self, context):
        scn = context.scene
        settings = get_scene_settings()

        layout = self.layout
        col = layout.column()
        col.use_property_split = True 
        col.prop(self, 'template')
        col.prop(self, 'start_number')
        if self.channel_name == 'Shots':
            col.prop(self, 'by_sequence')
        #col.prop(self, 'selected_only')

    def execute(self, context):
        scn = context.scene   
        #strips = get_strips(channel=self.channel_name, selected_only=self.selected_only)
        strips = get_strips(channel=self.channel_name, selected_only=True)
        
        rename_strips(
            strips, self.template,
            increment=self.increment, start_number=self.start_number,
            by_sequence=self.by_sequence
        )

        return {"FINISHED"}
        

class VSETB_OT_render(Operator):
    bl_idname = "vse_toolbox.strips_render"
    bl_label = "Render Shots Strips"
    bl_description = "Render Shots Strips"
    bl_options = {"REGISTER", "UNDO"}

    selected_only : BoolProperty(name="Selected Only", default=False)

    @classmethod
    def poll(cls, context):
        return True

    def invoke(self, context, event):
        scn = context.scene
        settings = get_scene_settings()
        
        return context.window_manager.invoke_props_dialog(self)

    def draw(self, context):
        scn = context.scene
        settings = get_scene_settings()

        layout = self.layout
        col = layout.column()
        col.use_property_split = True
        col.use_property_decorate = False
        col.prop(settings, 'channel', text='Channel')
        col.prop(self, 'selected_only')

    def execute(self, context):
        scn = context.scene
        settings = get_scene_settings()
        strips = get_strips(channel=settings.channel, selected_only=self.selected_only)
        
        render_strips(strips)

        return {"FINISHED"}


class VSETB_OT_set_sequencer(Operator):
    bl_idname = "vse_toolbox.set_sequencer"
    bl_label = "Set Sequencer"
    bl_description = "Set resolution, frame end and channel names"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        return True

    def execute(self, context):
        scn = context.scene   

        set_channels()
        movie = get_strips(channel='Movie')
        if movie:
            movie = movie[0]
            movie.transform.scale_x = movie.transform.scale_y = 1
            elem = movie.strip_elem_from_frame(scn.frame_current)
            scn.render.resolution_x = elem.orig_width
            scn.render.resolution_y = elem.orig_height
        else:
            self.report({'INFO'}, f'Cannot set Resolution. No Movie Found.')
        
        return {"FINISHED"}


def get_asset_items(self, context):
    settings = get_scene_settings()
    project = settings.active_project

    return [(a.id, a.label, '', i) for i, a in enumerate(project.assets)]


class VSETB_OT_casting_add(Operator):
    bl_idname = "vse_toolbox.casting_add"
    bl_label = "Casting Add"
    bl_description = "Add Asset to Castin"
    bl_options = {"REGISTER", "UNDO"}
    bl_property = "asset_name"

    asset_name : EnumProperty(name='', items=get_asset_items)
    @classmethod
    def poll(cls, context):
        active_strip = context.scene.sequence_editor.active_strip
        if active_strip:
            return True

    def invoke(self, context, event):
        context.window_manager.invoke_search_popup(self)
        return {'FINISHED'}

    def execute(self, context):
        scn = context.scene
        active_strip = scn.sequence_editor.active_strip

        settings = get_scene_settings()
        strip_settings = get_strip_settings()

        project = settings.active_project

        if strip_settings.casting.get(self.asset_name):
            asset = project.assets[self.asset_name]
            self.report({'WARNING'}, f"Asset {asset.label} already casted.") 
            return {"CANCELLED"}

        item = strip_settings.casting.add()
        asset = project.assets[self.asset_name]

        item.name = asset.name 
        item.id = project.assets[self.asset_name].id

        item['_name'] = asset.label

        strip_settings.casting.update()
        
        return {"FINISHED"}


class VSETB_OT_casting_remove(Operator):
    bl_idname = "vse_toolbox.casting_remove"
    bl_label = "Remove Item from Casting"
    bl_description = "Remove Item from Casting"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        active_strip = context.scene.sequence_editor.active_strip
        if active_strip:
            return True

    def invoke(self, context, event):
        scn = context.scene
        
        strip_settings = get_strip_settings()
        idx = strip_settings.casting_index

        try:
            item = strip_settings.casting[idx]
        except IndexError:
            pass
        else:
            item = strip_settings.casting[strip_settings.casting_index]
            label = item.asset.label if item.asset.label else 'Empty'
            info = f"Item {label} removed from casting"
            strip_settings.casting.remove(idx)

            if strip_settings.casting_index == 0:
                strip_settings.casting_index = 0
            else:
                strip_settings.casting_index -= 1
                
            self.report({'INFO'}, info)

        return {"FINISHED"}


class VSETB_OT_casting_move(Operator):
    bl_idname = "vse_toolbox.casting_move"
    bl_label = "Casting Actions"
    bl_description = "Actions to Add, Remove, Move casting items"
    bl_options = {"REGISTER", "UNDO"}

    action: EnumProperty(
        items=(
            ('UP', "Up", ""),
            ('DOWN', "Down", ""),
        )
    )

    asset_name : StringProperty()

    @classmethod
    def poll(cls, context):
        active_strip = context.scene.sequence_editor.active_strip
        if active_strip:
            return True

    def invoke(self, context, event):
        scn = context.scene

        strip_settings = get_strip_settings()
        idx = strip_settings.casting_index

        try:
            item = strip_settings.casting[idx]
        except IndexError:
            pass
        else:
            if self.action == 'DOWN' and idx < len(strip_settings.casting) - 1:
                item_next = strip_settings.casting[idx+1].name
                strip_settings.casting.move(idx, idx+1)
                strip_settings.casting_index += 1

            elif self.action == 'UP' and idx >= 1:
                item_prev = strip_settings.casting[idx-1].name
                strip_settings.casting.move(idx, idx-1)
                strip_settings.casting_index -= 1

            info = f"Item {item.asset.label} moved to position {(item.asset.label, strip_settings.casting_index + 1)}"
            self.report({'INFO'}, info)

        return {"FINISHED"}


class VSETB_OT_copy_casting(Operator):
    bl_idname = "vse_toolbox.copy_casting"
    bl_label = "Casting Actions"
    bl_description = "Copy Casting from active strip"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        active_strip = context.scene.sequence_editor.active_strip
        strip_settings = get_strip_settings()
        
        if active_strip and strip_settings.casting:
            return True

    def execute(self, context):
        scn = context.scene
        
        active_strip = scn.sequence_editor.active_strip
        strip_settings = get_strip_settings()
        
        datas = [c.to_dict() for c in strip_settings.casting]
        CASTING_BUFFER.write_text(json.dumps(datas), encoding='utf-8')
        
        return {"FINISHED"}


class VSETB_OT_paste_casting(Operator):
    bl_idname = "vse_toolbox.paste_casting"
    bl_label = "Casting Actions"
    bl_description = "Copy Casting from active strip"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        active_strip = context.scene.sequence_editor.active_strip
        if active_strip:
            return True

    def execute(self, context):
        scn = context.scene
        strip_settings = get_strip_settings()

        if not CASTING_BUFFER.exists():
            self.report({'ERROR'}, f'No Casting to copy.')
            return {"CANCELLED"}

        casting_datas = json.loads(CASTING_BUFFER.read_text())
        for casting_item in casting_datas:
            for strip in context.selected_sequences:
                strip_settings = strip.vsetb_strip_settings
    
                if strip_settings.casting.get(casting_item['name']):
                    continue

                item = strip.vsetb_strip_settings.casting.add()
                for k, v in casting_item.items():
                    setattr(item, k, v)

                strip_settings.casting.update()

        return {"FINISHED"}


class VSETB_OT_set_stamps(Operator):
    bl_idname = "vse_toolbox.set_stamps"
    bl_label = "Set Stamps"
    bl_description = "Set Stamps on Video"
    bl_options = {"REGISTER", "UNDO"}

    def execute(self, context):
        scn = context.scene
        settings = get_scene_settings()
        #strip_settings = get_strip_settings()
        channel_index = get_channel_index('Stamps')

        for strip in get_strips('Stamps'):
            if strip.type == 'META':
                scn.sequence_editor.sequences.remove(strip)

        bpy.ops.sequencer.select_all(action='DESELECT')

        # Project Name
        project_strip_stamp = new_text_strip(
            'project_name_stamp', channel=1, start=scn.frame_start, end=scn.frame_end,
            text=settings.active_project.name, font_size=24,
            x=0.01, y=0.01, align_x='LEFT', align_y='BOTTOM', select=True,
            box_color=(0, 0, 0, 0.5), box_margin=0.005,
        )

        # Shot Name
        shot_strip_stamp = new_text_strip(
            f'shot_name_stamp', channel=2, start=scn.frame_start, end=scn.frame_end,
            text='{active_shot_name}', font_size=24,
            y=0.01, align_y='BOTTOM', select=True,
            box_color=(0, 0, 0, 0.5), box_margin=0.005,
        )

        # Frame Range
        frame_strip_stamp = new_text_strip(
            'frame_range_stamp', channel=3, start=scn.frame_start, end=scn.frame_end,
            text='{active_shot_frame} / {active_shot_duration}', font_size=24,
            x=0.99, y=0.01, align_x='RIGHT', align_y='BOTTOM', select=True,
            box_color=(0, 0, 0, 0.5), box_margin=0.005,
        )

        # for shot_strip in get_strips('Shots'):            
        #     # Shot Name
        #     shot_strip_stamp = new_text_strip(
        #         f'{shot_strip.name}_stamp', channel=3, start=shot_strip.frame_final_start, end=shot_strip.frame_final_end,
        #         text=shot_strip.name, font_size=24,
        #         y=0.01, align_y='BOTTOM', select=True,
        #         box_color=(0, 0, 0, 0.35), box_margin=0.005,
        #     )

        bpy.ops.sequencer.meta_make()
        stamps_strip = context.active_sequence_strip
        stamps_strip.name = 'Stamps'
        stamps_strip.channel = channel_index

        #stamps_strip = scn.sequence_editor.sequences.new_meta('Stamps', scn.frame_start, scn.frame_end)
        #stamps_strip.channel = get_channel_index('Stamps')
        scn.frame_set(scn.frame_current)  # For update stamps

        return {"FINISHED"}


classes = (
    VSETB_OT_auto_select_files,
    VSETB_OT_casting_add,
    VSETB_OT_casting_remove,
    VSETB_OT_casting_move,
    VSETB_OT_copy_casting,
    VSETB_OT_paste_casting,
    VSETB_OT_export_csv,
    VSETB_OT_import_files,
    VSETB_OT_load_assets,
    VSETB_OT_load_projects,
    VSETB_OT_new_episode,
    VSETB_OT_reload_addon,
    VSETB_OT_rename,
    VSETB_OT_render,
    VSETB_OT_set_sequencer,
    VSETB_OT_tracker_connect,
    VSETB_OT_set_stamps
)

def register(): 
    for cls in classes:
        bpy.utils.register_class(cls)

def unregister():
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)