vse_toolbox/operators/imports.py

from pathlib import Path
from os.path import expandvars
import re
import glob
from tempfile import gettempdir

import bpy
from bpy.types import Operator, UIList
from bpy.props import (CollectionProperty, BoolProperty, EnumProperty, StringProperty, IntProperty)

from vse_toolbox.constants import (EDITS, EDIT_SUFFIXES, MOVIES, MOVIE_SUFFIXES,
    SOUNDS, SOUND_SUFFIXES)

from vse_toolbox.sequencer_utils import (clean_sequencer, import_edit, import_movie,
    import_sound, get_strips, get_channel_index, get_empty_channel, scale_clip_to_fit)
       
from vse_toolbox.bl_utils import (get_scene_settings, get_addon_prefs, get_scene_settings, abspath)
from vse_toolbox.file_utils import install_module, parse, find_last, expand


class VSETB_OT_auto_select_files(Operator):
    bl_idname = "vse_toolbox.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(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.sort(reverse=True)
        movies.sort(reverse=True)
        sounds.sort(reverse=True)

        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)
    match_by : EnumProperty(name='Match By', items=[('NAME', 'Name', ''), ('INDEX', 'Index', '')])

    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('vse_toolbox.auto_select_files', text='Auto Select')

        row = 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 = layout.row(align=True)
        row.prop(self, 'match_by', expand=True)

        layout.separator()

        row = 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 = 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):
        sequencer = context.scene.sequence_editor.sequences

        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", match_by=self.match_by)
        
        if self.import_movie:
            print(f'[>.] Loading Movie from: {str(movie_filepath)}')

            for strip in get_strips(channel='Movie'):
                sequencer.remove(strip)

            import_movie(movie_filepath)

        if self.import_sound or (not self.import_sound and self.import_movie):

            for strip in get_strips(channel='Audio'):
                sequencer.remove(strip)

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

        return {"FINISHED"}

class VSETB_UL_import_task(UIList):
    def draw_item(self, context, layout, data, item, icon, active_data,
                  active_propname, index):

        layout.separator(factor=0.5)
        layout.prop(item, 'import_enabled', text='')
        layout.label(text=item.name)
        

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

    return [(t, t, '') for t in project.task_types.keys()]
class VSETB_OT_select_task(Operator):
    bl_idname = "vse_toolbox.select_task"
    bl_label = "Select Task"
    bl_description = "Select Task"
    bl_options = {"REGISTER", "UNDO"}

    bl_property = "task"

    task: EnumProperty(name="My Search", items=get_task_items)

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

    def execute(self, context):
        scn = context.scene
        settings = get_scene_settings()
        project = settings.active_project

        project.task_types[self.task].import_enabled = True

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


class VSETB_OT_unselect_task(Operator):
    bl_idname = "vse_toolbox.unselect_task"
    bl_label = "Unselect Task"
    bl_description = "Unselect Task"
    bl_options = {"REGISTER", "UNDO"}

    task : StringProperty()

    def execute(self, context):
        scn = context.scene
        settings = get_scene_settings()
        project = settings.active_project

        project.task_types[self.task].import_enabled = False

        return {'FINISHED'}


class VSETB_OT_add_import_template(Operator):
    bl_idname = "vse_toolbox.add_import_template"
    bl_label = "Add Template"
    bl_options = {"REGISTER", "UNDO"}

    def execute(self, context):
        scn = context.scene
        settings = get_scene_settings()
        project = settings.active_project

        project.import_shots.video_templates.add()

        return {'FINISHED'}
    

class VSETB_OT_remove_import_template(Operator):
    bl_idname = "vse_toolbox.remove_import_template"
    bl_label = "Remove Template"
    bl_options = {"REGISTER", "UNDO"}

    index : IntProperty()

    def execute(self, context):
        scn = context.scene
        settings = get_scene_settings()
        project = settings.active_project

        project.import_shots.video_templates.remove(self.index)

        return {'FINISHED'}


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

    return [(t, t, '') for t in project.sequences.keys()]
class VSETB_OT_select_sequence(Operator):
    bl_idname = "vse_toolbox.select_sequence"
    bl_label = "Select Sequence"
    bl_description = "Select Sequence"
    bl_options = {"REGISTER", "UNDO"}

    bl_property = "sequence"

    sequence: EnumProperty(name="My Search", items=get_sequence_items)

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

    def execute(self, context):
        scn = context.scene
        settings = get_scene_settings()
        project = settings.active_project

        project.sequences[self.sequence].import_enabled = True

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


class VSETB_OT_unselect_sequence(Operator):
    bl_idname = "vse_toolbox.unselect_sequence"
    bl_label = "Unselect sequence"
    bl_description = "Unselect sequence"
    bl_options = {"REGISTER", "UNDO"}

    sequence : StringProperty()

    def execute(self, context):
        scn = context.scene
        settings = get_scene_settings()
        project = settings.active_project

        project.sequences[self.sequence].import_enabled = False

        return {'FINISHED'}

class VSETB_OT_import_shots(Operator):
    bl_idname = "vse_toolbox.import_shots"
    bl_label = "Import Shots"
    bl_description = "Import Shots for disk or tracker"
    bl_options = {"REGISTER", "UNDO"}

    @classmethod
    def poll(cls, context):
        return get_scene_settings().active_project

    def set_sequencer_channels(self, task_types):
        scn = bpy.context.scene
        channels = ['Shots', 'Sequences', 'Stamps']
        #channel_index = len(task_types * 3) + len(channels) + 1

        channel_index = 1

        for task_type in task_types:
            audio_channel = scn.sequence_editor.channels[channel_index]
            audio_channel.name = f'{task_type} Audio'
            audio_channel.mute = (task_type != task_types[-1])

            video_channel = scn.sequence_editor.channels[channel_index+1]
            video_channel.name = f'{task_type} Video'
        
            channel_index += 3

        for channel in channels:
            scn.sequence_editor.channels[channel_index].name = channel
        
            channel_index += 1

    # def get_preview_dir(self):
    #     preview_dir = Path(bpy.app.tempdir, 'previews')
    #     if bpy.data.filepath:
    #         preview_dir = Path(bpy.data.filepath).parent / 'previews'
    #     return preview_dir

    def download_preview(self, task_type, shot):
        prefs = get_addon_prefs()
        tracker = prefs.tracker

        settings = get_scene_settings()
        project = settings.active_project

        task = tracker.get_task(task_type.id or task_type.name, entity=shot)
        last_comment = tracker.get_last_comment_with_preview(task)
        if not last_comment:
            return

        last_preview = last_comment['previews'][0]

        ext = last_preview['extension']
        shot_name = shot['name']
        sequence_name = f"{shot['sequence_name']}_"
        if shot_name.startswith(sequence_name):
            shot_name = shot_name.replace(sequence_name, '', 1)

        preview_dir = abspath(project.import_shots.previews_folder)#self.get_preview_dir()

        filepath = Path(preview_dir, f'{sequence_name}{shot_name}_{task_type.name}.{ext}')
        filepath.parent.mkdir(parents=True, exist_ok=True)

        tracker.download_preview_file(last_preview, str(filepath))

        return filepath

    def conform_render(self, clip):
        scn = bpy.context.scene
        scn.render.resolution_x = clip.elements[0].orig_width
        scn.render.resolution_y = clip.elements[0].orig_height
        scn.render.fps = int(clip.elements[0].orig_fps)

        scn.view_settings.view_transform = 'Standard'

    def find_shot_preview(self, sequence, shot_name, task_type):
        settings = get_scene_settings()
        project = settings.active_project
        import_shots = project.import_shots
        project_templates = {t.name: t.value for t in project.templates}

        for template in [import_shots.video_template] + list(import_shots.video_templates.keys()):
            
            template = expand(template, **project_templates)

            format_data = project.format_data
            format_data.update(parse(project.sequence_template, sequence.name))
            format_data.update(parse(project.shot_template, shot_name))
            # Normalize name in snake_case for now
            format_data.update(task=task_type.name.lower().replace(' ', '_'))

            if last_preview := find_last(template, **format_data):
                return last_preview

    def import_casting(self, shot, strip):
        prefs = get_addon_prefs()
        tracker = prefs.tracker
        strip_settings = strip.vsetb_strip_settings

        casting_data = tracker.get_casting(shot)

        for asset_data in casting_data:
            item = strip_settings.casting.add()
            item.name = asset_data['asset_name'] 
            item.id = asset_data['asset_id']
            item['_name'] = asset_data['asset_name']

        strip_settings.casting.update()

    def execute(self, context):
        scn = context.scene
        settings = get_scene_settings()
        project = settings.active_project
        import_shots = project.import_shots

        prefs = get_addon_prefs()
        tracker = prefs.tracker
        tracker.admin_connect()

        task_types = [t for t in project.task_types if t.import_enabled]
        sequences = [s for s in project.sequences if s.import_enabled]

        conformed = False

        scn.sequence_editor_clear()
        scn.sequence_editor_create()

        self.set_sequencer_channels([t.name for t in task_types])
        frame_index = 1
        for sequence in sequences:
            shots_data = tracker.get_shots(sequence=sequence.id)
            sequence_start = frame_index
            for shot_data in shots_data:
                frames = int(shot_data['nb_frames'])
                frame_end = frame_index + frames

                strip = scn.sequence_editor.sequences.new_effect(
                    name=shot_data['name'],
                    type='COLOR',
                    channel=get_channel_index('Shots'),
                    frame_start=frame_index,
                    frame_end=frame_index + frames
                )
                strip.blend_alpha = 0
                strip.color = (0.5, 0.5, 0.5)

                self.import_casting(shot_data, strip)

                for task_type in task_types:
                    if import_shots.import_source == 'DISK':
                        preview = self.find_shot_preview(sequence, shot_data['name'], task_type)          
                    else:
                        preview = self.download_preview(task_type, shot_data)
                        if not preview:
                            print(f'No preview found for shot {shot_data["name"]}')

                    if not preview:
                        continue
                    
                    print(f'Loading Preview from {preview}')
                    channel_index = get_channel_index(f'{task_type.name} Video')
                    video_clip = import_movie(preview, frame_start=frame_index, frame_end=frame_end)
                    video_clip.channel = channel_index

                    if video_clip.frame_offset_end:
                        video_clip.color_tag = 'COLOR_01'

                    if not conformed:
                        self.conform_render(video_clip)
                        conformed = True

                    scale_clip_to_fit(video_clip)

                    # Load Audio
                    channel_index = get_channel_index(f'{task_type.name} Audio')
                    audio_clip = import_sound(preview, frame_start=frame_index, 
                        frame_end=frame_end)
                    audio_clip.channel = channel_index
                    if video_clip.frame_offset_end:
                        audio_clip.color_tag = 'COLOR_01'
                    
                frame_index += frames

            strip = scn.sequence_editor.sequences.new_effect(
                name=sequence.name,
                type='COLOR',
                channel=get_channel_index('Sequences'),
                frame_start=sequence_start,
                frame_end=frame_index
            )
            strip.blend_alpha = 0
            strip.color = (0.25, 0.25, 0.25)

        scn.frame_start = 1
        scn.frame_end = frame_index -1

        #bpy.ops.vse_toolbox.set_stamps()

        return {'FINISHED'}

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

        layout = self.layout
        
        col = layout.column(align=False)
        col.use_property_split = True
        col.use_property_decorate = False

        row = col.row(align=True)
        row.prop(import_shots, "import_source", text='Videos Source', expand=True)

        if import_shots.import_source == 'DISK':
            #col.prop(import_shots, "sequence_dir_template", text='Sequence Dir')
            #col.prop(import_shots, "shot_folder_template", text='Shot Folder')
            row = col.row(align=True)
            row.prop(import_shots, "video_template", text='')
            row.operator("vse_toolbox.add_import_template", text='', icon='ADD')

            for i, template in enumerate(import_shots.video_templates):
                row = col.row(align=True)
                row.prop(template, "name", text='')
                row.operator("vse_toolbox.remove_import_template", text='', icon='REMOVE').index=i

            col.separator()
        else:
            col.prop(import_shots, "previews_folder", text='Previews Folder')
            col.separator()
            #if bpy.data.filepath:
                #col.label(text=f'  {self.get_preview_dir()}')
            #else:
            #    col.label(icon="ERROR", text="Save your Blender to keep the previews")

        row = col.row(align=True)
        row.prop(import_shots, 'import_task', text='Import Tasks')
        row.operator("vse_toolbox.select_task", text='', icon='ADD')

        tasks = [t for t in project.task_types if t.import_enabled]

        if import_shots.import_task == "FROM_LIST":
            if tasks:
                split = col.split(factor=0.4)
                split.row()
                box = split.box()
                box_col = box.column(align=True)

                for task_type in tasks:
                    row = box_col.row(align=True)
                    sub = row.row(align=True)
                    sub.enabled = False
                    sub.alignment = 'LEFT'
                    sub.scale_x = 0.15
                    sub.prop(task_type, 'color', text='')

                    op = row.operator("vse_toolbox.unselect_task", text=task_type.name)
                    op.task = task_type.name

                    row.separator(factor=0.5)
                    op = row.operator("vse_toolbox.unselect_task", text='', icon='REMOVE', emboss=False)
                    op.task = task_type.name
            else:
                split = col.split(factor=0.4)
                split.row()
                row = split.row()
                row.label(icon="ERROR")
                row.label(text='Add at least one task')

        # Choose Sequences
        row = col.row(align=True)
        row.prop(import_shots, 'import_sequence', text='Import Sequences')
        row.operator("vse_toolbox.select_sequence", text='', icon='ADD')

        sequences = [s for s in project.sequences if s.import_enabled]
        if import_shots.import_sequence == "FROM_LIST":
            if sequences:
                split = col.split(factor=0.4)
                split.row()
                box = split.box()
                box_col = box.column(align=True)

                for sequence in sequences:
                    row = box_col.row(align=True)
                    op = row.operator("vse_toolbox.unselect_sequence", text=sequence.name)
                    op.sequence = sequence.name

                    row.separator(factor=0.5)
                    op = row.operator("vse_toolbox.unselect_sequence", text='', icon='REMOVE', emboss=False)
                    op.sequence = sequence.name
            else:
                split = col.split(factor=0.4)
                split.row()
                row = split.row()
                row.label(icon="ERROR")
                row.label(text='Add at least one Sequence')

    def invoke(self, context, event):
        scn = context.scene
        settings = get_scene_settings()
        project = settings.active_project

        tmp_dir = Path(gettempdir(), 'reviews')

        if bpy.data.filepath and Path(project.import_shots.previews_folder) == tmp_dir:
            project.import_shots.previews_folder = '//sources'

        if not bpy.data.filepath and project.import_shots.previews_folder == '//sources':
            project.import_shots.previews_folder = str(tmp_dir)

        # if not bpy.data.filepath:
        #     self.report({"ERROR"}, "Save your Blender file first")
        #     return {"CANCELLED"}
        
        return context.window_manager.invoke_props_dialog(self, width=350)

    def check(self, context):
        return True


classes = (
    VSETB_OT_select_sequence,
    VSETB_OT_unselect_sequence,
    VSETB_OT_unselect_task,
    VSETB_OT_select_task,
    VSETB_OT_add_import_template,
    VSETB_OT_remove_import_template,
    VSETB_UL_import_task,
    VSETB_OT_auto_select_files,
    VSETB_OT_import_files,
    VSETB_OT_import_shots,
)

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

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