"""Load datasets from any supported format (HuggingFace, local, HDF5, RLDS, ROS bags, etc.).

This module provides backward-compatible functions that dispatch to the
pluggable loader system in orbit.loaders. All existing code that imports
load_dataset_info() and load_episode_data() continues to work unchanged.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import pandas as pd

logger = logging.getLogger(__name__)


@dataclass
class DatasetInfo:
    """Metadata about a dataset."""

    repo_id: str
    num_episodes: int
    num_frames: int
    fps: float
    robot_type: str
    camera_names: list[str]
    video_keys: list[str]
    shapes: dict[str, list[int]] = field(default_factory=dict)
    duration_seconds: float = 0.0


def load_dataset_info(source: str, format: str | None = None, **kwargs) -> DatasetInfo:
    """Load dataset metadata — auto-detects format.

    Supports: HuggingFace repo IDs, local LeRobot dirs, RLDS, HDF5, ROS bags, raw dirs.

    Args:
        source: HuggingFace repo ID (e.g. "lerobot/pusht"), local path, or file.
        format: Explicit format override (e.g. "hdf5", "rlds", "rosbag").
        **kwargs: Extra arguments passed to the loader (e.g. action_topic for ROS bags).

    Returns:
        DatasetInfo with parsed metadata.

    Raises:
        ValueError: If the source can't be loaded or format is unknown.
        ConnectionError: On network failures (HuggingFace only).
    """
    from orbit.loaders import detect_and_load_info

    return detect_and_load_info(source, format=format, **kwargs)


def load_episode_data(
    source: str,
    episode_indices: list[int] | None = None,
    max_episodes: int | None = None,
    format: str | None = None,
    **kwargs,
) -> pd.DataFrame:
    """Load episode data — auto-detects format.

    Args:
        source: HuggingFace repo ID, local path, or file.
        episode_indices: Specific episodes to load. If None and the dataset
            has >100 episodes, 100 episodes are sampled evenly.
        max_episodes: Override max episodes to sample. Defaults to 100.
        format: Explicit format override.
        **kwargs: Extra arguments passed to the loader.

    Returns:
        DataFrame with columns: episode_index, frame_index, action, state,
        timestamp.

    Raises:
        ValueError: If the source can't be loaded.
        ConnectionError: On network failures (HuggingFace only).
    """
    from orbit.loaders import detect_and_load_episodes

    return detect_and_load_episodes(
        source,
        episode_indices=episode_indices,
        max_episodes=max_episodes,
        format=format,
        **kwargs,
    )


# Keep the sampling utility available for backward compatibility
def _sample_episode_indices(total: int, max_samples: int = 100) -> list[int]:
    """Select evenly-spaced episode indices when total exceeds max_samples."""
    if total <= max_samples:
        return list(range(total))
    step = total / max_samples
    return [int(i * step) for i in range(max_samples)]
