vse_toolbox/operators/spreadsheet.py

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

import csv
from datetime import datetime
import os
from pathlib import Path

import bpy
from bpy.types import (Operator, Menu, OperatorFileListElement)
from bpy.props import (EnumProperty, StringProperty, CollectionProperty)
from bl_operators.presets import AddPresetBase


from vse_toolbox.sequencer_utils import (get_strips, get_strip_sequence_name)
from vse_toolbox.bl_utils import (get_addon_prefs, get_scene_settings)
from vse_toolbox.file_utils import (open_file, install_module, fuzzy_match)


class VSETB_MT_export_spreadsheet_presets(Menu): 
    bl_label = 'Presets' 
    preset_subdir = 'vse_toolbox' 
    preset_operator = 'script.execute_preset' 
    draw = Menu.draw_preset


class VSETB_OT_add_export_spreadsheet_preset(AddPresetBase, Operator):

    bl_idname = "vse_toolbox.add_spreadsheet_preset"
    bl_label = "Add Spreadsheet Preset"
    bl_description = "Add Spreadsheet Preset"
    bl_options = {"REGISTER", "UNDO"}

    preset_menu = 'VSETB_MT_export_spreadsheet_presets'

    #preset_menu = 'VSETB_OT_MT_spreadsheet_presets'

    # Common variable used for all preset values
    #C.scene.vsetb_settings.active_project.spreadsheet_options
    preset_defines = [
                        'scene = bpy.context.scene',
                        'settings = scene.vsetb_settings',
                        'project = settings.active_project'
                     ]

    # Properties to store in the preset
    preset_values = [
                        'project.spreadsheet_export',
                    ]

    # Directory to store the presets
    preset_subdir = 'vse_toolbox/spreadsheet_export'



class VSETB_OT_spreadsheet_cell_move(Operator):
    bl_idname = "vse_toolbox.spreadsheet_move"
    bl_label = "Move Spreadsheet items"
    bl_description = "Move Spreadsheet items"
    bl_options = {"REGISTER", "UNDO"}

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

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

        spreadsheet = project.spreadsheet_export
        cells = spreadsheet.cells
        idx = spreadsheet.cell_index

        try:
            item = cells[idx]
        except IndexError:
            pass
        else:
            if self.direction == 'DOWN' and idx < len(cells) - 1:
                item_next = cells[idx+1].name
                cells.move(idx, idx+1)
                spreadsheet.cell_index += 1

            elif self.direction == 'UP' and idx >= 1:
                item_prev = cells[idx-1].name
                cells.move(idx, idx-1)
                spreadsheet.cell_index -= 1

        return {"FINISHED"}


class VSETB_OT_spreadsheet_from_file(Operator):
    bl_idname = "vse_toolbox.spreadsheet_from_file"
    bl_label = "Read Spreadsheet Column"
    bl_description = "Read Spreadsheet Column"
    bl_options = {"REGISTER", "UNDO"}

    filepath : StringProperty()

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

        if not self.filepath:
            self.report({'ERROR'}, 'No filepath provided')
            return {'CANCELLED'}
        
        if not Path(self.filepath).exists():
            self.report({'ERROR'}, f'Filepath {self.filepath} not exists')
            return {'CANCELLED'}

        filepath = Path(self.filepath)

        if filepath.suffix.lower() == '.csv':
            with open(filepath, newline='') as csvfile:
                reader = csv.reader(csvfile)
                next(reader)
                data = dict(reader)


        elif filepath.suffix.lower() == '.xlsx':
            raise Exception('Not Supported')
            workbook = openpyxl.load_workbook(filepath, read_only=True)
            sheet = next((ws for ws in workbook.worksheets if ws.title.lower() == 'export'), None)
        
        
        else:
            self.report({'ERROR'}, f'File extension {filepath.suffix} should be in [.csv, .xlsx]')
            return {'CANCELLED'}

        return {"FINISHED"}
    

class VSETB_OT_spreadsheet_from_clipboard(Operator):
    bl_idname = "vse_toolbox.spreadsheet_from_clipboard"
    bl_label = "Read Spreadsheet from clipboard"
    bl_description = "Read Spreadsheet from clipboard"
    bl_options = {"REGISTER", "UNDO"}

    def execute(self, context):
        scn = context.scene
        project = get_scene_settings().active_project
        import_cells = project.spreadsheet_import.cells
        import_cells.clear()

        spreadsheet = context.window_manager.clipboard

        cell_types = project.get_cell_types()
        rows = list(csv.reader(spreadsheet.splitlines(), delimiter='\t'))
        for cell_name in rows[0]:
            if not cell_name:
                continue

            cell = import_cells.add()
            cell.name = cell_name
            cell.import_name = max(cell_types.keys(), key=lambda x: fuzzy_match(cell_name, x))
            cell.enabled = True
        
        project.spreadsheet_import.use_custom_cells = True

        return {"FINISHED"}

class VSETB_OT_import_spreadsheet(Operator):
    bl_idname = "vse_toolbox.import_spreadsheet"
    bl_label = "Import Spreadsheet"
    bl_description = "Create strips from nb frames with casting and custom data"
    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=OperatorFileListElement)

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

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

        context.window_manager.fileselect_add(self)
        return {'RUNNING_MODAL'}

    def draw(self, context):
        scn = context.scene
        settings = get_scene_settings()
        project = settings.active_project
        spreadsheet = project.spreadsheet_import
        #options = project.spreadsheet_options

        layout = self.layout

        row = layout.row(align=True)
        row.label(text='Source')
        #row.alignment='RIGHT'
        row.operator("vse_toolbox.spreadsheet_from_clipboard", text='Clipboard', icon='PASTEDOWN')
        row.operator("vse_toolbox.spreadsheet_from_file", text='File', icon='FILE')

        col = layout.column(align=False)

        row = col.row(align=True, heading='Custom Asset Name')
        row.use_property_split = True
        row.use_property_decorate = False
        row.prop(spreadsheet, 'use_custom_name', text='')
        sub = row.row(align=False)
        sub.enabled = spreadsheet.use_custom_name
        sub.prop(spreadsheet, 'custom_name', text='')
        sub.label(icon='BLANK1')

        row = layout.row()
        row.template_list("VSETB_UL_spreadsheet_import", "spreadsheet_import", spreadsheet, "cells", spreadsheet, "cell_index", rows=8)

        col_tool = row.column(align=True)

        #bpy.types.VSETB_PT_presets.draw_panel_header(col_tool)
        #col_tool.operator('wm.call_menu', icon="PRESET").name = 'VSETB_MT_spreadsheet_presets'
        #col_tool.operator('vse_toolbox.load_spreadsheet_preset', icon='PRESET', text="")
        op = col_tool.operator('wm.call_panel', icon="PRESET", emboss=False, text='')
        op.name = 'VSETB_PT_presets'
        op.keep_open = False

        #col_tool.separator()
        #col_tool.operator('vse_toolbox.spreadsheet_move', icon='TRIA_UP', text="").direction = 'UP'
        #col_tool.operator('vse_toolbox.spreadsheet_move', icon='TRIA_DOWN', text="").direction = 'DOWN'

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

        #col.separator()

        row = col.row(align=False)

        col.prop(spreadsheet, "separator", expand=True, text='Separator')
        if self.filepath.endswith('.csv'):
            col.prop(spreadsheet, "delimiter", expand=True, text='Delimiter')
        

        col.prop(spreadsheet, 'import_casting', text='Import Casting')
        col.prop(spreadsheet, 'import_custom_data', text='Import Custom Data')
        col.prop(spreadsheet, 'update_edit', text='Update Edit')

        col.separator()

    def execute(self, context):

        return {"FINISHED"}
    

class VSETB_OT_export_spreadsheet(Operator):
    bl_idname = "vse_toolbox.export_spreadsheet"
    bl_label = "Export Spreadsheet"
    bl_description = "Export Shot data in a table as a csv or an xlsl"
    bl_options = {"REGISTER", "UNDO"}

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

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

        return context.window_manager.invoke_props_dialog(self)

    def draw(self, context):
        scn = context.scene
        settings = get_scene_settings()
        project = settings.active_project
        spreadsheet = project.spreadsheet_export
        #options = project.spreadsheet_options

        layout = self.layout
        row = layout.row(align=True)
        row.prop(spreadsheet, 'use_custom_cells', text='Custom Cells')

        col = layout.column(align=False)
        col.enabled = spreadsheet.use_custom_cells

        row = col.row(align=True, heading='Custom Asset Name')
        row.use_property_split = True
        row.use_property_decorate = False
        row.prop(spreadsheet, 'use_custom_name', text='')
        sub = row.row(align=False)
        sub.enabled = spreadsheet.use_custom_name
        sub.prop(spreadsheet, 'custom_name', text='')
        sub.label(icon='BLANK1')
        
        row = col.row()
        row.template_list("VSETB_UL_spreadsheet_export", "spreadsheet_export", spreadsheet, "cells", spreadsheet, "cell_index", rows=8)

        col_tool = row.column(align=True)

        #bpy.types.VSETB_PT_presets.draw_panel_header(col_tool)
        #col_tool.operator('wm.call_menu', icon="PRESET").name = 'VSETB_MT_spreadsheet_presets'
        #col_tool.operator('vse_toolbox.load_spreadsheet_preset', icon='PRESET', text="")
        op = col_tool.operator('wm.call_panel', icon="PRESET", emboss=False, text='')
        op.name = 'VSETB_PT_presets'
        op.keep_open = False

        col_tool.separator()
        col_tool.operator('vse_toolbox.spreadsheet_move', icon='TRIA_UP', text="").direction = 'UP'
        col_tool.operator('vse_toolbox.spreadsheet_move', icon='TRIA_DOWN', text="").direction = 'DOWN'

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

        col.separator()

        row = col.row(align=False)
        row.prop(spreadsheet, "format", expand=True, text='Format')
        row.prop(spreadsheet, 'show_settings', text='', icon='PREFERENCES')
        if spreadsheet.show_settings:
            col.prop(spreadsheet, "separator", expand=True, text='Separator')
            if spreadsheet.format == 'CSV':
                col.prop(spreadsheet, "delimiter", expand=True, text='Delimiter')
        
        col.separator()
        col.prop(spreadsheet, 'open_folder', text='Open Folder')
        col.prop(spreadsheet, 'export_path', text='Export Path')

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

        cells = [cell for cell in spreadsheet.cells if cell.enabled]
        rows = []

        # Header
        rows.append([cell.export_name for cell in cells])

        separator = spreadsheet.separator.replace('\\n', '\n').replace('\\t', '\t').replace('\\r', '\r')
        delimiter = spreadsheet.delimiter.replace('\\n', '\n').replace('\\t', '\t').replace('\\r', '\r')

        for strip in get_strips('Shots'):
            row = []
            for cell in cells:
                if cell.type == "METADATA":
                    row += [getattr(strip.vsetb_strip_settings.metadata, cell.field_name)]
                elif cell.type == "ASSET_TYPE":
                    asset_castings = []
                    for asset_casting in strip.vsetb_strip_settings.casting:
                        asset = asset_casting.asset
                        if not asset.asset_type == cell.name:
                            continue
                        
                        if spreadsheet.use_custom_name:
                            if asset.get('metadata', {}).get(spreadsheet.custom_name):
                                asset_castings.append(asset['metadata'][spreadsheet.custom_name])
                            else:
                                self.report({'ERROR'}, f'The asset {asset.tracker_name} has no data {spreadsheet.custom_name}')
                        else:
                            asset_castings.append(asset.tracker_name)

                    row += [separator.join(asset_castings)]
                elif cell.field_name == 'EPISODE':
                    row += [settings.active_episode.name]
                elif cell.field_name == 'SEQUENCE':
                    row += [get_strip_sequence_name(strip)]
                elif cell.field_name == 'SHOT':
                    row += [strip.name]
                elif cell.field_name == 'DESCRIPTION':
                    row += [strip.vsetb_strip_settings.description]
                elif cell.field_name == 'FRAMES':
                    row += [strip.frame_final_duration]
            
            rows.append(row)
        
        #print(rows)

        export_path = Path(os.path.abspath(bpy.path.abspath(spreadsheet.export_path)))
        export_name = export_path.name

        if export_path.suffix or export_name.endswith('{ext}'):
            export_path = export_path.parent

        else: # It's a directory
            if project.type == 'TVSHOW':
                export_name = '{date}_{project}_{episode}_{tracker}_shots.{ext}'
            else:
                export_name = '{date}_{project}_{tracker}_shots.{ext}'

        date = datetime.now().strftime('%Y_%m_%d')
        project_name = project.name.replace(' ', '_').lower()
        episode_name = episode.name.replace(' ', '_').lower() if episode else 'episode'
        ext = spreadsheet.format.lower()

        export_name = export_name.format(date=date, project=project_name, 
            episode=episode_name, tracker=settings.tracker_name.lower(), ext=ext)

        export_path = export_path / export_name

        #2023_04_11_kitsu_boris_ep01_shots
        export_path.parent.mkdir(parents=True, exist_ok=True)

        if spreadsheet.format == 'CSV':
            print('Writing .csv file to', export_path)
            with open(str(export_path), 'w', newline='\n', encoding='utf-8') as f:
                writer = csv.writer(f, delimiter=spreadsheet.delimiter)
                for row in rows:
                    writer.writerow(row)

        elif spreadsheet.format == 'XLSX':
            try:
                import openpyxl
            except ModuleNotFoundError:
                self.report({'INFO'}, 'Installing openpyxl')
                openpyxl = install_module('openpyxl')

            from openpyxl import Workbook

            workbook = Workbook()
            worksheet = workbook.active
            for row in rows:
                worksheet.append(row)

            for col in worksheet.columns:
                letter = col[0].column_letter
                worksheet.column_dimensions[letter].auto_size = True

            # Save the file
            workbook.save(str(export_path))

        if spreadsheet.open_folder:
            open_file(export_path, select=True)

        return {"FINISHED"}
    

classes = (
    VSETB_OT_add_export_spreadsheet_preset,
    VSETB_MT_export_spreadsheet_presets,
    VSETB_OT_spreadsheet_from_file,
    VSETB_OT_spreadsheet_from_clipboard,
    VSETB_OT_spreadsheet_cell_move,
    VSETB_OT_export_spreadsheet,
    VSETB_OT_import_spreadsheet
)

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

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