import typing as t

from datasets import load_dataset
import pyarrow as pa

from sarus_data_spec.arrow.schema import type_from_arrow_schema
from sarus_data_spec.manager.ops.base import (
    DatasetImplementation,
    DatasetStaticChecker,
)
from sarus_data_spec.schema import schema
import sarus_data_spec.typing as st


class HuggingfaceStaticChecker(DatasetStaticChecker):
    def __init__(self, dataset: st.Dataset):
        super().__init__(dataset)
        spec = dataset.protobuf().spec
        assert spec.WhichOneof("spec") == "huggingface"
        self.ds_name = spec.huggingface.name
        self.split = spec.huggingface.split

    async def schema(self) -> st.Schema:
        ds = load_dataset(path=self.ds_name, split=self.split)
        data_type = type_from_arrow_schema(ds.data.schema)
        return schema(dataset=self.dataset, schema_type=data_type)


class Huggingface(DatasetImplementation):
    def __init__(self, dataset: st.Dataset):
        super().__init__(dataset)
        spec = dataset.protobuf().spec
        assert spec.WhichOneof("spec") == "huggingface"
        self.ds_name = spec.huggingface.name
        self.split = spec.huggingface.split

    async def to_arrow(
        self, batch_size: int
    ) -> t.AsyncIterator[pa.RecordBatch]:
        ds = load_dataset(path=self.ds_name, split=self.split)

        async def async_iterator() -> t.AsyncIterator[pa.RecordBatch]:
            for batch in ds.data.to_batches(max_chunksize=batch_size):
                yield batch

        return async_iterator()