vse_toolbox/operators/tracker.py

import os
from pathlib import Path

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

from vse_toolbox.constants import (ASSET_PREVIEWS, PREVIEWS_DIR, ASSET_ITEMS)
    
from vse_toolbox.sequencer_utils import (get_strips, get_strip_render_path, get_strip_sequence_name)

from vse_toolbox.bl_utils import (get_addon_prefs, get_scene_settings)
from vse_toolbox.file_utils import (norm_name,)
from bpy.app.handlers import persistent


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_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']
            
            #for key, value in asset_data.get('data', {}).items():
            asset['metadata'] = asset_data.get('data', {})

            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)
        
        set_asset_items()

        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

        # old_project_name = settings.project_name.replace(' ', '_').upper()

        # old_episode_name = None
        # if settings.active_project:
        #     old_episode_name = settings.active_project.episode_name.replace(' ', '_').upper()
       


        #settings.projects.clear()
        tracker.connect()

        project_datas = tracker.get_projects()
        for project_data in project_datas:
            project = settings.projects.get(project_data['name'])
            if not project:
                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':
                episode_datas = tracker.get_episodes(project_data)
                for episode_data in episode_datas:
                    episode = project.episodes.get(episode_data['name'])
                    
                    if not episode:
                        episode = project.episodes.add()

                    episode.name = episode_data['name']
                    episode.id = episode_data['id']
                
                # Clear deleted episodes
                ep_names = [e['name'] for e in episode_datas]
                for ep in reversed(project.episodes):
                    if ep.name not in ep_names:
                        project.episodes.remove(list(project.episodes).index(ep))

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

                if prefs.sort_metadata_items:
                    metadata_data['choices'].sort()

                for choice in metadata_data['choices']:
                    choice_item = metadata_type.choices.add()
                    choice_item.name = choice
                
                metadata_type['entity_type'] = metadata_data['entity_type'].upper()
            
            project.task_statuses.clear()
            for status_data in tracker.get_task_statuses(project_data):
                #print(metadata_data)
                task_status = project.task_statuses.add()
                task_status.name = status_data['short_name'].upper()

            project.task_types.clear()
            for task_type_data in tracker.get_shot_task_types(project_data):
                task_type = project.task_types.add()
                task_type.name = task_type_data['name']

            project.asset_types.clear()
            for asset_type_data in tracker.get_asset_types(project_data):
               asset_type = project.asset_types.add()
               asset_type.name = asset_type_data['name']
            
            project.set_spreadsheet()

        bpy.ops.vse_toolbox.load_settings()

        # Remove deleted projects
        project_names = [p['name'] for p in project_datas]
        for project in reversed(settings.projects):
            if project.name not in project_names:
                settings.projects.remove(list(settings.projects).index(project))


        #if settings.active_project:
        #    settings.active_project.set_strip_metadata()
        

        self.report({"INFO"}, 'Successfully Load Tracker Projects')

        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'}


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

    status_items = [('CURRENT', 'Current', '')]

    if project:
        status_items += [(t.name, t.name, '') for t in project.task_statuses]
    return status_items

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

    if not project:
        return [('NONE', 'None', '')]

    return [(t.name, t.name, '') for t in project.task_types]

class VSETB_OT_upload_to_tracker(Operator):
    bl_idname = "vse_toolbox.upload_to_tracker"
    bl_label = "Upload to tracker"
    bl_description = "Upload selected strip to tracker"
    bl_options = {"REGISTER", "UNDO"}

    task : EnumProperty(items=get_task_type_items)
    status : EnumProperty(items=get_task_status_items)
    comment : StringProperty()
    add_preview : BoolProperty(default=True)
    preview_mode : EnumProperty(items=[(m, m.title().replace('_', ' '), '') 
        for m in ('ONLY_NEW', 'REPLACE', 'ADD')])
    set_main_preview : BoolProperty(default=True)
    casting : BoolProperty(default=True)
    custom_data : BoolProperty(default=True)

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

    def invoke(self, context, event):
        prefs = get_addon_prefs()
        settings = get_scene_settings()
        project = settings.active_project

        tracker = prefs.tracker
        tracker.connect()
        
        #self.bl_label = f"Upload to {settings.tracker_name.title()}"

        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, 'task', text='Task')
        col.prop(self, 'status', text='Status')
        col.prop(self, 'comment', text='Comment')
        row = col.row(heading='Add Preview')
        row.prop(self, 'add_preview', text='')
        row.prop(self, 'preview_mode', text='')
        col.separator()
        col.prop(self, 'casting', text='Casting')
        col.prop(self, 'custom_data', text='Custom Data')
        col.prop(self, 'set_main_preview', text='Set Main Preview')

    def execute(self, context):
        #self.report({'ERROR'}, f'Export not implemented yet.')
        prefs = get_addon_prefs()
        settings = get_scene_settings()
        project = settings.active_project
        episode = None
        if settings.active_episode:
            episode = settings.active_episode.id

        format_data = {**settings.format_data, **project.format_data}


        tracker = prefs.tracker

        status = self.status
        if status == 'CURRENT':
            status = None

        for strip in get_strips(channel='Shots', selected_only=True):
            strip_settings = strip.vsetb_strip_settings
            sequence_name = get_strip_sequence_name(strip)
            shot_name = strip.name
            sequence = tracker.get_sequence(sequence_name, episode=episode)
            metadata = strip_settings.metadata.to_dict()
            #print(metadata)

            if not sequence:
                self.report({"INFO"}, f'Create sequence {sequence_name} in Kitsu')
                sequence = tracker.new_sequence(sequence_name, episode=episode)

            shot = tracker.get_shot(shot_name, sequence=sequence)
            if not shot:
                self.report({"INFO"}, f'Create shot {shot_name} in Kitsu')
                shot = tracker.new_shot(shot_name, sequence=sequence)

            task = tracker.get_task(self.task, entity=shot)
            if not task:
                task = tracker.new_task(shot, task_type=self.task)

            # print('\n', 'task comment')
            # print(task['last_comment'])

            preview = None
            if self.add_preview:
                
                strip_data = {**format_data, **strip_settings.format_data}
                preview = project.render_strip_template.format(**strip_data)
                preview = Path(os.path.abspath(bpy.path.abspath(preview)))
                #print(preview)
                if not preview.exists():
                    print(f'The preview {preview} not exists')
                    preview = None

                elif task.get('last_comment') and task['last_comment']['previews']:
                    if self.preview_mode == 'REPLACE':
                        tracker.remove_comment(task['last_comment'])
                    elif self.preview_mode == 'ONLY_NEW':
                        preview = None

            #print(f'{preview=}')
            #print(f'{status=}')
            if status or preview:
                tracker.new_comment(task, comment=self.comment, status=status, preview=preview, set_main_preview=self.set_main_preview)


            if self.custom_data:
                metadata = strip_settings.metadata.to_dict()
                description = strip_settings.description
                tracker.update_data(shot, metadata, frames=strip.frame_final_duration, description=description)
            
            if self.casting:
                casting = [{'asset_id': a.id, 'nb_occurences': a.instance} for a in strip_settings.casting]
                tracker.update_casting(shot, casting)


        return {"FINISHED"}
    

@persistent
def set_asset_items(scene=None):
    ASSET_ITEMS.clear()
    
    settings = get_scene_settings()
    if settings.active_project:
        ASSET_ITEMS.extend([(a.id, a.label, '', i) for i, a in enumerate(settings.active_project.assets)])


classes = (
    VSETB_OT_load_assets,
    VSETB_OT_load_projects,
    VSETB_OT_new_episode,
    VSETB_OT_tracker_connect,
    VSETB_OT_upload_to_tracker,
)

def register():
    if not bpy.app.background:
        bpy.app.handlers.load_post.append(set_asset_items)

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

def unregister():
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)
    
    if not bpy.app.background:
        bpy.app.handlers.load_post.remove(set_asset_items)