"""Base class for all dataset format loaders."""

from __future__ import annotations

from abc import ABC, abstractmethod

import pandas as pd

from orbit.analyzer.dataset_loader import DatasetInfo


class BaseLoader(ABC):
    """Abstract base class for dataset format loaders.

    All loaders produce the same output contract:
    - DatasetInfo metadata dataclass
    - pd.DataFrame with columns: [episode_index, frame_index, action, state, timestamp]
    """

    @abstractmethod
    def can_load(self, source: str) -> bool:
        """Return True if this loader can handle the given source."""
        ...

    @abstractmethod
    def load_info(self, source: str) -> DatasetInfo:
        """Load dataset metadata without loading full episode data."""
        ...

    @abstractmethod
    def load_episodes(
        self,
        source: str,
        episode_indices: list[int] | None = None,
        max_episodes: int | None = None,
    ) -> pd.DataFrame:
        """Load episode data as DataFrame.

        Returns DataFrame with columns:
            episode_index (int), frame_index (int), action (list/array),
            state (list/array or None), timestamp (float)
        """
        ...


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


def normalize_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    """Ensure DataFrame has the expected columns in the correct order."""
    if df.empty:
        return pd.DataFrame(
            columns=["episode_index", "frame_index", "action", "state", "timestamp"]
        )

    # Map observation.state -> state if needed
    if "state" not in df.columns and "observation.state" in df.columns:
        df = df.rename(columns={"observation.state": "state"})

    expected = ["episode_index", "frame_index", "action", "state", "timestamp"]
    for col in expected:
        if col not in df.columns:
            df[col] = None

    return df[expected + [c for c in df.columns if c not in expected]]
