

import sys
from pathlib import Path
#sys.path.append(str(Path(__file__).parents[3]))

from asset_library.action.concat_preview import mosaic_export
from asset_library.common.file_utils import open_file
from asset_library.action.functions import reset_bone, get_keyframes
from asset_library.common.functions import read_catalog

import bpy
import argparse
import json
import re
import shutil
import subprocess

from tempfile import gettempdir



def rm_tree(pth):
    pth = Path(pth)
    for child in pth.glob('*'):
        if child.is_file():
            child.unlink()
        else:
            rm_tree(child)
    pth.rmdir()

def render_preview(directory, asset_catalog, render_actions, publish_actions, remove_folder):

    scn = bpy.context.scene
    rnd = bpy.context.scene.render
    rnd.resolution_x = rnd.resolution_y = 512

    report = []
    blendfile = Path(bpy.data.filepath)
    asset_catalog_data = read_catalog(asset_catalog)

    anim_render_dir = Path(gettempdir()) / 'actionlib_render'   #/tmp/actionlib_render. Removed at the end
    anim_render_dir.mkdir(exist_ok=True, parents=True)

    preview_render_dir = Path(directory) / 'preview'

    if preview_render_dir.exists() and remove_folder:
        rm_tree(preview_render_dir)

    preview_render_dir.mkdir(exist_ok=True, parents=True)
    for i in ('anim', 'pose'):
        Path(preview_render_dir / i).mkdir(exist_ok=True, parents=True)
    
    for f in preview_render_dir.rglob('*'):
        if f.is_dir():
            print(f'{f} is dir. Skipped.')
            continue
        if all(i not in f.parts for i in ('anim', 'pose')) and f.parent.parts[-1] != 'preview':
            print(f'{f} is out of pipe. Approved or Rtk pictures. Skipped.')
            continue
        if not any(f.stem.endswith(a) for a in publish_actions):
            print(f'{str(f)} not in publish actions anymore. Removing...')
            f.unlink()

    # Set Scene
    # ----------
    # Scene Setting
    scn.use_preview_range = True
    scn.eevee.use_gtao = True
    scn.tool_settings.use_keyframe_insert_auto = False
    
    # Render Setting
    rnd.engine = 'BLENDER_EEVEE'
    rnd.use_simplify = False
    rnd.use_stamp_date = True
    rnd.use_stamp_time = True
    rnd.use_stamp_render_time = False
    rnd.use_stamp_frame = True
    rnd.use_stamp_frame_range = False
    rnd.use_stamp_memory = False
    rnd.use_stamp_hostname = False
    rnd.use_stamp_camera = True
    rnd.use_stamp_lens = False
    rnd.use_stamp_scene = False
    rnd.use_stamp_marker = False
    rnd.use_stamp_filename = False
    rnd.use_stamp_sequencer_strip = False
    rnd.use_stamp_note = True
    rnd.use_stamp = True
    rnd.stamp_font_size = 16
    rnd.use_stamp_labels = False
    rnd.image_settings.file_format = 'JPEG'

    # Viewport Look
    # ----------
    """
    # Eevee
    for screen in bpy.data.screens:
        for area in screen.areas:
            for space in area.spaces:
                if space.type == 'VIEW_3D':
                    space.overlay.show_overlays = False
                    space.shading.type = 'RENDERED'
                    space.shading.use_scene_lights_render = False
                    space.shading.use_scene_world_render = False
                    space.region_3d.view_perspective = 'CAMERA'
    
    """
    # Cycles Mat Shading
    for a in bpy.context.screen.areas:
        if a.type == 'VIEW_3D':
            a.spaces[0].overlay.show_overlays = False
            a.spaces[0].region_3d.view_perspective = 'CAMERA'
            a.spaces[0].shading.show_cavity = True
            a.spaces[0].shading.cavity_type = 'WORLD'
            a.spaces[0].shading.cavity_ridge_factor = 0.75
            a.spaces[0].shading.cavity_valley_factor = 1.0
    

    # Add Subsurf
    # -----------
    deform_ob = [m.object for o in scn.objects \
            for m in o.modifiers if m.type == 'MESH_DEFORM'
    ]
    deform_ob += [m.target for o in scn.objects \
            for m in o.modifiers if m.type == 'SURFACE_DEFORM'
    ]
    
    objects = [o for o in bpy.context.scene.objects if (o.type == 'MESH' 
        and o not in deform_ob and o not in bpy.context.scene.collection.objects[:])
    ]
    
    for o in objects:
        subsurf = False
        for m in o.modifiers:
            if m.type == 'SUBSURF':
                m.show_viewport = m.show_render
                m.levels = m.render_levels
                subsurf = True
                break

        if not subsurf:
            subsurf = o.modifiers.new('', 'SUBSURF')
            subsurf.show_viewport = subsurf.show_render
            subsurf.levels = subsurf.render_levels


    # Loop through action and render
    # ------------------------------
    rig = next((o for o in scn.objects if o.type == 'ARMATURE'), None)
    # actions = [a for a in bpy.data.actions if a.asset_data]


    rig.animation_data_create()
    for action_name in render_actions:
        action = bpy.data.actions.get(action_name)

        if not action:
            print(f'\'{action_name}\' not found.')
            continue

        print(f"-- Current --: {action.name}")
    
        rnd.stamp_note_text = '{type} : {pose_name}'
        action_data = action.asset_data
    
        if 'camera' not in action_data.keys():
            report.append(f"'{action.name}' has no CameraData.")
            continue
        
        catalog_name = next((v['name'] for v in asset_catalog_data.values() if action_data.catalog_id == v['id']), None)
        pose_name = '/'.join([*catalog_name.split('-'), action.name])
        filename = bpy.path.clean_name(f'{catalog_name}_{action.name}')
        ext = 'jpg'

        rig.animation_data.action = None
        bpy.context.view_layer.update()
        for b in rig.pose.bones:
            if re.match('^[A-Z]+\.', b.name):
                continue
            reset_bone(b)

        rest_pose = None
        if isinstance(action.asset_data.get('rest_pose'), str):
            rest_pose = bpy.data.actions.get(action.asset_data['rest_pose'])
        
        rig.animation_data.action = rest_pose
        bpy.context.view_layer.update()
        
        rig.animation_data.action = action

        if 'camera' in action.asset_data.keys():
            action_cam = bpy.data.objects.get(action.asset_data['camera'], '')
            if action_cam:
                scn.camera = action_cam

        # Is Anim
        if not action_data['is_single_frame'] or 'anim' in action_data.tags.keys():
            keyframes = get_keyframes(action)
            if not keyframes:
                continue
            anim_start = keyframes[0]
            anim_end = keyframes[-1]

            if anim_start < scn.frame_start:
                report.append(f"Issue found for '{action.name}'. Has keyframes before 'Start Frame'.")
                continue
            
            scn.frame_preview_start = anim_start
            scn.frame_preview_end = anim_end

            rnd.stamp_note_text = rnd.stamp_note_text.format(
                type='ANIM', 
                pose_name=pose_name,
            )

            rnd.filepath = f'{str(anim_render_dir)}/{filename}_####.{ext}'
            
            bpy.ops.render.opengl(animation=True)

            ffmpeg_cmd = [
                'ffmpeg', '-y', 
                '-start_number', f'{anim_start:04d}',
                '-i', rnd.filepath.replace('####', '%04d'),
                '-c:v', 'libx264', 
                str((preview_render_dir/'anim'/filename).with_suffix('.mov')),
            ]
            subprocess.call(ffmpeg_cmd)
            
        # Is Pose
        elif action_data['is_single_frame'] or 'pose' in action_data.tags.keys():
            scn.frame_preview_start = scn.frame_preview_end = scn.frame_start
            
            rnd.stamp_note_text = rnd.stamp_note_text.format(
                type='POSE', 
                pose_name=pose_name,
            )

            rnd.filepath = f'{str(preview_render_dir)}/pose/{filename}_####.{ext}'
            
            bpy.ops.render.opengl(animation=True)

            filename = rnd.filepath.replace('####', f'{scn.frame_preview_end:04d}')
            Path(filename).rename(re.sub('_[0-9]{4}.', '.', filename))

    shutil.rmtree(anim_render_dir)
    
    # Report
    # ------
    if report:
        report_file = blendfile.parent / Path(f'{blendfile.stem}report').with_suffix('.txt')
        if not report_file.exists():
            report_file.touch(exist_ok=False)

        report_file.write_text('-')
        report_file.write_text('\n'.join(report))
        
        result = report_file

    else:
        result = preview_render_dir

    open_file(result)

    files = [str(f) for f in sorted((preview_render_dir/'pose').glob('*.jpg'))]

    mosaic_export(
        files=files, catalog_data=asset_catalog_data,
        row=2, columns=2, auto_calculate=True,
        bg_color=(0.18, 0.18, 0.18,), resize_output=100
    )

    bpy.ops.wm.quit_blender()



if __name__ == '__main__' :
    parser = argparse.ArgumentParser(description='Add Comment To the tracker',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('--directory')
    parser.add_argument('--asset-catalog')
    parser.add_argument('--render-actions', nargs='+')
    parser.add_argument('--publish-actions', nargs='+')
    parser.add_argument('--remove-folder', type=json.loads, default='false')

    if '--' in sys.argv :
        index = sys.argv.index('--')
        sys.argv = [sys.argv[index-1], *sys.argv[index+1:]]

    args = parser.parse_args()
    render_preview(**vars(args))
