"""Generate turntable Level Sequences (and optionally queue Movie Render Queue jobs).

Run this from Unreal's Output Log with::

    exec(open(r".../Content/Python/generate_blueprint_turntable.py").read())
    generate_turntables("/Game/Blueprints/BP_MyAsset")

The script creates/saves a small presentation map, so the bindings in the
sequences and MRQ jobs are reproducible.  Open any generated sequence in
Sequencer to show the normal camera, camera-cut and transform tracks.
"""
import sys
import unreal


# Project-specific defaults.  They can be overridden in generate_turntables().
DEFAULT_OUTPUT_ROOT = "/Game/Sequences/Turntables"
DEFAULT_STAGE_MAP = "/Game/Maps/TurntableStage"
DEFAULT_RENDER_DIRECTORY = "{project_dir}/Saved/TurntableRenders"
FPS = 30
TURN_DURATION_SECONDS = 5.0
QUEUE_RENDER_IMMEDIATELY = False  # True starts MRQ immediately after queueing.
RENDER_ALL_SKELETON_ANIMATIONS = False


def _log(message):
    unreal.log("[Turntable] " + message)


def _asset_name(path):
    return path.rsplit("/", 1)[-1].split(".")[0]


def _ensure_folder(path):
    if not unreal.EditorAssetLibrary.does_directory_exist(path):
        unreal.EditorAssetLibrary.make_directory(path)


def _create_or_open_stage(stage_map):
    """Open the dedicated map; creating it only on the first run."""
    if unreal.EditorAssetLibrary.does_asset_exist(stage_map):
        unreal.EditorLoadingAndSavingUtils.load_map(stage_map)
    else:
        _ensure_folder(stage_map.rsplit("/", 1)[0])
        unreal.EditorLevelLibrary.new_level(stage_map)
        unreal.EditorLevelLibrary.save_current_level()


def _find_skeletal_components(actor):
    components = actor.get_components_by_class(unreal.SkeletalMeshComponent)
    return [component for component in components if component.get_editor_property("skeletal_mesh")]


def _animations_set_on_blueprint(actor):
    """Return AnimSequences explicitly selected on components in the Blueprint.

    Anim Blueprints are procedural at runtime and do not expose a finite list of
    clips, so they are deliberately not guessed here.  A component using
    "Use Animation Asset" is exported as one sequence per assigned clip.
    """
    animations = []
    for component in _find_skeletal_components(actor):
        try:
            animation_data = component.get_editor_property("animation_data")
            animation = animation_data.get_editor_property("anim_to_play")
            if animation and animation not in animations:
                animations.append(animation)
        except Exception:
            pass
    return animations


def _all_skeleton_animations(actor):
    """Optional broad mode: render every AnimSequence sharing the actor skeleton."""
    skeletons = set()
    for component in _find_skeletal_components(actor):
        mesh = component.get_editor_property("skeletal_mesh")
        skeleton = mesh.get_editor_property("skeleton") if mesh else None
        if skeleton:
            skeletons.add(skeleton.get_path_name())

    animations = []
    registry = unreal.AssetRegistryHelpers.get_asset_registry()
    for data in registry.get_assets_by_class("AnimSequence", True):
        animation = data.get_asset()
        if animation and animation.get_editor_property("skeleton").get_path_name() in skeletons:
            animations.append(animation)
    return animations


def _set_transform_keys(track, frames, start_location, start_rotation):
    section = track.add_section()
    section.set_range(0, frames)
    # UE 5.8 exposes section channels through the SequencerScripting extension
    # library (MovieScene3DTransformSection itself has no get_channels method).
    channels = unreal.MovieSceneSectionExtensions.get_channels_by_type(
        section, unreal.MovieSceneScriptingDoubleChannel)
    if len(channels) != 9:
        raise RuntimeError("Expected 9 transform channels, got {0}".format(len(channels)))
    # Translation XYZ, rotation XYZ, scale XYZ.  Key all channels explicitly so
    # the result remains readable/editable in Sequencer.
    values = [start_location.x, start_location.y, start_location.z,
              start_rotation.roll, start_rotation.pitch, start_rotation.yaw,
              1.0, 1.0, 1.0]
    for index, value in enumerate(values):
        channels[index].add_key(unreal.FrameNumber(0), value)
        end_value = start_rotation.yaw + 360.0 if index == 5 else value
        channels[index].add_key(unreal.FrameNumber(frames), end_value)


def _add_camera(sequence, camera, frames):
    camera_binding = sequence.add_possessable(camera)
    cut_track = sequence.add_master_track(unreal.MovieSceneCameraCutTrack)
    cut = cut_track.add_section()
    cut.set_range(0, frames)
    # Camera cuts require a MovieSceneObjectBindingID, not the binding proxy.
    cut.set_camera_binding_id(sequence.get_binding_id(camera_binding))


def _add_animation_track(sequence, actor, animation, frames):
    """Bind the matching skeletal component and play its animation clip."""
    for component in _find_skeletal_components(actor):
        mesh = component.get_editor_property("skeletal_mesh")
        skeleton = mesh.get_editor_property("skeleton") if mesh else None
        if skeleton != animation.get_editor_property("skeleton"):
            continue
        binding = sequence.add_possessable(component)
        track = binding.add_track(unreal.MovieSceneSkeletalAnimationTrack)
        section = track.add_section()
        section.set_range(0, frames)
        params = section.get_editor_property("params")
        params.set_editor_property("animation", animation)
        section.set_editor_property("params", params)
        return True
    return False


def _make_sequence(sequence_name, output_folder, actor, camera, animation=None):
    tools = unreal.AssetToolsHelpers.get_asset_tools()
    existing_path = output_folder + "/" + sequence_name
    if unreal.EditorAssetLibrary.does_asset_exist(existing_path):
        unreal.EditorAssetLibrary.delete_asset(existing_path)
    sequence = tools.create_asset(sequence_name, output_folder, unreal.LevelSequence,
                                  unreal.LevelSequenceFactoryNew())
    frames = max(1, round(FPS * max(TURN_DURATION_SECONDS,
                                   animation.get_play_length() if animation else 0.0)))
    sequence.set_display_rate(unreal.FrameRate(FPS, 1))
    sequence.set_playback_start(0)
    sequence.set_playback_end(frames)

    actor_binding = sequence.add_possessable(actor)
    transform_track = actor_binding.add_track(unreal.MovieScene3DTransformTrack)
    _set_transform_keys(transform_track, frames, actor.get_actor_location(), actor.get_actor_rotation())
    _add_camera(sequence, camera, frames)
    if animation:
        _add_animation_track(sequence, actor, animation, frames)

    unreal.EditorAssetLibrary.save_asset(sequence.get_path_name())
    return sequence


def _queue_render(sequence, stage_map, directory):
    subsystem = unreal.get_editor_subsystem(unreal.MoviePipelineQueueSubsystem)
    queue = subsystem.get_queue()
    job = queue.allocate_new_job(unreal.MoviePipelineExecutorJob)
    job.job_name = sequence.get_name()
    job.sequence = unreal.SoftObjectPath(sequence.get_path_name())
    job.map = unreal.SoftObjectPath(stage_map)
    config = job.get_configuration()
    output = config.find_or_add_setting_by_class(unreal.MoviePipelineOutputSetting)
    output.output_directory = unreal.DirectoryPath(directory)
    output.file_name_format = "{job_name}/{frame_number}"
    output.output_resolution = unreal.IntPoint(1920, 1080)
    config.find_or_add_setting_by_class(unreal.MoviePipelineImageSequenceOutput_PNG)
    return subsystem


def generate_turntables(blueprint_path, output_root=DEFAULT_OUTPUT_ROOT,
                        stage_map=DEFAULT_STAGE_MAP, render_directory=DEFAULT_RENDER_DIRECTORY,
                        render_immediately=QUEUE_RENDER_IMMEDIATELY,
                        render_all_skeleton_animations=RENDER_ALL_SKELETON_ANIMATIONS):
    """Create a turntable plus one turntable per explicitly assigned animation.

    `blueprint_path` must be an asset path, e.g. /Game/Blueprints/BP_Robot.
    The return value is the generated LevelSequence list.
    """
    blueprint = unreal.load_asset(blueprint_path)
    if not blueprint or not isinstance(blueprint, unreal.Blueprint):
        raise RuntimeError("Expected a Blueprint asset path, got: " + blueprint_path)

    bp_name = _asset_name(blueprint_path)
    # A separate stage per asset prevents other turntable actors from appearing
    # in the same render.  Supplying a custom stage_map preserves that exact map.
    if stage_map == DEFAULT_STAGE_MAP:
        stage_map = stage_map + "_" + bp_name
    _create_or_open_stage(stage_map)
    folder = output_root + "/" + bp_name
    _ensure_folder(folder)

    # Re-running the script replaces only its own presentation actors in the
    # selected stage map.  This keeps repeated demos deterministic.
    generated_labels = {"TT_" + bp_name, "TT_Camera_" + bp_name}
    for existing_actor in unreal.EditorLevelLibrary.get_all_level_actors():
        if existing_actor.get_actor_label() in generated_labels:
            unreal.EditorLevelLibrary.destroy_actor(existing_actor)

    actor = unreal.EditorLevelLibrary.spawn_actor_from_object(blueprint, unreal.Vector(0, 0, 0))
    actor.set_actor_label("TT_" + bp_name)
    origin, extent = actor.get_actor_bounds(False)
    distance = max(extent.x, extent.y, extent.z) * 3.0 + 150.0
    camera = unreal.EditorLevelLibrary.spawn_actor_from_class(
        unreal.CineCameraActor, unreal.Vector(origin.x - distance, origin.y, origin.z), unreal.Rotator(0, 0, 0))
    camera.set_actor_label("TT_Camera_" + bp_name)
    camera.set_actor_rotation(
        unreal.MathLibrary.find_look_at_rotation(camera.get_actor_location(), origin), False)
    camera.get_cine_camera_component().set_editor_property("current_focal_length", 50.0)

    sequences = [_make_sequence("LS_TT_" + bp_name, folder, actor, camera)]
    animations = _all_skeleton_animations(actor) if render_all_skeleton_animations else _animations_set_on_blueprint(actor)
    for animation in animations:
        name = "LS_TT_{0}_{1}".format(bp_name, animation.get_name())
        sequences.append(_make_sequence(name, folder, actor, camera, animation))

    unreal.EditorLevelLibrary.save_current_level()
    subsystem = None
    for sequence in sequences:
        subsystem = _queue_render(sequence, stage_map, render_directory)
    if render_immediately and subsystem:
        subsystem.render_queue_with_executor(unreal.MoviePipelinePIEExecutor)

    _log("Created {0} sequence(s) and queued {0} PNG job(s) for {1}.".format(len(sequences), bp_name))
    return sequences


def _run_from_console_arguments():
    """Support `py "script.py" /Game/Path/BP_Name` in the Unreal console."""
    for argument in sys.argv[1:]:
        # Dragging an asset into the Output Log can produce
        # /Script/Engine.Blueprint'/Game/Foo/BP_X.BP_X'; accept that form too.
        game_index = argument.find("/Game/")
        if game_index < 0:
            continue
        asset_path = argument[game_index:].split("'", 1)[0].split(".", 1)[0]
        generate_turntables(asset_path)
        return


_run_from_console_arguments()