"""
Dataset hub for managing and loading physical reasoning datasets.

This module provides a clean, simple interface for loading physical reasoning datasets.
"""

from pathlib import Path
from typing import Any, Dict, List, Optional, Type, Union

from prkit.prkit_core import PRKitLogger
from prkit.prkit_core.domain import PhysicalDataset
from prkit.prkit_datasets.downloaders import (
    PHYBenchDownloader,
    PhyXDownloader,
    PhysReasonDownloader,
    SeePhysDownloader,
    UGPhysicsDownloader,
)
from prkit.prkit_datasets.downloaders.base_downloader import BaseDownloader
from prkit.prkit_datasets.loaders import (
    JEEBenchLoader,
    PHYBenchLoader,
    PhyXLoader,
    PhysReasonLoader,
    SeePhysLoader,
    TPBenchLoader,
    UGPhysicsLoader,
)
from prkit.prkit_datasets.loaders.base_loader import BaseDatasetLoader


class DatasetHub:
    """
    Simple hub for loading physical reasoning datasets.

    This class provides a clean, intuitive interface similar to Hugging Face's datasets library.
    It combines registry functionality with a user-friendly API.

    Usage:
        # Simple loading
        dataset = DatasetHub.load("ugphysics")

        # With options
        dataset = DatasetHub.load("ugphysics", split="test", sample_size=100)

        # With variant
        dataset = DatasetHub.load("ugphysics", variant="mini", split="test")

        # List available datasets
        print(DatasetHub.list_available())

        # Get dataset info
        info = DatasetHub.get_info("ugphysics")

        # Register custom loader
        DatasetHub.register("custom", CustomLoader)
    """

    # Class-level registry of dataset loaders
    _loaders: Dict[str, Type[BaseDatasetLoader]] = {}
    # Class-level registry of dataset downloaders
    _downloaders: Dict[str, Type[BaseDownloader]] = {}
    _logger = PRKitLogger.get_logger(__name__)

    @classmethod
    def _register_default_loaders(cls):
        """Register the default dataset loaders."""
        cls.register("phybench", PHYBenchLoader)
        cls.register("phyx", PhyXLoader)
        cls.register("seephys", SeePhysLoader)
        cls.register("ugphysics", UGPhysicsLoader)
        cls.register("jeebench", JEEBenchLoader)
        cls.register("tpbench", TPBenchLoader)
        cls.register("physreason", PhysReasonLoader)

    @classmethod
    def _register_default_downloaders(cls):
        """Register the default dataset downloaders."""
        cls.register_downloader("phybench", PHYBenchDownloader)
        cls.register_downloader("phyx", PhyXDownloader)
        cls.register_downloader("physreason", PhysReasonDownloader)
        cls.register_downloader("seephys", SeePhysDownloader)
        cls.register_downloader("ugphysics", UGPhysicsDownloader)
        # Add more downloaders as they are implemented

    @classmethod
    def register(cls, name: str, loader_class: Type[BaseDatasetLoader]):
        """Register a new dataset loader."""
        cls._loaders[name] = loader_class

    @classmethod
    def register_downloader(cls, name: str, downloader_class: Type[BaseDownloader]):
        """Register a new dataset downloader."""
        cls._downloaders[name] = downloader_class

    @classmethod
    def _get_downloader(cls, name: str) -> Optional[BaseDownloader]:
        """Get a dataset downloader by name."""
        if not cls._downloaders:
            cls._register_default_downloaders()

        if name not in cls._downloaders:
            return None

        return cls._downloaders[name]()

    @classmethod
    def _get_loader(cls, name: str) -> BaseDatasetLoader:
        """Get a dataset loader by name."""
        if not cls._loaders:
            cls._register_default_loaders()

        if name not in cls._loaders:
            available = ", ".join(cls._loaders.keys())
            raise ValueError(
                f"Unknown dataset: {name}. Available datasets: {available}"
            )

        return cls._loaders[name]()

    @classmethod
    def load(
        cls,
        dataset_name: str,
        data_dir: Union[str, Path, None] = None,
        sample_size: Optional[int] = None,
        auto_download: bool = False,
        **kwargs,
    ) -> PhysicalDataset:
        """
        Load a physical reasoning dataset.

        Args:
            dataset_name: Name of the dataset ('ugphysics', 'phybench', 'seephys', etc.)
            data_dir: Path to the data directory (None = auto-detect)
            sample_size: Number of problems to load (None = all)
            auto_download: If True, automatically download the dataset if it doesn't exist
            **kwargs: Additional arguments for the specific loader (e.g., split, variant, etc.)

        Returns:
            PhysicalDataset: Loaded dataset

        Raises:
            ValueError: If dataset name is unknown, or if variant/split is invalid
            FileNotFoundError: If data directory doesn't exist and auto_download=False
            RuntimeError: If auto_download=True but download fails

        Examples:
            >>> # Load UGPhysics dataset (uses default variant and split)
            >>> dataset = DatasetHub.load("ugphysics")
            >>> print(f"Loaded {len(dataset)} problems")

            >>> # Load with sample size
            >>> dataset = DatasetHub.load("ugphysics", sample_size=50)

            >>> # Load specific split
            >>> dataset = DatasetHub.load("ugphysics", split="test")

            >>> # Load with variant and auto-download
            >>> dataset = DatasetHub.load("physreason", variant="full", auto_download=True)
        """
        # Get the appropriate loader
        loader = cls._get_loader(dataset_name)

        # Handle variant: use default if not provided, validate if provided
        variant = kwargs.pop("variant", None)
        if variant is None:
            variant = loader.get_default_variant()
            if variant is not None:
                cls._logger.info(
                    f"Using default variant '{variant}' for dataset '{dataset_name}'"
                )
            else:
                cls._logger.error(f"No default variant set for dataset '{dataset_name}'")
                raise ValueError(f"No default variant set and no variant provided for dataset '{dataset_name}'")
        else:
            # Validate explicitly provided variant
            try:
                loader.validate_variant(variant)
            except ValueError as e:
                cls._logger.error(
                    f"Invalid variant '{variant}' for dataset '{dataset_name}'. {e}"
                )
                raise

        # Handle split: use default if not provided, validate if provided
        split = kwargs.pop("split", None)
        if split is None:
            split = loader.get_default_split()
            if split is not None:
                cls._logger.info(
                    f"Using default split '{split}' for dataset '{dataset_name}'"
                )
            else:
                cls._logger.error(f"No default split set for dataset '{dataset_name}'")
                raise ValueError(f"No default split set and no split provided for dataset '{dataset_name}'")
        else:
            # Validate explicitly provided split
            try:
                loader.validate_split(split)
            except ValueError as e:
                cls._logger.error(
                    f"Invalid split '{split}' for dataset '{dataset_name}'. {e}"
                )
                raise

        # Resolve actual values for logging (including defaults)
        if data_dir is None:
            # Resolve default data directory
            actual_data_dir = loader.resolve_data_dir(None, dataset_name)
        else:
            actual_data_dir = data_dir
        
        if variant is None:
            actual_variant = loader.get_default_variant()
        else:
            actual_variant = variant
        
        if split is None:
            actual_split = loader.get_default_split()
        else:
            actual_split = split
        
        if sample_size is None:
            actual_sample_size = "full"
        else:
            actual_sample_size = sample_size

        # Log detailed loading arguments with actual resolved values
        cls._logger.info(f"Loading dataset '{dataset_name}' with the following arguments:")
        cls._logger.info(f"  - dataset_name: {dataset_name}")
        cls._logger.info(f"  - data_dir: {actual_data_dir}")
        cls._logger.info(f"  - variant: {actual_variant}")
        cls._logger.info(f"  - split: {actual_split}")
        cls._logger.info(f"  - sample_size: {actual_sample_size}")
        cls._logger.info(f"  - auto_download: {auto_download}")
        if kwargs:
            cls._logger.info(f"  - additional kwargs: {kwargs}")

        load_kwargs = {
            "data_dir": actual_data_dir,
            "variant": actual_variant,
            "split": actual_split,
            "sample_size": sample_size,
            "auto_download": auto_download,
            **kwargs,
        }

        # Try to load the dataset
        try:
            dataset = loader.load(**load_kwargs)
            return dataset
        except FileNotFoundError as e:
            # If dataset doesn't exist and auto_download is enabled, try to download
            if auto_download:
                cls._logger.info(
                    "Dataset not found. Attempting to download %s...", dataset_name
                )
                downloader = cls._get_downloader(dataset_name)

                if downloader is None:
                    cls._logger.warning(
                        "No downloader available for %s. Cannot auto-download.",
                        dataset_name,
                    )
                    raise FileNotFoundError(
                        f"Dataset not found and no downloader available for {dataset_name}. "
                        f"Please download the dataset manually or implement a downloader."
                    ) from e

                # Extract variant from load_kwargs (which may have been set to default)
                variant = load_kwargs.get("variant")
                if variant is None:
                    # Try to get default from loader
                    variant = loader.get_default_variant()
                    if variant is None:
                        variant = "full"  # Fallback default
                
                # Extract split if downloader needs it
                split = load_kwargs.get("split")
                if split is None:
                    split = loader.get_default_split()
                
                try:
                    # Download the dataset
                    download_kwargs = {"data_dir": data_dir, "force": False}
                    if variant is not None:
                        download_kwargs["variant"] = variant
                    if split is not None:
                        download_kwargs["split"] = split
                    
                    download_path = downloader.download(**download_kwargs)
                    cls._logger.info(
                        "Successfully downloaded %s to %s", dataset_name, download_path
                    )

                    # Retry loading after download
                    dataset = loader.load(**load_kwargs)
                    return dataset
                except Exception as download_error:
                    cls._logger.error(
                        "Failed to download %s: %s", dataset_name, download_error
                    )
                    raise RuntimeError(
                        f"Auto-download failed for {dataset_name}: {download_error}"
                    ) from download_error
            else:
                # Re-raise the original FileNotFoundError
                raise

    @classmethod
    def list_available(cls) -> List[str]:
        """List all available dataset names."""
        if not cls._loaders:
            cls._register_default_loaders()
        return list(cls._loaders.keys())

    @classmethod
    def get_info(cls, dataset_name: str) -> Dict[str, Any]:
        """Get information about a specific dataset."""
        loader = cls._get_loader(dataset_name)
        return loader.get_info()

    @classmethod
    def get_loader_info(cls, dataset_name: str) -> Dict[str, Any]:
        """Get detailed information about a dataset loader including supported parameters."""
        loader = cls._get_loader(dataset_name)
        info = loader.get_info()

        # Add loader class information
        info["loader_class"] = loader.__class__.__name__
        info["loader_module"] = loader.__class__.__module__

        return info
