"""
Pose format loading, saving, and conversion.

Supports:
- WiLoR H5 files (MANO hand model)
- NLF/SMPL-X H5 files (full body model)
- Generic numpy files (.npy, .npz)

New formats can be registered using the @register_format decorator.
"""

from __future__ import annotations

from collections.abc import Callable
from pathlib import Path
from typing import Any

import numpy as np

from sltk.data.core import PoseSequence
from sltk.exceptions import FeatureNotFoundError, FormatError, UnsupportedFormatError

# Format registry
_FORMAT_LOADERS: dict[str, Callable] = {}
_FORMAT_SAVERS: dict[str, Callable] = {}
_FORMAT_INFO: dict[str, dict[str, Any]] = {}


def register_format(name: str, **info):
    """
    Decorator to register a format loader/saver.

    Example:
        @register_format("my_format", num_keypoints=21, description="Custom format")
        class MyFormat:
            @staticmethod
            def load(path, fps, **kwargs) -> PoseSequence:
                ...

            @staticmethod
            def save(poses: PoseSequence, path, **kwargs):
                ...
    """

    def decorator(cls):
        if hasattr(cls, "load"):
            _FORMAT_LOADERS[name] = cls.load
        if hasattr(cls, "save"):
            _FORMAT_SAVERS[name] = cls.save
        _FORMAT_INFO[name] = info
        return cls

    return decorator


def load_poses(
    path: str | Path,
    format: str,
    fps: float,
    **kwargs,
) -> PoseSequence:
    """
    Load poses from file.

    Args:
        path: Path to pose file (.h5, .npy, .npz, .pkl)
        format: Format identifier ('wilor', 'smplx', 'mediapipe')
        fps: Frames per second
        **kwargs: Format-specific arguments

    Returns:
        PoseSequence
    """
    path = Path(path)

    if format in _FORMAT_LOADERS:
        return _FORMAT_LOADERS[format](path, fps, **kwargs)

    # Try generic numpy loading as fallback
    return _load_generic(path, format, fps, **kwargs)


def save_poses(poses: PoseSequence, path: str | Path, **kwargs):
    """Save poses to file."""
    path = Path(path)

    if poses.format in _FORMAT_SAVERS:
        _FORMAT_SAVERS[poses.format](poses, path, **kwargs)
    else:
        # Generic numpy save
        _save_generic(poses, path, **kwargs)


def convert_poses(poses: PoseSequence, target_format: str) -> PoseSequence:
    """
    Convert poses to a different format.

    Currently only supports relabelling (same format family or explicit opt-in).
    Full keypoint mapping between different pose formats (e.g. MediaPipe -> SMPL-X)
    is not yet implemented.

    Args:
        poses: Source pose sequence.
        target_format: Target format string.

    Returns:
        A new PoseSequence with the target format label and copied data.

    Raises:
        NotImplementedError: If source and target formats require keypoint remapping
            (i.e. they belong to different format families).
    """
    # If same format, just return a copy
    if poses.format == target_format:
        return poses.copy()

    # Determine format families for safety check
    # Formats in the same family share keypoint layout and can be relabelled
    _FAMILY = {
        "mediapipe": "mediapipe",
        "mediapipe_holistic": "mediapipe",
        "wilor": "mano",
        "mano": "mano",
        "smplx": "smplx",
        "nlf": "smplx",
        "generic": "generic",
    }
    src_family = _FAMILY.get(poses.format, poses.format)
    tgt_family = _FAMILY.get(target_format, target_format)

    if src_family != tgt_family:
        raise NotImplementedError(
            f"Conversion from '{poses.format}' to '{target_format}' requires "
            f"keypoint remapping between format families '{src_family}' -> '{tgt_family}', "
            f"which is not yet implemented."
        )

    # Same family: safe to relabel without keypoint remapping
    return PoseSequence(
        data=poses.data.copy(),
        fps=poses.fps,
        format=target_format,
        confidence=poses.confidence.copy() if poses.confidence is not None else None,
        metadata=poses.metadata.copy(),
    )


def list_formats() -> dict[str, dict[str, Any]]:
    """List registered formats and their capabilities."""
    result = {}
    all_formats = set(_FORMAT_LOADERS.keys()) | set(_FORMAT_SAVERS.keys())

    for fmt in sorted(all_formats):
        result[fmt] = {
            "can_load": fmt in _FORMAT_LOADERS,
            "can_save": fmt in _FORMAT_SAVERS,
            **_FORMAT_INFO.get(fmt, {}),
        }

    return result


def _load_generic(path: Path, format: str, fps: float, **kwargs) -> PoseSequence:
    """Generic loader for numpy-based formats."""
    suffix = path.suffix.lower()

    if suffix == ".npy":
        data = np.load(path)
    elif suffix == ".npz":
        loaded = np.load(path)
        data = loaded.get("poses", loaded.get("data", loaded[loaded.files[0]]))
    elif suffix == ".pkl":
        # SECURITY: Pickle deserialization is disabled due to RCE vulnerability.
        # Pickle files can execute arbitrary code during deserialization.
        # Convert pickle files to .npy or .npz format instead:
        #   import pickle, numpy as np
        #   data = pickle.load(open('file.pkl', 'rb'))
        #   np.save('file.npy', data['poses'])  # or data if it's already an array
        raise UnsupportedFormatError(
            ".pkl",
            supported_formats=[".npy", ".npz", ".h5"],
        )
    else:
        raise UnsupportedFormatError(suffix, supported_formats=[".npy", ".npz", ".h5", ".pkl"])

    # Ensure 3D if needed
    if data.ndim == 2:
        # Assume (T, N*C) -> reshape to (T, N, 3) assuming 3 coords
        num_coords = 3
        if data.shape[1] % num_coords == 0:
            num_keypoints = data.shape[1] // num_coords
            data = data.reshape(data.shape[0], num_keypoints, num_coords)

    return PoseSequence(data=data, fps=fps, format=format, **kwargs)


def _save_generic(poses: PoseSequence, path: Path, **kwargs):
    """Generic saver."""
    suffix = path.suffix.lower()

    if suffix == ".npy":
        np.save(path, poses.data)
    elif suffix == ".npz":
        np.savez(path, poses=poses.data, fps=poses.fps)
    else:
        raise UnsupportedFormatError(suffix, supported_formats=[".npy", ".npz"])


# ============================================================================
# Built-in format handlers
# ============================================================================


@register_format(
    "wilor",
    description="WiLoR MANO hand model",
    num_keypoints=21,
    keypoint_type="hand",
)
class WiLoRFormat:
    """
    WiLoR H5 format (MANO hand model).

    H5 Structure:
        - frame_idx: (num_frames, 2) - sparse index [start_idx, count]
        - kpts_2d: (N, 21, 2) - 2D keypoints
        - kpts_3d: (N, 21, 3) - 3D keypoints
        - right: (N,) - handedness flag
        - mano/global_orient: (N, 1, 3, 3) - rotation matrices
        - mano/hand_pose: (N, 15, 3, 3) - joint rotations
        - mano/betas: (N, 10) - shape parameters
    """

    @staticmethod
    def load(
        path: Path,
        fps: float,
        keypoints: str = "3d",
        hand: str = "both",
        **kwargs,
    ) -> PoseSequence:
        """
        Load WiLoR poses from H5 file.

        Args:
            path: Path to H5 file
            fps: Frames per second
            keypoints: '2d' or '3d'
            hand: 'left', 'right', or 'both'
        """
        try:
            import hdf5plugin  # For LZ4 compression support
        except ImportError:
            pass
        import h5py

        with h5py.File(path, "r") as f:
            frame_idx = f["frame_idx"][:]
            num_frames = frame_idx.shape[0]

            kpts_key = f"kpts_{keypoints}"
            if kpts_key not in f:
                raise FeatureNotFoundError(kpts_key)

            kpts = f[kpts_key][:]
            right_flags = f["right"][:] if "right" in f else None

            # Get metadata with numpy type conversion
            from sltk.io.h5 import numpy_to_python

            metadata = {
                "source_path": str(path),
                "num_detections": int(kpts.shape[0]),
            }
            for attr in ["video_name", "fps", "resolution"]:
                if attr in f.attrs:
                    metadata[attr] = numpy_to_python(f.attrs[attr])

        # Reconstruct frame-aligned data using optimized approach
        # Use frame_idx to map detections back to frames
        num_keypoints = kpts.shape[1]
        num_coords = kpts.shape[2]

        # Initialize output array
        if hand == "both":
            # Stack left and right hands: (T, 42, C)
            data = np.zeros((num_frames, num_keypoints * 2, num_coords), dtype=np.float32)
        else:
            data = np.zeros((num_frames, num_keypoints, num_coords), dtype=np.float32)

        confidence = np.zeros(num_frames, dtype=np.float32)

        # Pre-compute valid frames mask for efficiency
        start_indices = frame_idx[:, 0]
        counts = frame_idx[:, 1]
        valid_mask = (start_indices >= 0) & (counts > 0)
        valid_frames = np.where(valid_mask)[0]

        # Vectorize confidence assignment
        confidence[valid_frames] = 1.0

        if hand == "both" and right_flags is not None:
            # Need to separate left/right - can't fully vectorize due to variable counts
            for frame_num in valid_frames:
                start_idx = int(start_indices[frame_num])
                count = int(counts[frame_num])
                end_idx = start_idx + count
                frame_kpts = kpts[start_idx:end_idx]
                frame_right = right_flags[start_idx:end_idx]

                # Use vectorized assignment where possible
                right_mask = frame_right.astype(bool)
                left_mask = ~right_mask

                if right_mask.any():
                    data[frame_num, :num_keypoints] = frame_kpts[right_mask][0]
                if left_mask.any():
                    data[frame_num, num_keypoints:] = frame_kpts[left_mask][0]
        elif len(valid_frames) > 0:
            # Single hand case - can be more vectorized
            # For simplicity and reliability, use efficient loop with pre-computed indices
            for frame_num in valid_frames:
                start_idx = int(start_indices[frame_num])
                count = int(counts[frame_num])
                end_idx = start_idx + count
                frame_kpts = kpts[start_idx:end_idx]

                if right_flags is not None:
                    frame_right = right_flags[start_idx:end_idx]
                    if hand == "right":
                        mask = frame_right.astype(bool)
                    else:  # hand == "left"
                        mask = ~frame_right.astype(bool)
                    filtered = frame_kpts[mask]
                else:
                    filtered = frame_kpts

                if len(filtered) > 0:
                    data[frame_num] = filtered[0]

        return PoseSequence(
            data=data,
            fps=fps,
            format="wilor",
            confidence=confidence,
            metadata=metadata,
        )


@register_format(
    "smplx",
    description="SMPL-X full body model",
    num_keypoints=55,
    keypoint_type="body",
)
class SMPLXFormat:
    """
    SMPL-X/NLF H5 format (full body model).

    H5 Structure:
        - frame_idx: (num_frames, 2) - sparse index
        - joints2d: (N, 55, 2) - 2D joints
        - joints3d: (N, 55, 3) - 3D joints
        - smplx/global_orient: (N, 3) - axis-angle
        - smplx/body_pose: (N, 63) - axis-angle (21 joints × 3)
        - smplx/betas: (N, 10) - shape parameters
    """

    @staticmethod
    def load(
        path: Path,
        fps: float,
        keypoints: str = "3d",
        **kwargs,
    ) -> PoseSequence:
        """
        Load SMPL-X poses from H5 file.

        Args:
            path: Path to H5 file
            fps: Frames per second
            keypoints: '2d' or '3d'
        """
        try:
            import hdf5plugin  # For LZ4 compression support
        except ImportError:
            pass
        import h5py

        with h5py.File(path, "r") as f:
            frame_idx = f["frame_idx"][:]
            num_frames = frame_idx.shape[0]

            joints_key = f"joints{keypoints}"
            if joints_key not in f:
                raise FeatureNotFoundError(joints_key)

            joints = f[joints_key][:]

            # Get uncertainty if available
            uncertainties = None
            if "joint_uncertainties" in f:
                uncertainties = f["joint_uncertainties"][:]

            # Get metadata with numpy type conversion
            from sltk.io.h5 import numpy_to_python

            metadata = {
                "source_path": str(path),
                "num_detections": int(joints.shape[0]),
            }
            for attr in ["video_name", "fps", "resolution", "fov"]:
                if attr in f.attrs:
                    metadata[attr] = numpy_to_python(f.attrs[attr])

        # Reconstruct frame-aligned data using vectorized operations
        num_keypoints = joints.shape[1]
        num_coords = joints.shape[2]
        data = np.zeros((num_frames, num_keypoints, num_coords), dtype=np.float32)
        confidence = np.zeros((num_frames, num_keypoints), dtype=np.float32)

        # Vectorized sparse-to-dense conversion
        # Find valid frames (where we have detections)
        start_indices = frame_idx[:, 0]
        counts = frame_idx[:, 1]
        valid_mask = (start_indices >= 0) & (counts > 0)
        valid_frames = np.where(valid_mask)[0]

        if len(valid_frames) > 0:
            # Get first detection index for each valid frame
            detection_indices = start_indices[valid_frames].astype(np.intp)

            # Vectorized assignment - assign all valid frames at once
            data[valid_frames] = joints[detection_indices]

            if uncertainties is not None:
                # Convert uncertainty to confidence (inverse relationship)
                confidence[valid_frames] = 1.0 / (1.0 + uncertainties[detection_indices])
            else:
                confidence[valid_frames] = 1.0

        return PoseSequence(
            data=data,
            fps=fps,
            format="smplx",
            confidence=confidence,
            metadata=metadata,
        )


@register_format(
    "mediapipe",
    description="MediaPipe Holistic",
    num_keypoints=543,
    keypoint_type="holistic",
)
class MediaPipeFormat:
    """MediaPipe Holistic format (33 pose + 21×2 hands + 468 face)."""

    @staticmethod
    def load(path: Path, fps: float, **kwargs) -> PoseSequence:
        """Load MediaPipe poses from numpy file."""
        data = np.load(path)

        # MediaPipe might be saved in various shapes
        if data.ndim == 2:
            # Try to infer structure
            if data.shape[1] == 543 * 4:
                data = data.reshape(-1, 543, 4)
            elif data.shape[1] == 543 * 3:
                data = data.reshape(-1, 543, 3)

        return PoseSequence(
            data=data,
            fps=fps,
            format="mediapipe",
            metadata={"source_path": str(path)},
        )
