"""LeRobot dataset loader for local directories."""

from __future__ import annotations

import json
import logging
from pathlib import Path

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__)


class LeRobotLocalLoader(BaseLoader):
    """Load LeRobot v2/v3 datasets from local directories."""

    def can_load(self, source: str) -> bool:
        """True if source is a directory containing meta/info.json."""
        p = Path(source).expanduser()
        return p.is_dir() and (p / "meta" / "info.json").exists()

    def load_info(self, source: str) -> DatasetInfo:
        """Load dataset metadata from local meta/info.json."""
        root = Path(source).expanduser().resolve()
        info_path = root / "meta" / "info.json"

        if not info_path.exists():
            raise ValueError(f"Not a LeRobot dataset: {source} (no meta/info.json)")

        with open(info_path) as f:
            info = json.load(f)

        num_episodes = info.get("total_episodes", 0)
        num_frames = info.get("total_frames", 0)
        fps = float(info.get("fps", 0))
        robot_type = info.get("robot_type", "unknown")

        video_keys: list[str] = []
        camera_names: list[str] = []
        shapes: dict[str, list[int]] = {}

        features = info.get("features", {})
        for key, feat in features.items():
            if feat.get("dtype") == "video":
                video_keys.append(key)
                parts = key.split(".")
                camera_names.append(parts[-1])
            shape = feat.get("shape")
            if shape is not None:
                shapes[key] = shape

        # Override episode count from episodes metadata
        ep_parquet = root / "meta" / "episodes" / "chunk-000" / "file-000.parquet"
        ep_jsonl = root / "meta" / "episodes.jsonl"

        if ep_parquet.exists():
            try:
                ep_df = pd.read_parquet(ep_parquet)
                if len(ep_df) > 0:
                    num_episodes = len(ep_df)
            except Exception as e:
                logger.debug("Failed to read episodes parquet: %s", e)
        elif ep_jsonl.exists():
            with open(ep_jsonl) as f:
                lines = [line.strip() for line in f if line.strip()]
            if lines:
                num_episodes = len(lines)

        duration_seconds = num_frames / fps if fps > 0 else 0.0

        return DatasetInfo(
            repo_id=str(root),
            num_episodes=num_episodes,
            num_frames=num_frames,
            fps=fps,
            robot_type=robot_type,
            camera_names=camera_names,
            video_keys=video_keys,
            shapes=shapes,
            duration_seconds=duration_seconds,
        )

    def load_episodes(
        self,
        source: str,
        episode_indices: list[int] | None = None,
        max_episodes: int | None = None,
    ) -> pd.DataFrame:
        """Load episode data from local parquet files."""
        root = Path(source).expanduser().resolve()
        info = self.load_info(source)

        if max_episodes is None:
            max_episodes = 100

        # Discover parquet data files
        data_dir = root / "data"
        if not data_dir.exists():
            logger.warning("No data/ directory in %s", source)
            return normalize_dataframe(pd.DataFrame())

        parquet_files = sorted(data_dir.rglob("*.parquet"))

        if not parquet_files:
            logger.warning("No parquet files in %s/data/", source)
            return normalize_dataframe(pd.DataFrame())

        # Determine format
        is_per_episode = any("episode_" in f.name for f in parquet_files[:5])

        if is_per_episode:
            return self._load_per_episode(
                parquet_files, info, episode_indices, max_episodes
            )
        else:
            return self._load_chunk(parquet_files, info, episode_indices, max_episodes)

    def _load_per_episode(
        self,
        parquet_files: list[Path],
        info: DatasetInfo,
        episode_indices: list[int] | None,
        max_episodes: int,
    ) -> pd.DataFrame:
        """Load v3 per-episode parquet files."""
        import re

        # Build index of episode -> file
        ep_file_map: dict[int, Path] = {}
        for f in parquet_files:
            m = re.search(r"episode_(\d+)\.parquet$", f.name)
            if m:
                ep_file_map[int(m.group(1))] = f

        target_indices = episode_indices
        if target_indices is None:
            all_indices = sorted(ep_file_map.keys())
            if len(all_indices) > max_episodes:
                target_indices = sample_episode_indices(len(all_indices), max_episodes)
                # Map to actual episode indices
                target_indices = [all_indices[i] for i in target_indices if i < len(all_indices)]
            else:
                target_indices = all_indices

        frames: list[pd.DataFrame] = []
        for ep_idx in target_indices:
            if ep_idx not in ep_file_map:
                continue
            df = pd.read_parquet(ep_file_map[ep_idx])
            if "episode_index" not in df.columns:
                df["episode_index"] = ep_idx
            frames.append(df)

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

        return normalize_dataframe(pd.concat(frames, ignore_index=True))

    def _load_chunk(
        self,
        parquet_files: list[Path],
        info: DatasetInfo,
        episode_indices: list[int] | None,
        max_episodes: int,
    ) -> pd.DataFrame:
        """Load v2 chunk-based parquet files."""
        # Limit to first 3 chunks for speed
        files_to_load = parquet_files[:3] if len(parquet_files) > 3 else parquet_files

        frames: list[pd.DataFrame] = []
        for f in files_to_load:
            df = pd.read_parquet(f)
            frames.append(df)

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

        result = pd.concat(frames, ignore_index=True)

        if "episode_index" in result.columns:
            if episode_indices is not None:
                result = result[result["episode_index"].isin(episode_indices)]
            elif info.num_episodes > max_episodes:
                sampled = sample_episode_indices(info.num_episodes, max_episodes)
                result = result[result["episode_index"].isin(sampled)]

        return normalize_dataframe(result)
