"""RLDS / Open X-Embodiment dataset loader (TFRecord format)."""

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_tfrecord():
    """Check if tfrecord package is available."""
    try:
        import tfrecord  # noqa: F401

        return True
    except ImportError:
        return False


def _check_tensorflow():
    """Check if tensorflow_datasets is available."""
    try:
        import tensorflow_datasets  # noqa: F401

        return True
    except ImportError:
        return False


class RLDSLoader(BaseLoader):
    """Load RLDS / Open X-Embodiment datasets from TFRecord files."""

    def can_load(self, source: str) -> bool:
        """True if source is a directory containing .tfrecord files."""
        p = Path(source).expanduser()
        if not p.is_dir():
            return False
        # Check for tfrecord files
        tfrecords = list(p.glob("*.tfrecord*"))
        if tfrecords:
            return True
        # Check for TFDS dataset_info.json
        if (p / "dataset_info.json").exists():
            return True
        # Check in subdirectories (e.g., split dirs like train/)
        for subdir in p.iterdir():
            if subdir.is_dir() and list(subdir.glob("*.tfrecord*")):
                return True
        return False

    def load_info(self, source: str) -> DatasetInfo:
        """Load metadata from RLDS dataset."""
        root = Path(source).expanduser().resolve()

        # Try TFDS dataset_info.json
        info_path = root / "dataset_info.json"
        num_episodes = 0
        num_frames = 0
        fps = 10.0  # RLDS doesn't always specify FPS
        robot_type = "unknown"
        shapes: dict[str, list[int]] = {}

        if info_path.exists():
            with open(info_path) as f:
                ds_info = json.load(f)

            # Parse TFDS metadata
            splits = ds_info.get("splits", [])
            for split in splits:
                if isinstance(split, dict):
                    num_episodes += split.get("statistics", {}).get(
                        "numExamples", split.get("numShards", 0)
                    )

            # Try to get feature info
            features = ds_info.get("features", {})
            if "steps" in features:
                steps_feat = features["steps"]
                if isinstance(steps_feat, dict):
                    feature_info = steps_feat.get("feature", {})
                    if "action" in feature_info:
                        action_shape = feature_info["action"].get("shape", {}).get(
                            "dimensions", []
                        )
                        if action_shape:
                            shapes["action"] = [int(d) for d in action_shape]
                    obs = feature_info.get("observation", {})
                    if isinstance(obs, dict) and "state" in obs:
                        state_shape = obs["state"].get("shape", {}).get(
                            "dimensions", []
                        )
                        if state_shape:
                            shapes["state"] = [int(d) for d in state_shape]

        # If no metadata, count tfrecord files as rough episode estimate
        if num_episodes == 0:
            tfrecords = self._find_tfrecords(root)
            num_episodes = max(len(tfrecords), 1)

        return DatasetInfo(
            repo_id=str(root),
            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 TFRecord files."""
        if max_episodes is None:
            max_episodes = 100

        # Try tensorflow_datasets first, then fall back to tfrecord
        if _check_tensorflow():
            return self._load_with_tfds(source, episode_indices, max_episodes)
        elif _check_tfrecord():
            return self._load_with_tfrecord(source, episode_indices, max_episodes)
        else:
            raise ImportError(
                "RLDS format detected but neither 'tensorflow-datasets' nor 'tfrecord' "
                "is installed.\n  Install with: pip install orbit-robotics[rlds]"
            )

    def _find_tfrecords(self, root: Path) -> list[Path]:
        """Find all tfrecord files in the dataset directory."""
        tfrecords = sorted(root.glob("*.tfrecord*"))
        if not tfrecords:
            # Check subdirectories (split dirs)
            for subdir in sorted(root.iterdir()):
                if subdir.is_dir():
                    tfrecords.extend(sorted(subdir.glob("*.tfrecord*")))
        return tfrecords

    def _load_with_tfds(
        self,
        source: str,
        episode_indices: list[int] | None,
        max_episodes: int,
    ) -> pd.DataFrame:
        """Load using tensorflow_datasets."""
        import tensorflow_datasets as tfds

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

        try:
            builder = tfds.builder_from_directory(str(root))
            ds = builder.as_dataset(split="all")
        except Exception:
            # Try loading as a standard TFDS dataset
            try:
                ds = tfds.load(str(root), split="train")
            except Exception as e:
                raise ValueError(f"Could not load RLDS dataset from {source}: {e}")

        rows: list[dict] = []
        episode_idx = 0

        for episode in ds:
            if episode_indices is not None and episode_idx not in episode_indices:
                episode_idx += 1
                continue
            if episode_indices is None and episode_idx >= max_episodes:
                break

            steps = episode.get("steps", episode)
            frame_idx = 0
            for step in steps:
                action = step["action"].numpy() if "action" in step else None
                obs = step.get("observation", {})
                state = obs["state"].numpy() if "state" in obs else None

                rows.append(
                    {
                        "episode_index": episode_idx,
                        "frame_index": frame_idx,
                        "action": action.tolist() if action is not None else None,
                        "state": state.tolist() if state is not None else None,
                        "timestamp": frame_idx / 10.0,  # Default 10 FPS
                    }
                )
                frame_idx += 1

            episode_idx += 1

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

        return normalize_dataframe(pd.DataFrame(rows))

    def _load_with_tfrecord(
        self,
        source: str,
        episode_indices: list[int] | None,
        max_episodes: int,
    ) -> pd.DataFrame:
        """Load using the lightweight tfrecord package."""
        from tfrecord.reader import tfrecord_loader

        root = Path(source).expanduser().resolve()
        tfrecords = self._find_tfrecords(root)

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

        if episode_indices is not None:
            target = set(episode_indices)
        else:
            target = None

        rows: list[dict] = []
        episode_idx = 0

        for tf_path in tfrecords:
            if target is not None and episode_idx not in target:
                episode_idx += 1
                continue
            if target is None and episode_idx >= max_episodes:
                break

            try:
                for frame_idx, record in enumerate(
                    tfrecord_loader(str(tf_path), None, None)
                ):
                    action = record.get("action") or record.get("steps/action")
                    state = (
                        record.get("observation/state")
                        or record.get("steps/observation/state")
                    )

                    rows.append(
                        {
                            "episode_index": episode_idx,
                            "frame_index": frame_idx,
                            "action": (
                                list(action) if action is not None else None
                            ),
                            "state": list(state) if state is not None else None,
                            "timestamp": frame_idx / 10.0,
                        }
                    )
            except Exception as e:
                logger.warning("Error reading %s: %s", tf_path, e)

            episode_idx += 1

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

        return normalize_dataframe(pd.DataFrame(rows))
