"""Universal dataset loader registry with auto-detection.

Usage:
    from orbit.loaders import detect_and_load_info, detect_and_load_episodes

    info = detect_and_load_info("lerobot/pusht")                # HuggingFace
    info = detect_and_load_info("./my-dataset/")                # Local LeRobot
    info = detect_and_load_info("./demo.hdf5")                  # RoboMimic HDF5
    info = detect_and_load_info("./recording.mcap")             # ROS 2 MCAP
    info = detect_and_load_info("./data/")                      # Raw numpy/csv/parquet

    # Universal loader with intelligent array profiling:
    from orbit.loaders.universal import universal_load
    info, df, profiles = universal_load("lerobot/pusht")
"""

from __future__ import annotations

import logging
from pathlib import Path

import pandas as pd

from orbit.analyzer.dataset_loader import DatasetInfo
from orbit.loaders.base import BaseLoader

logger = logging.getLogger(__name__)

# Lazy-loaded loader instances (created on first use)
_loader_cache: dict[str, BaseLoader] = {}


def _get_loaders() -> list[tuple[str, BaseLoader]]:
    """Return all registered loaders in detection priority order."""
    from orbit.loaders.directory import DirectoryLoader
    from orbit.loaders.hdf5 import HDF5Loader
    from orbit.loaders.lerobot_hub import LeRobotHubLoader
    from orbit.loaders.lerobot_local import LeRobotLocalLoader
    from orbit.loaders.rlds import RLDSLoader
    from orbit.loaders.rosbag import RosBagLoader

    # Order: most specific first
    return [
        ("lerobot_local", LeRobotLocalLoader()),
        ("hdf5", HDF5Loader()),
        ("rosbag", RosBagLoader()),
        ("rlds", RLDSLoader()),
        ("directory", DirectoryLoader()),
        ("lerobot_hub", LeRobotHubLoader()),
    ]


# Map of format names to loader classes for explicit --format override
FORMAT_NAMES: dict[str, type] = {}


def _init_format_names():
    global FORMAT_NAMES
    if FORMAT_NAMES:
        return
    from orbit.loaders.directory import DirectoryLoader
    from orbit.loaders.hdf5 import HDF5Loader
    from orbit.loaders.lerobot_hub import LeRobotHubLoader
    from orbit.loaders.lerobot_local import LeRobotLocalLoader
    from orbit.loaders.rlds import RLDSLoader
    from orbit.loaders.rosbag import RosBagLoader

    FORMAT_NAMES = {
        "lerobot": LeRobotHubLoader,
        "lerobot-hub": LeRobotHubLoader,
        "lerobot-local": LeRobotLocalLoader,
        "rlds": RLDSLoader,
        "hdf5": HDF5Loader,
        "rosbag": RosBagLoader,
        "directory": DirectoryLoader,
    }


def detect_loader(
    source: str,
    format: str | None = None,
    **loader_kwargs,
) -> BaseLoader:
    """Detect the appropriate loader for the given source.

    Args:
        source: Path, file, or HuggingFace repo ID.
        format: Explicit format override (bypasses auto-detection).
        **loader_kwargs: Extra kwargs passed to loader constructor
            (e.g., action_topic, state_topic for RosBagLoader).

    Returns:
        A BaseLoader instance that can handle the source.

    Raises:
        ValueError: If no loader can handle the source.
    """
    # Explicit format override
    if format is not None:
        _init_format_names()
        fmt = format.lower().strip()
        if fmt == "auto":
            format = None  # Fall through to auto-detection
        elif fmt in FORMAT_NAMES:
            return FORMAT_NAMES[fmt](**loader_kwargs)
        else:
            raise ValueError(
                f"Unknown format '{format}'. "
                f"Supported formats: {', '.join(sorted(FORMAT_NAMES.keys()))}"
            )

    # Auto-detection: try each loader in priority order
    for name, loader in _get_loaders():
        try:
            if loader.can_load(source):
                # If loader needs kwargs (e.g. RosBagLoader), reinstantiate
                if loader_kwargs and hasattr(loader.__class__.__init__, '__code__'):
                    try:
                        return loader.__class__(**loader_kwargs)
                    except TypeError:
                        pass
                return loader
        except Exception:
            # If a loader's can_load raises (e.g. missing dependency), skip it
            continue

    # Provide helpful error message
    p = Path(source).expanduser()
    if p.exists():
        if p.is_dir():
            contents = list(p.iterdir())[:10]
            file_types = set(f.suffix for f in contents if f.is_file())
            raise ValueError(
                f"Could not detect dataset format in '{source}'.\n"
                f"  Directory contains: {file_types or 'no files'}\n"
                f"  Supported formats: LeRobot (meta/info.json), RLDS (.tfrecord), "
                f"numpy (.npz/.npy), CSV, parquet\n"
                f"  Tip: Use --format to specify the format explicitly."
            )
        else:
            raise ValueError(
                f"Could not detect dataset format for '{source}' "
                f"(extension: {p.suffix}).\n"
                f"  Supported file types: .hdf5, .h5, .bag, .mcap\n"
                f"  Tip: Use --format to specify the format explicitly."
            )
    else:
        # Not a local path — maybe a HuggingFace repo ID?
        if "/" in source:
            raise ValueError(
                f"Dataset '{source}' not found locally and could not be loaded "
                f"from HuggingFace Hub.\n"
                f"  Check the repo ID or provide a valid local path."
            )
        raise ValueError(
            f"'{source}' is not a valid path or HuggingFace repo ID.\n"
            f"  Examples:\n"
            f"    orbit analyze lerobot/pusht          (HuggingFace)\n"
            f"    orbit analyze ./my-dataset/           (local LeRobot)\n"
            f"    orbit analyze ./demo.hdf5             (RoboMimic HDF5)\n"
            f"    orbit analyze ./recording.mcap        (ROS 2 bag)"
        )


def detect_and_load_info(
    source: str,
    format: str | None = None,
    **loader_kwargs,
) -> DatasetInfo:
    """Auto-detect format and load dataset metadata.

    Args:
        source: Path, file, or HuggingFace repo ID.
        format: Explicit format override (bypasses auto-detection).
        **loader_kwargs: Extra kwargs passed to loader constructor.

    Returns:
        DatasetInfo with parsed metadata.
    """
    loader = detect_loader(source, format=format, **loader_kwargs)
    return loader.load_info(source)


def detect_and_load_episodes(
    source: str,
    episode_indices: list[int] | None = None,
    max_episodes: int | None = None,
    format: str | None = None,
    **loader_kwargs,
) -> pd.DataFrame:
    """Auto-detect format and load episode data.

    Args:
        source: Path, file, or HuggingFace repo ID.
        episode_indices: Specific episodes to load.
        max_episodes: Maximum episodes to sample.
        format: Explicit format override.
        **loader_kwargs: Extra kwargs passed to loader constructor.
            action_key / state_key are extracted and applied as column
            renames after loading.

    Returns:
        DataFrame with columns: [episode_index, frame_index, action, state, timestamp]
    """
    # Extract column override keys before passing to loader constructor
    action_key = loader_kwargs.pop("action_key", None)
    state_key = loader_kwargs.pop("state_key", None)

    loader = detect_loader(source, format=format, **loader_kwargs)
    df = loader.load_episodes(source, episode_indices=episode_indices, max_episodes=max_episodes)

    # Apply user-specified column renames if the canonical names are missing
    if action_key and "action" not in df.columns and action_key in df.columns:
        df = df.rename(columns={action_key: "action"})
    if state_key and "state" not in df.columns and state_key in df.columns:
        df = df.rename(columns={state_key: "state"})

    return df
