"""HDF5 dataset loader for RoboMimic / robosuite demonstrations."""

from __future__ import annotations

import json
import logging
from pathlib import Path

import numpy as np
import pandas as pd

from orbit.analyzer.dataset_loader import DatasetInfo
from orbit.loaders.base import BaseLoader, normalize_dataframe, sample_episode_indices

logger = logging.getLogger(__name__)


def _check_h5py():
    try:
        import h5py  # noqa: F401

        return True
    except ImportError:
        return False


class HDF5Loader(BaseLoader):
    """Load RoboMimic / robosuite HDF5 demonstration files."""

    def can_load(self, source: str) -> bool:
        """True if source is an HDF5 file containing data/demo_0."""
        p = Path(source).expanduser()
        if not p.is_file():
            return False
        if p.suffix not in (".hdf5", ".h5"):
            return False
        if not _check_h5py():
            return False
        # Quick check for RoboMimic structure
        import h5py

        try:
            with h5py.File(str(p), "r") as f:
                return "data" in f and "demo_0" in f["data"]
        except Exception as e:
            logger.debug("HDF5 probe failed for %s: %s", source, e)
            return False

    def load_info(self, source: str) -> DatasetInfo:
        """Load metadata from HDF5 file."""
        if not _check_h5py():
            raise ImportError(
                "HDF5 format detected but 'h5py' is not installed.\n"
                "  Install with: pip install orbit-robotics[hdf5]"
            )

        import h5py

        p = Path(source).expanduser().resolve()

        with h5py.File(str(p), "r") as f:
            if "data" not in f:
                raise ValueError(f"Not a RoboMimic HDF5 file: {source} (no 'data' group)")

            data_group = f["data"]
            demo_names = sorted(
                [k for k in data_group.keys() if k.startswith("demo_")]
            )

            num_episodes = len(demo_names)
            num_frames = 0
            action_dim = 0
            state_dim = 0

            for name in demo_names:
                demo = data_group[name]
                if "actions" in demo:
                    n = demo["actions"].shape[0]
                    num_frames += n
                    action_dim = demo["actions"].shape[1]
                if "states" in demo and state_dim == 0:
                    state_dim = demo["states"].shape[1]

            # Try to get FPS and robot type from env_args
            fps = 20.0  # RoboMimic default
            robot_type = "unknown"

            if "data" in f.attrs:
                try:
                    env_args = json.loads(f.attrs.get("env_args", "{}"))
                    robot_type = (
                        env_args.get("env_kwargs", {}).get("robots", [robot_type])
                    )
                    if isinstance(robot_type, list):
                        robot_type = robot_type[0] if robot_type else "unknown"
                except (json.JSONDecodeError, TypeError):
                    pass
            elif "env_args" in f.attrs:
                try:
                    env_args = json.loads(f.attrs["env_args"])
                    robot_type = (
                        env_args.get("env_kwargs", {}).get("robots", [robot_type])
                    )
                    if isinstance(robot_type, list):
                        robot_type = robot_type[0] if robot_type else "unknown"
                except (json.JSONDecodeError, TypeError):
                    pass

        shapes: dict[str, list[int]] = {}
        if action_dim > 0:
            shapes["action"] = [action_dim]
        if state_dim > 0:
            shapes["state"] = [state_dim]

        return DatasetInfo(
            repo_id=str(p),
            num_episodes=num_episodes,
            num_frames=num_frames,
            fps=fps,
            robot_type=robot_type,
            camera_names=[],
            video_keys=[],
            shapes=shapes,
            duration_seconds=num_frames / fps if fps > 0 else 0.0,
        )

    def load_episodes(
        self,
        source: str,
        episode_indices: list[int] | None = None,
        max_episodes: int | None = None,
    ) -> pd.DataFrame:
        """Load episode data from HDF5 file."""
        if not _check_h5py():
            raise ImportError(
                "HDF5 format detected but 'h5py' is not installed.\n"
                "  Install with: pip install orbit-robotics[hdf5]"
            )

        import h5py

        if max_episodes is None:
            max_episodes = 100

        p = Path(source).expanduser().resolve()

        rows: list[dict] = []

        with h5py.File(str(p), "r") as f:
            data_group = f["data"]
            demo_names = sorted(
                [k for k in data_group.keys() if k.startswith("demo_")]
            )

            # Determine FPS
            fps = 20.0
            if "env_args" in f.attrs:
                try:
                    env_args = json.loads(f.attrs["env_args"])
                    fps = float(env_args.get("env_kwargs", {}).get("control_freq", fps))
                except (json.JSONDecodeError, TypeError, ValueError):
                    pass

            # Select episodes
            if episode_indices is not None:
                target_demos = [
                    demo_names[i] for i in episode_indices if i < len(demo_names)
                ]
            elif len(demo_names) > max_episodes:
                sampled = sample_episode_indices(len(demo_names), max_episodes)
                target_demos = [demo_names[i] for i in sampled]
            else:
                target_demos = demo_names

            # Discover action/state keys from the first demo.
            # RoboMimic uses "actions"/"states", but other formats may
            # use "action", "delta_action", "joint_command", etc.
            _ACTION_KEYS = [
                "actions", "action", "delta_action", "joint_command",
                "joint_cmd", "end_pose", "delta_ee", "ee_command",
            ]
            _STATE_KEYS = [
                "states", "state", "obs/joint_pos", "obs/state",
                "obs/robot_state", "obs/proprioception", "obs/qpos",
            ]
            action_key_found: str | None = None
            state_key_found: str | None = None
            if target_demos:
                first_demo = data_group[target_demos[0]]
                all_keys = list(first_demo.keys())
                # Also check obs/ subgroup keys
                if "obs" in first_demo:
                    all_keys += [f"obs/{k}" for k in first_demo["obs"].keys()]
                for k in _ACTION_KEYS:
                    if k in all_keys or k in first_demo:
                        action_key_found = k
                        break
                for k in _STATE_KEYS:
                    if k in all_keys or k in first_demo:
                        state_key_found = k
                        break

            for demo_name in target_demos:
                demo = data_group[demo_name]
                # Extract demo index number
                ep_idx = int(demo_name.replace("demo_", ""))

                actions = None
                states = None
                if action_key_found:
                    try:
                        obj = demo[action_key_found]
                        actions = obj[:] if hasattr(obj, "shape") else None
                    except KeyError:
                        pass
                if state_key_found:
                    try:
                        obj = demo[state_key_found]
                        states = obj[:] if hasattr(obj, "shape") else None
                    except KeyError:
                        pass

                n_frames = actions.shape[0] if actions is not None else 0

                for frame_idx in range(n_frames):
                    rows.append(
                        {
                            "episode_index": ep_idx,
                            "frame_index": frame_idx,
                            "action": (
                                actions[frame_idx].tolist()
                                if actions is not None
                                else None
                            ),
                            "state": (
                                states[frame_idx].tolist()
                                if states is not None
                                else None
                            ),
                            "timestamp": frame_idx / fps,
                        }
                    )

        if not rows:
            return normalize_dataframe(pd.DataFrame())

        return normalize_dataframe(pd.DataFrame(rows))
