import logging
from typing import Optional

from .. import OnDemandCluster

from .table import Table

logger = logging.getLogger(__name__)


class HuggingFaceTable(Table):
    DEFAULT_FOLDER_PATH = "/runhouse/huggingface-tables"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @staticmethod
    def from_config(config: dict, dryrun=True):
        """Load config values into the object."""
        if isinstance(config["system"], dict):
            config["system"] = OnDemandCluster.from_config(
                config["system"], dryrun=dryrun
            )
        return HuggingFaceTable(**config, dryrun=dryrun)

    def write(self):
        hf_dataset = None
        if self._cached_data is not None:
            import datasets

            if isinstance(self.data, datasets.Dataset):
                arrow_table = self.data.data.table
                ray_dataset = self.ray_dataset_from_arrow(arrow_table)
                self.data, hf_dataset = ray_dataset, self.data
                self.write_ray_dataset(self.data)
            elif isinstance(self.data, datasets.DatasetDict):
                # TODO [JL] Add support for dataset dict
                raise NotImplementedError(
                    "Runhouse does not currently support DatasetDict objects, please convert to "
                    "a Dataset before saving."
                )
            else:
                raise TypeError(
                    "Unsupported data type for HuggingFaceTable. Please use a Dataset"
                )

            logger.info(f"Saved {str(self)} to: {self.fsspec_url}")

        # Restore the original dataset
        if hf_dataset is not None:
            self.data = self.to_dataset(hf_dataset)

        return self

    def fetch(self, **kwargs):
        # Read as pyarrow table, then convert back to HF dataset
        arrow_table = super().fetch(**kwargs)
        self._cached_data = self.to_dataset(arrow_table)
        return self._cached_data

    def stream(
        self,
        batch_size: int,
        drop_last: bool = False,
        shuffle_seed: Optional[int] = None,
        shuffle_buffer_size: Optional[int] = None,
        prefetch_blocks: Optional[int] = None,
        as_dict: bool = True,
    ):
        """Stream data as either Dataset object or dict (as generated by ray iter)"""
        for batch in super().stream(
            batch_size, drop_last, shuffle_seed, shuffle_buffer_size, prefetch_blocks
        ):
            yield batch if as_dict else self.to_dataset(batch)

    @staticmethod
    def to_dataset(data):
        """Convert to a HuggingFace Dataset. Relevant when fetching the data or when choosing to stream the data
        in as a Dataset."""
        import pandas as pd
        import pyarrow as pa
        import ray.data
        from datasets import Dataset

        if isinstance(data, Dataset):
            return data

        elif isinstance(data, dict):
            return Dataset.from_dict(data)

        elif isinstance(data, pd.DataFrame):
            return Dataset.from_pandas(data)

        elif isinstance(data, (pa.Table, ray.data.Dataset)):
            return Dataset.from_pandas(data.to_pandas())

        else:
            raise TypeError(
                f"Data must be a dict, Pandas DataFrame, ray Dataset, or PyArrow table, not {type(data)}"
            )
