import pickle

import smplx
import torch
from smplx.lbs import vertices2joints
from smplx.utils import MANOOutput, to_tensor
from smplx.vertex_ids import vertex_ids


# SECURITY: Restricted unpickler to prevent arbitrary code execution
# Only allows numpy arrays and basic types needed for joint regressor data
class _RestrictedUnpickler(pickle.Unpickler):
    """Unpickler that only allows safe types for MANO joint regressor."""

    ALLOWED_MODULES = {
        "numpy": {"ndarray", "dtype", "float64", "float32", "int64", "int32"},
        "numpy.core.multiarray": {"_reconstruct", "scalar"},
        "builtins": {"dict", "list", "tuple", "str", "bytes", "int", "float", "bool"},
        "_codecs": {"encode"},
    }

    def find_class(self, module, name):
        # Allow numpy reconstruction
        if module.startswith("numpy"):
            import numpy
            if hasattr(numpy, name):
                return getattr(numpy, name)
            parts = module.split(".")
            obj = numpy
            for part in parts[1:]:
                if hasattr(obj, part):
                    obj = getattr(obj, part)
            if hasattr(obj, name):
                return getattr(obj, name)

        if module in self.ALLOWED_MODULES:
            allowed = self.ALLOWED_MODULES[module]
            if name in allowed:
                return super().find_class(module, name)

        raise ValueError(
            f"Restricted unpickling: '{module}.{name}' is not allowed. "
            "Only numpy arrays and basic Python types are permitted."
        )


def _safe_load_pickle(filepath, encoding="latin1"):
    """Safely load a pickle file with restricted unpickling."""
    with open(filepath, "rb") as f:
        return _RestrictedUnpickler(f, encoding=encoding).load()


class MANO(smplx.MANOLayer):
    def __init__(self, *args, joint_regressor_extra: str | None = None, **kwargs):
        """
        Extension of the official MANO implementation to support more joints.
        Args:
            Same as MANOLayer.
            joint_regressor_extra (str): Path to extra joint regressor.
        """
        super().__init__(*args, **kwargs)
        mano_to_openpose = [
            0,
            13,
            14,
            15,
            16,
            1,
            2,
            3,
            17,
            4,
            5,
            6,
            18,
            10,
            11,
            12,
            19,
            7,
            8,
            9,
            20,
        ]

        # 2, 3, 5, 4, 1
        if joint_regressor_extra is not None:
            # SECURITY: Use restricted unpickler to prevent arbitrary code execution
            self.register_buffer(
                "joint_regressor_extra",
                torch.tensor(
                    _safe_load_pickle(joint_regressor_extra, encoding="latin1"),
                    dtype=torch.float32,
                ),
            )
        self.register_buffer(
            "extra_joints_idxs", to_tensor(list(vertex_ids["mano"].values()), dtype=torch.long)
        )
        self.register_buffer("joint_map", torch.tensor(mano_to_openpose, dtype=torch.long))

    def forward(self, *args, **kwargs) -> MANOOutput:
        """
        Run forward pass. Same as MANO and also append an extra set of joints if joint_regressor_extra is specified.
        """
        mano_output = super().forward(*args, **kwargs)
        extra_joints = torch.index_select(mano_output.vertices, 1, self.extra_joints_idxs)
        joints = torch.cat([mano_output.joints, extra_joints], dim=1)
        joints = joints[:, self.joint_map, :]
        if hasattr(self, "joint_regressor_extra"):
            extra_joints = vertices2joints(self.joint_regressor_extra, mano_output.vertices)
            joints = torch.cat([joints, extra_joints], dim=1)
        mano_output.joints = joints
        return mano_output
