"""LeRobot dataset loader for HuggingFace Hub."""

from __future__ import annotations

import json
import logging
import re

import numpy as np
import pandas as pd
from huggingface_hub import hf_hub_download
from huggingface_hub.hf_api import HfApi
from huggingface_hub.utils import (
    EntryNotFoundError,
    HfHubHTTPError,
    RepositoryNotFoundError,
)

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

logger = logging.getLogger(__name__)


class LeRobotHubLoader(BaseLoader):
    """Load LeRobot datasets from HuggingFace Hub."""

    def can_load(self, source: str) -> bool:
        """True if source looks like a HuggingFace repo ID (owner/dataset)."""
        if not source or "/" not in source:
            return False
        # Must not be a local filesystem path
        if source.startswith("/") or source.startswith(".") or source.startswith("~"):
            return False
        # Must look like owner/name (exactly one slash, no path separators beyond that)
        parts = source.split("/")
        return len(parts) == 2 and all(p for p in parts)

    def load_info(self, source: str) -> DatasetInfo:
        """Fetch dataset metadata from HuggingFace Hub."""
        try:
            info_path = hf_hub_download(
                repo_id=source,
                filename="meta/info.json",
                repo_type="dataset",
            )
        except RepositoryNotFoundError:
            raise ValueError(f"Dataset '{source}' not found on HuggingFace Hub")
        except EntryNotFoundError:
            raise ValueError(f"'{source}' does not appear to be a LeRobot dataset")
        except (HfHubHTTPError, Exception) as e:
            if isinstance(e, ValueError):
                raise
            raise ConnectionError(
                f"Network error downloading metadata from '{source}': {e}"
            ) from e

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

        # Download episodes metadata (try v3 parquet first, then v1 jsonl)
        episodes_parquet_path = None
        episodes_path = None
        try:
            episodes_parquet_path = hf_hub_download(
                repo_id=source,
                filename="meta/episodes/chunk-000/file-000.parquet",
                repo_type="dataset",
            )
        except (EntryNotFoundError, Exception) as e:
            logger.debug("Episodes parquet not found, trying jsonl: %s", e)

        if episodes_parquet_path is None:
            try:
                episodes_path = hf_hub_download(
                    repo_id=source,
                    filename="meta/episodes.jsonl",
                    repo_type="dataset",
                )
            except EntryNotFoundError:
                episodes_path = None
            except (HfHubHTTPError, Exception) as e:
                if isinstance(e, (ValueError, EntryNotFoundError)):
                    raise
                raise ConnectionError(
                    f"Network error downloading episodes from '{source}': {e}"
                ) from e

        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 if available
        if episodes_parquet_path is not None:
            try:
                ep_df = pd.read_parquet(episodes_parquet_path)
                if len(ep_df) > 0:
                    num_episodes = len(ep_df)
            except Exception as e:
                logger.debug("Failed to read episodes parquet: %s", e)
        elif episodes_path is not None:
            from pathlib import Path

            ep_path = Path(episodes_path)
            if ep_path.exists():
                with open(ep_path) 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

        # Ensure canonical action/state keys exist in shapes
        normalize_shapes(shapes)

        return DatasetInfo(
            repo_id=source,
            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:
        """Download and load episode parquet files from a LeRobot dataset."""
        info = self.load_info(source)

        if max_episodes is None:
            max_episodes = 100

        # For large datasets, try direct filename construction (fast path)
        target_indices = episode_indices
        if target_indices is None and info.num_episodes > max_episodes:
            target_indices = sample_episode_indices(info.num_episodes, max_episodes)

        if target_indices is not None and info.num_episodes > 200:
            constructed = self._try_direct_episode_download(source, target_indices)
            if constructed:
                logger.info(
                    "Using direct episode download for %d/%d episodes",
                    len(target_indices),
                    info.num_episodes,
                )
                result = self._load_chunk_format(source, constructed, info, None, max_episodes)
                if not result.empty:
                    return normalize_dataframe(result)

        # Fallback: discover files
        data_files = self._discover_data_files(source, max_files=500)

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

        is_chunk_format = any("chunk-" in f for f in data_files[:5])

        if is_chunk_format:
            result = self._load_chunk_format(
                source, data_files, info, episode_indices, max_episodes
            )
        else:
            result = self._load_per_episode_format(
                source, data_files, info, episode_indices, max_episodes
            )

        return normalize_dataframe(result)

    # --- Private helpers (extracted from dataset_loader.py) ---

    def _discover_data_files(
        self, repo_id: str, max_files: int = 500, max_iterations: int = 5000
    ) -> list[str]:
        api = HfApi()
        try:
            all_files = api.list_repo_files(repo_id, repo_type="dataset")
        except Exception:
            return []

        parquet_files = []
        iterations = 0
        for f in all_files:
            iterations += 1
            if iterations > max_iterations:
                break
            if f.startswith("data/") and f.endswith(".parquet"):
                parquet_files.append(f)
                if len(parquet_files) >= max_files:
                    break

        return sorted(parquet_files)

    def _construct_episode_filenames(
        self,
        repo_id: str,
        episode_indices: list[int],
        data_files_sample: list[str],
    ) -> list[str]:
        if not data_files_sample:
            return []

        sample = data_files_sample[0]

        # v3 per-episode: data/chunk-NNN/episode_NNNNNN.parquet
        m = re.match(r"data/chunk-(\d+)/episode_(\d+)\.parquet", sample)
        if m:
            ep_per_chunk = 1000
            chunk_nums = set()
            ep_nums = set()
            for f in data_files_sample[:50]:
                m2 = re.match(r"data/chunk-(\d+)/episode_(\d+)\.parquet", f)
                if m2:
                    chunk_nums.add(int(m2.group(1)))
                    ep_nums.add(int(m2.group(2)))
            if len(chunk_nums) > 1:
                sorted_eps = sorted(ep_nums)
                for i in range(1, len(sorted_eps)):
                    if sorted_eps[i] - sorted_eps[i - 1] > 1:
                        ep_per_chunk = sorted_eps[i]
                        break
            elif len(ep_nums) > 1:
                max_ep = max(ep_nums)
                if max_ep >= 1000:
                    ep_per_chunk = 1000

            filenames = []
            for ep_idx in episode_indices:
                chunk = ep_idx // ep_per_chunk
                filenames.append(f"data/chunk-{chunk:03d}/episode_{ep_idx:06d}.parquet")
            return filenames

        # v1: data/episode_NNNNNN.parquet
        m = re.match(r"data/episode_(\d+)\.parquet", sample)
        if m:
            return [f"data/episode_{ep_idx:06d}.parquet" for ep_idx in episode_indices]

        return []

    def _try_direct_episode_download(
        self, repo_id: str, episode_indices: list[int]
    ) -> list[str] | None:
        ep_per_chunk = 1000
        test_ep = episode_indices[0]
        test_chunk = test_ep // ep_per_chunk
        test_filename = f"data/chunk-{test_chunk:03d}/episode_{test_ep:06d}.parquet"

        try:
            hf_hub_download(
                repo_id=repo_id, filename=test_filename, repo_type="dataset"
            )
        except (EntryNotFoundError, Exception):
            test_filename = f"data/episode_{test_ep:06d}.parquet"
            try:
                hf_hub_download(
                    repo_id=repo_id, filename=test_filename, repo_type="dataset"
                )
                return [f"data/episode_{ep:06d}.parquet" for ep in episode_indices]
            except (EntryNotFoundError, Exception):
                return None

        return [
            f"data/chunk-{ep // ep_per_chunk:03d}/episode_{ep:06d}.parquet"
            for ep in episode_indices
        ]

    @staticmethod
    def _extract_episode_index_from_filename(filename: str) -> int | None:
        match = re.search(r"episode_(\d+)\.parquet$", filename)
        if match:
            return int(match.group(1))
        return None

    def _load_chunk_format(
        self,
        repo_id: str,
        data_files: list[str],
        info: DatasetInfo,
        episode_indices: list[int] | None,
        max_episodes: int,
    ) -> pd.DataFrame:
        is_per_episode = any("episode_" in f for f in data_files[:5])

        if is_per_episode:
            target_indices = episode_indices
            if target_indices is None and info.num_episodes > max_episodes:
                target_indices = sample_episode_indices(info.num_episodes, max_episodes)

            if target_indices is not None:
                target_set = set(target_indices)
                data_files = [
                    f
                    for f in data_files
                    if self._extract_episode_index_from_filename(f) in target_set
                ]

        if not is_per_episode and len(data_files) > 3:
            data_files = data_files[:3]

        frames: list[pd.DataFrame] = []
        for filename in data_files:
            try:
                parquet_path = hf_hub_download(
                    repo_id=repo_id, filename=filename, repo_type="dataset"
                )
            except EntryNotFoundError:
                continue
            except (HfHubHTTPError, Exception) as e:
                if isinstance(e, (ValueError, EntryNotFoundError)):
                    raise
                raise ConnectionError(
                    f"Network error downloading {filename} from '{repo_id}': {e}"
                ) from e

            df = pd.read_parquet(parquet_path)
            frames.append(df)

        if not frames:
            return pd.DataFrame()

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

        if "episode_index" not in result.columns:
            return result

        if not is_per_episode:
            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 result

    def _load_per_episode_format(
        self,
        repo_id: str,
        data_files: list[str],
        info: DatasetInfo,
        episode_indices: list[int] | None,
        max_episodes: int,
    ) -> pd.DataFrame:
        if episode_indices is None:
            episode_indices = sample_episode_indices(info.num_episodes, max_episodes)

        frames: list[pd.DataFrame] = []
        for ep_idx in episode_indices:
            filename = f"data/episode_{ep_idx:06d}.parquet"
            if filename not in data_files:
                continue
            try:
                parquet_path = hf_hub_download(
                    repo_id=repo_id, filename=filename, repo_type="dataset"
                )
            except EntryNotFoundError:
                continue
            except (HfHubHTTPError, Exception) as e:
                if isinstance(e, (ValueError, EntryNotFoundError)):
                    raise
                raise ConnectionError(
                    f"Network error downloading episode {ep_idx} from '{repo_id}': {e}"
                ) from e

            df = pd.read_parquet(parquet_path)
            if "episode_index" not in df.columns:
                df["episode_index"] = ep_idx
            frames.append(df)

        if not frames:
            return pd.DataFrame()

        return pd.concat(frames, ignore_index=True)
