"""
Base dataset class for SLTK.

All dataset loaders inherit from BaseDataset and implement
the abstract methods for loading samples.
"""

from abc import ABC, abstractmethod
from collections.abc import Iterator
from pathlib import Path
from typing import Any

from sltk.data.core import Sample
from sltk.exceptions import (
    ConfigurationError,
    DatasetNotFoundError,
    SampleNotFoundError,
)


class BaseDataset(ABC):
    """
    Abstract base class for all SLTK datasets.

    Subclasses must implement:
        - _load_index(): Build list of sample IDs
        - _load_sample(id): Load a single sample

    Attributes:
        NAME: Dataset identifier
        TASKS: Supported tasks (e.g., ['translation', 'recognition'])
        LANGUAGES: Sign languages in dataset (e.g., ['bsl', 'asl'])

    Example:
        >>> class MyDataset(BaseDataset):
        ...     NAME = "mydataset"
        ...     def _load_index(self): ...
        ...     def _load_sample(self, sample_id): ...
        >>>
        >>> dataset = MyDataset("/path/to/data", split="train")
        >>> sample = dataset[0]

    Using unified annotations:
        >>> dataset = MyDataset(
        ...     "/path/to/data",
        ...     unified_annotations="/path/to/annotations/mydataset",
        ... )
    """

    NAME: str = "base"
    TASKS: list[str] = []
    LANGUAGES: list[str] = []
    CITATION: str = ""

    # Default paths (can be overridden)
    DEFAULT_VIDEO_ROOT: Path | None = None
    DEFAULT_FEATURE_ROOT: Path | None = None
    DEFAULT_ANNOTATION_ROOT: Path | None = None

    def __init__(
        self,
        root: str | Path,
        split: str = "train",
        feature_root: str | Path | None = None,
        load_poses: bool = False,
        pose_format: str | None = None,
        fps: float = 25.0,
        unified_annotations: str | Path | None = None,
        video_root: str | Path | None = None,
    ):
        """
        Initialize dataset.

        Args:
            root: Path to dataset root directory
            split: Data split ('train', 'val', 'test')
            feature_root: Optional path to pre-extracted features
            load_poses: Whether to load pose data with each sample
            pose_format: Pose format to load ('wilor', 'smplx', etc.)
            fps: Frames per second for pose data
            unified_annotations: Path to unified annotation directory.
                If provided, samples are loaded from unified JSON files.
            video_root: Override video root path for resolving video paths.
        """
        self.root = Path(root)
        self.split = split
        self.feature_root = Path(feature_root) if feature_root else self.DEFAULT_FEATURE_ROOT
        self.load_poses = load_poses
        self.pose_format = pose_format
        self.fps = fps

        # Unified annotation support
        self.unified_annotations = (
            Path(unified_annotations) if unified_annotations else self.DEFAULT_ANNOTATION_ROOT
        )
        self.video_root = Path(video_root) if video_root else self.root
        self._use_unified = (
            self.unified_annotations is not None and self.unified_annotations.exists()
        )
        self._unified_manifest: dict[str, Any] | None = None

        if not self.root.exists():
            raise DatasetNotFoundError(self.NAME, str(self.root))

        # Load unified manifest if available
        if self._use_unified:
            self._load_unified_manifest()

        self._index = self._load_index()
        # Set for O(1) membership checks (list kept for ordered iteration)
        self._index_set: set[str] = set(self._index)

        # Feature path cache: maps sample_id -> Path (or None if not found)
        # Avoids repeated filesystem existence checks
        self._feature_path_cache: dict[str, Path | None] = {}

    @abstractmethod
    def _load_index(self) -> list[str]:
        """
        Build list of sample IDs for current split.

        Returns:
            List of unique sample identifiers
        """
        ...

    @abstractmethod
    def _load_sample(self, sample_id: str) -> Sample:
        """
        Load a single sample by ID.

        Args:
            sample_id: Unique sample identifier

        Returns:
            Sample instance with available data
        """
        ...

    def __len__(self) -> int:
        return len(self._index)

    def __getitem__(self, idx: int) -> Sample:
        if idx < 0 or idx >= len(self._index):
            raise IndexError(f"Index {idx} out of range for dataset of size {len(self)}")
        sample_id = self._index[idx]
        return self._load_sample(sample_id)

    def __iter__(self) -> Iterator[Sample]:
        for idx in range(len(self)):
            yield self[idx]

    def get_by_id(self, sample_id: str) -> Sample:
        """Load sample by ID directly."""
        if sample_id not in self._index_set:  # O(1) lookup using set
            raise SampleNotFoundError(sample_id, self.NAME)
        return self._load_sample(sample_id)

    @property
    def sample_ids(self) -> list[str]:
        """List of all sample IDs."""
        return self._index.copy()

    @property
    def num_samples(self) -> int:
        """Number of samples in dataset."""
        return len(self._index)

    def get_split_ids(self, split: str) -> list[str]:
        """
        Get sample IDs for a specific split.

        Override this method if split filtering is needed.
        """
        if split == self.split:
            return self._index.copy()
        # Would need to reload with different split
        raise NotImplementedError("Changing split requires re-initialization")

    def summary(self) -> dict[str, Any]:
        """Get dataset summary statistics."""
        summary = {
            "name": self.NAME,
            "split": self.split,
            "num_samples": len(self),
            "root": str(self.root),
            "feature_root": str(self.feature_root) if self.feature_root else None,
            "tasks": self.TASKS,
            "languages": self.LANGUAGES,
            "using_unified_annotations": self._use_unified,
        }

        if self._use_unified and self._unified_manifest:
            stats = self._unified_manifest.get("statistics", {})
            summary["unified_stats"] = {
                "total_samples": stats.get("total_samples"),
                "unique_glosses": stats.get("unique_glosses"),
                "unique_signers": stats.get("unique_signers"),
            }

        return summary

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__name__}("
            f"name={self.NAME!r}, "
            f"split={self.split!r}, "
            f"num_samples={len(self)})"
        )

    # =========================================================================
    # Unified Annotation Support
    # =========================================================================

    def _load_unified_manifest(self) -> None:
        """Load the unified annotation manifest if available."""
        if self.unified_annotations is None:
            return

        manifest_path = self.unified_annotations / "manifest.json"
        if manifest_path.exists():
            import json

            with open(manifest_path, encoding="utf-8") as f:
                self._unified_manifest = json.load(f)

    def _load_unified_index(self) -> list[str]:
        """
        Load sample IDs from unified annotations.

        This is used when unified_annotations is set. Override _load_index()
        to call this method when unified annotations should be used.

        Returns:
            List of sample IDs for the current split
        """
        if not self._use_unified or self.unified_annotations is None:
            return []

        # Use manifest if available
        if self._unified_manifest:
            samples = self._unified_manifest.get("samples", [])
            if self.split:
                return [s["sample_id"] for s in samples if s.get("split") == self.split]
            return [s["sample_id"] for s in samples]

        # Fall back to scanning files
        from sltk.io.annotations import list_samples

        return list_samples(self.unified_annotations, split=self.split)

    def _load_from_unified(self, sample_id: str) -> Sample:
        """
        Load a sample from unified annotation format.

        Args:
            sample_id: Sample identifier

        Returns:
            Sample loaded from unified annotation file

        Raises:
            FileNotFoundError: If annotation file not found
        """
        if self.unified_annotations is None:
            raise ConfigurationError("unified_annotations not set")

        from sltk.io.annotations import read_unified_annotation

        annotation_path = self.unified_annotations / f"{sample_id}.json"

        if not annotation_path.exists():
            raise SampleNotFoundError(sample_id, f"{self.NAME} unified annotations")

        sample = read_unified_annotation(
            annotation_path,
            load_video_path=True,
            video_root=self.video_root,
        )

        # Load poses if requested
        if self.load_poses and self.pose_format and self.feature_root:
            poses = self._load_poses_for_sample(sample_id)
            if poses is not None:
                sample.poses = poses

        return sample

    def _load_poses_for_sample(self, sample_id: str):
        """
        Load pose data for a sample.

        Args:
            sample_id: Sample identifier

        Returns:
            PoseSequence or None if not found
        """
        if not self.feature_root or not self.pose_format:
            return None

        from sltk.data.core import PoseSequence

        # Check cache first to avoid repeated filesystem existence checks
        cache_key = f"{self.pose_format}:{sample_id}"
        if cache_key in self._feature_path_cache:
            pose_path = self._feature_path_cache[cache_key]
            if pose_path is None:
                return None
            try:
                return PoseSequence.load(pose_path, format=self.pose_format, fps=self.fps)
            except Exception as e:
                import logging

                logger = logging.getLogger(__name__)
                logger.debug(f"Failed to load poses from {pose_path}: {e}")
                return None

        # Try common pose file patterns
        pose_patterns = [
            self.feature_root / self.NAME / self.pose_format / f"{sample_id}.h5",
            self.feature_root / self.NAME / self.pose_format / f"{sample_id}.npy",
            self.feature_root / self.pose_format / f"{sample_id}.h5",
        ]

        for pose_path in pose_patterns:
            if pose_path.exists():
                # Cache the found path
                self._feature_path_cache[cache_key] = pose_path
                try:
                    return PoseSequence.load(pose_path, format=self.pose_format, fps=self.fps)
                except Exception as e:
                    import logging

                    logger = logging.getLogger(__name__)
                    logger.debug(f"Failed to load poses from {pose_path}: {e}")

        # Cache negative result to avoid repeated checks
        self._feature_path_cache[cache_key] = None
        return None

    @property
    def has_unified_annotations(self) -> bool:
        """Check if unified annotations are available."""
        return self._use_unified

    def get_unified_statistics(self) -> dict[str, Any] | None:
        """Get statistics from unified manifest."""
        if self._unified_manifest:
            return self._unified_manifest.get("statistics")
        return None
