# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0. See the LICENSE file in this repository.

from dataclasses import replace

from diffusers.modular_pipelines import ModularPipeline, SequentialPipelineBlocks
from diffusers.modular_pipelines.modular_pipeline_utils import OutputParam
from diffusers.utils.dynamic_modules_utils import get_class_from_dynamic_module

from .decoders import MotionBricksDecodeStep, MotionBricksPostprocessStep
from .encoders import MotionBricksGuidesStep
from .generation import MotionBricksPoseTokensStep, MotionBricksRootStep

# The remote-code loader only copies modules imported here directly, so the shared helpers are imported too.
from .motion_utils import encode_guides  # noqa: F401


class MotionBricksBlocks(SequentialPipelineBlocks):
    """Motion in-betweening with NVIDIA MotionBricks on the SOMA skeleton.

    Every component loads from `components_repo`, one subfolder per component, and so does its code, which requires
    `trust_components_code=True`.
    """

    model_name = "motionbricks"
    block_classes = [
        MotionBricksGuidesStep(),
        MotionBricksRootStep(),
        MotionBricksPoseTokensStep(),
        MotionBricksDecodeStep(),
        MotionBricksPostprocessStep(),
    ]
    block_names = ["guides", "root", "pose_tokens", "decode", "postprocess"]

    _requirements = {"diffusers": ">=0.41.0"}

    def __init__(self, components_repo=None, components_revision=None, trust_components_code=False):
        super().__init__()
        self.components_repo = components_repo
        self.components_revision = components_revision
        self.component_types = {}
        if components_repo is None:
            return
        if not trust_components_code:
            raise ValueError("Set `trust_components_code=True` to load the model code from the components repository")
        index = ModularPipeline.load_config(str(components_repo), revision=components_revision)
        for spec in super().expected_components:
            module_name, class_name, _ = index[spec.name]
            self.component_types[spec.name] = get_class_from_dynamic_module(
                str(components_repo),
                subfolder=spec.name,
                module_file=f"{module_name}.py",
                class_name=class_name,
                revision=components_revision,
                trust_remote_code=True,
            )

    @property
    def expected_components(self):
        if self.components_repo is None:
            return super().expected_components
        return [
            replace(
                spec,
                type_hint=self.component_types.get(spec.name, spec.type_hint),
                pretrained_model_name_or_path=str(self.components_repo),
                subfolder=spec.name,
                revision=self.components_revision,
            )
            for spec in super().expected_components
        ]

    @property
    def description(self) -> str:
        return (
            "Generates the transition from the last poses of a context motion to target poses, choosing its length "
            "among the allowed ones."
        )

    @property
    def outputs(self):
        return [
            OutputParam(
                "motion",
                type_hint=list,
                description="One SOMA30 motion dict per batch item, 30 FPS, meters, +Y up.",
            ),
            OutputParam(
                "placed_target",
                type_hint=list,
                description="`target` as used, one dict per batch item, moved by `target_offset` and `target_turn`.",
            ),
        ]