"""
H5 file utilities for loading and saving pose features.

Supports WiLoR, SMPL-X, MediaPipe and other H5 formats with sparse frame_idx indexing.
"""

from pathlib import Path
from typing import TYPE_CHECKING, Any

import numpy as np

# Import hdf5plugin early to register LZ4 and other compression filters
# This must happen before h5py opens any files with these compressions
try:
    import hdf5plugin  # noqa: F401
except ImportError:
    pass  # hdf5plugin is optional, but needed for LZ4-compressed files

from sltk.exceptions import FormatError

if TYPE_CHECKING:
    from sltk.data.core import PoseSequence


def numpy_to_python(value: Any) -> Any:
    """
    Convert numpy types to Python native types for JSON serialization.

    Handles numpy arrays, scalars (int, float, bool), and bytes.
    Recursively converts dictionaries and lists.

    Args:
        value: Any value that may be a numpy type

    Returns:
        Python native type equivalent
    """
    if isinstance(value, np.ndarray):
        return value.tolist()
    elif isinstance(value, (np.integer, np.floating, np.bool_)):
        return value.item()
    elif isinstance(value, bytes):
        return value.decode("utf-8")
    elif isinstance(value, dict):
        return {k: numpy_to_python(v) for k, v in value.items()}
    elif isinstance(value, (list, tuple)):
        return [numpy_to_python(v) for v in value]
    return value


def load_h5_features(
    path: str | Path,
    features: list[str] | None = None,
    frames: list[int] | None = None,
) -> dict[str, np.ndarray]:
    """
    Load features from H5 file.

    Args:
        path: Path to H5 file
        features: Optional list of feature keys to load. If None, loads all.
        frames: Optional list of frame indices to load. If None, loads all.

    Returns:
        Dictionary mapping feature names to arrays

    Example:
        >>> data = load_h5_features("video_wilor.h5", features=["kpts_3d", "mano/hand_pose"])
        >>> print(data["kpts_3d"].shape)
        (1000, 21, 3)
    """
    import h5py

    path = Path(path)
    result = {}

    with h5py.File(path, "r") as f:
        # Determine which keys to load
        if features is None:
            features = _get_all_keys(f)

        frame_idx = f["frame_idx"][:] if "frame_idx" in f else None

        for key in features:
            if key not in f and "/" not in key:
                # Try to find in groups
                for group_name in f.keys():
                    if isinstance(f[group_name], h5py.Group):
                        full_key = f"{group_name}/{key}"
                        if full_key in f:
                            key = full_key
                            break

            if key in f:
                data = f[key][:]

                # If frames specified and we have frame_idx, filter
                if frames is not None and frame_idx is not None:
                    data = _extract_frames(data, frame_idx, frames)

                result[key.replace("/", "_")] = data

    return result


def get_h5_metadata(path: str | Path) -> dict[str, Any]:
    """
    Get metadata from H5 file attributes.

    Args:
        path: Path to H5 file

    Returns:
        Dictionary of metadata attributes

    Example:
        >>> meta = get_h5_metadata("video_wilor.h5")
        >>> print(meta["fps"], meta["num_frames"])
        25.0 1000
    """
    import h5py

    path = Path(path)
    metadata = {}

    with h5py.File(path, "r") as f:
        # Get file-level attributes with numpy type conversion
        for key in f.attrs.keys():
            metadata[key] = numpy_to_python(f.attrs[key])

        # Add computed metadata
        if "frame_idx" in f:
            metadata["num_frames"] = f["frame_idx"].shape[0]

        # Detect format type
        if "mano" in f:
            metadata["format"] = "wilor"
        elif "smplx" in f:
            metadata["format"] = "smplx"

        # List available features
        metadata["features"] = _get_all_keys(f)

    return metadata


def get_frame_detections(
    h5_file,
    frame_num: int,
    frame_idx: np.ndarray,
) -> tuple[int, int]:
    """
    Get detection indices for a specific frame using sparse frame_idx.

    Args:
        h5_file: Open h5py File object
        frame_num: Frame number to query
        frame_idx: The frame_idx array from the H5 file

    Returns:
        Tuple of (start_idx, end_idx) for slicing detection arrays.
        Returns (-1, -1) if no detections for this frame.

    Example:
        >>> with h5py.File("video.h5", "r") as f:
        ...     frame_idx = f["frame_idx"][:]
        ...     start, end = get_frame_detections(f, 100, frame_idx)
        ...     if start >= 0:
        ...         kpts = f["kpts_3d"][start:end]
    """
    if frame_num >= len(frame_idx):
        return (-1, -1)

    start_idx, count = frame_idx[frame_num]

    if start_idx < 0 or count == 0:
        return (-1, -1)

    return (int(start_idx), int(start_idx + count))


def _get_all_keys(h5_file, prefix: str = "") -> list[str]:
    """Recursively get all dataset keys from H5 file."""
    import h5py

    keys = []
    for key in h5_file.keys():
        full_key = f"{prefix}/{key}" if prefix else key
        if isinstance(h5_file[key], h5py.Group):
            keys.extend(_get_all_keys(h5_file[key], full_key))
        else:
            keys.append(full_key)
    return keys


def _extract_frames(
    data: np.ndarray,
    frame_idx: np.ndarray,
    frames: list[int],
) -> np.ndarray:
    """
    Extract data for specific frames using frame_idx mapping.

    This reconstructs frame-aligned data from sparse detection arrays.
    """
    # Determine output shape
    # For detection arrays, we need to handle variable detections per frame
    # For simplicity, take first detection per frame

    result_list = []
    for frame_num in frames:
        if frame_num >= len(frame_idx):
            continue

        start_idx, count = frame_idx[frame_num]
        if start_idx >= 0 and count > 0:
            result_list.append(data[start_idx])
        else:
            # No detection - use zeros
            result_list.append(np.zeros_like(data[0]))

    if result_list:
        return np.stack(result_list)
    return np.array([])


def save_poses_h5(
    poses: "PoseSequence",
    path: str | Path,
    metadata: dict[str, Any] | None = None,
    compress: bool = True,
) -> Path:
    """
    Save PoseSequence to H5 file.

    Args:
        poses: PoseSequence to save
        path: Output file path
        metadata: Optional additional metadata
        compress: Whether to use compression

    Returns:
        Path to saved file

    Example:
        >>> from sltk.data.core import PoseSequence
        >>> poses = PoseSequence(data=np.zeros((100, 33, 3)), fps=25.0, format="mediapipe")
        >>> save_poses_h5(poses, "output.h5")
    """
    import logging

    import h5py

    logger = logging.getLogger(__name__)

    # Determine compression settings
    compression_kwargs = {}
    if compress:
        try:
            import hdf5plugin

            # Use hdf5plugin filter directly - modern API passes filter object
            # The LZ4 filter object can be passed directly to create_dataset via **kwargs
            lz4_filter = hdf5plugin.LZ4()

            # hdf5plugin >= 4.0 returns filter objects that should be unpacked
            # Check if filter_options is a dict (newer API) or tuple/int (older API)
            if hasattr(lz4_filter, "filter_id"):
                filter_opts = lz4_filter.filter_options
                # Ensure filter_options is properly formatted for h5py
                if isinstance(filter_opts, dict):
                    # Newer hdf5plugin returns dict - extract the tuple
                    compression_kwargs = {
                        "compression": lz4_filter.filter_id,
                        "compression_opts": tuple(filter_opts.values()) if filter_opts else None,
                    }
                elif filter_opts is None or isinstance(filter_opts, (tuple, int)):
                    # Older format or no options
                    compression_kwargs = {
                        "compression": lz4_filter.filter_id,
                        "compression_opts": filter_opts,
                    }
                else:
                    # Unknown format - log warning and use default
                    logger.warning(
                        f"Unexpected filter_options type: {type(filter_opts)}, using gzip"
                    )
                    compression_kwargs = {"compression": "gzip", "compression_opts": 4}
            else:
                # Very old hdf5plugin without filter_id attribute
                compression_kwargs = {"compression": "gzip", "compression_opts": 4}

        except (ImportError, AttributeError, TypeError) as e:
            # Fallback to gzip if hdf5plugin not available or has issues
            logger.debug(f"hdf5plugin not available or error ({e}), using gzip compression")
            compression_kwargs = {"compression": "gzip", "compression_opts": 4}

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    with h5py.File(path, "w") as f:
        # Save pose data
        f.create_dataset("poses", data=poses.data, **compression_kwargs)

        # Save confidence if available
        if poses.confidence is not None:
            f.create_dataset("confidence", data=poses.confidence, **compression_kwargs)

        # Create frame_idx for compatibility (dense format)
        num_frames = poses.num_frames
        frame_idx = np.stack(
            [
                np.arange(num_frames),
                np.ones(num_frames, dtype=np.int32),
            ],
            axis=1,
        )
        f.create_dataset("frame_idx", data=frame_idx)

        # Save attributes
        f.attrs["fps"] = poses.fps
        f.attrs["format"] = poses.format
        f.attrs["num_frames"] = poses.num_frames
        f.attrs["num_keypoints"] = poses.num_keypoints
        f.attrs["num_coords"] = poses.num_coords

        # Save additional metadata
        if metadata:
            for key, value in metadata.items():
                if isinstance(value, (str, int, float, bool)):
                    f.attrs[key] = value
                elif isinstance(value, (list, dict)):
                    import json

                    f.attrs[key] = json.dumps(value)

        # Save pose-specific metadata
        if poses.metadata:
            meta_group = f.create_group("metadata")
            for key, value in poses.metadata.items():
                if isinstance(value, np.ndarray):
                    meta_group.create_dataset(key, data=value, **compression_kwargs)
                elif isinstance(value, (str, int, float, bool)):
                    meta_group.attrs[key] = value
                elif isinstance(value, (list, dict)):
                    import json

                    meta_group.attrs[key] = json.dumps(value)

    return path


def load_poses_h5(
    path: str | Path,
    fps: float | None = None,
) -> "PoseSequence":
    """
    Load PoseSequence from H5 file.

    Args:
        path: Path to H5 file
        fps: Override FPS (uses file attribute if None)

    Returns:
        PoseSequence instance

    Example:
        >>> poses = load_poses_h5("video_mediapipe.h5")
        >>> print(poses.num_frames, poses.fps)
        1000 25.0
    """
    import h5py

    from sltk.data.core import PoseSequence

    path = Path(path)

    with h5py.File(path, "r") as f:
        # Load pose data
        if "poses" in f:
            data = f["poses"][:]
        elif "kpts_3d" in f:
            # WiLoR format - need to handle sparse
            data = _load_sparse_poses(f)
        else:
            raise FormatError(f"No pose data found in {path}. Expected 'poses' or 'kpts_3d' key.")

        # Load confidence if available
        confidence = None
        if "confidence" in f:
            confidence = f["confidence"][:]

        # Get attributes
        file_fps = f.attrs.get("fps", 25.0)
        if isinstance(file_fps, bytes):
            file_fps = float(file_fps.decode())
        fps = fps or float(file_fps)

        pose_format = f.attrs.get("format", "unknown")
        if isinstance(pose_format, bytes):
            pose_format = pose_format.decode()

        # Load metadata
        metadata = {}
        if "metadata" in f:
            meta_group = f["metadata"]
            for key in meta_group.attrs.keys():
                value = meta_group.attrs[key]
                if isinstance(value, bytes):
                    value = value.decode()
                    try:
                        import json

                        value = json.loads(value)
                    except (json.JSONDecodeError, TypeError):
                        pass
                metadata[key] = value

    return PoseSequence(
        data=data,
        fps=fps,
        format=pose_format,
        confidence=confidence,
        metadata=metadata,
    )


def _load_sparse_poses(f) -> np.ndarray:
    """Load poses from sparse WiLoR/SMPLX format."""
    kpts_3d = f["kpts_3d"][:]
    frame_idx = f["frame_idx"][:] if "frame_idx" in f else None

    if frame_idx is None:
        return kpts_3d

    # Reconstruct dense array
    num_frames = len(frame_idx)
    num_kpts = kpts_3d.shape[1] if kpts_3d.ndim > 1 else 21
    num_coords = kpts_3d.shape[2] if kpts_3d.ndim > 2 else 3

    # Handle dual hand format (left + right)
    if "right" in f:
        right_mask = f["right"][:]
        # Create (T, 42, 3) for both hands
        data = np.zeros((num_frames, 42, num_coords), dtype=np.float32)

        for frame in range(num_frames):
            start, count = frame_idx[frame]
            start, count = int(start), int(count)
            if start >= 0 and count > 0:
                for i in range(min(count, 2)):  # Max 2 hands
                    det_idx = start + i
                    if det_idx < len(kpts_3d):
                        is_right = right_mask[det_idx] if det_idx < len(right_mask) else True
                        offset = 21 if is_right else 0
                        data[frame, offset : offset + 21] = kpts_3d[det_idx]
    else:
        # Single detection per frame format
        data = np.zeros((num_frames, num_kpts, num_coords), dtype=np.float32)

        for frame in range(num_frames):
            start, count = frame_idx[frame]
            start, count = int(start), int(count)
            if start >= 0 and count > 0:
                data[frame] = kpts_3d[start]

    return data
