"""
AIS Plugin for PyTorch

PyTorch Dataset and DataLoader for AIS.

Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
"""

from typing import List, Union
from torch.utils.data import Dataset

from aistore.sdk import Client
from aistore.pytorch.utils import list_objects_info

# pylint: disable=too-few-public-methods


class AISBaseClass:
    """
    A base class for creating AIS Datasets for PyTorch

    Args:
        client_url(str): AIS endpoint URL
        urls_list(str or List[str]): single or list of url prefixes to load data
    """

    def __init__(self, client_url: str, urls_list: Union[str, List[str]]) -> None:
        self.client = Client(client_url)
        if isinstance(urls_list, str):
            urls_list = [urls_list]
        self._object_info = {}
        self._object_info = list_objects_info(self.client, urls_list)


# pylint: disable=unused-variable


class AISDataset(AISBaseClass, Dataset):
    """
    A map-style dataset for objects in AIS.
    If `etl_name` is provided, that ETL must already exist on the AIStore cluster

    Args:
        client_url (str): AIS endpoint URL
        urls_list (str or List[str]): single or list of url prefixes to load data
        etl_name (str, optional): Optional etl on the AIS cluster to apply to each object
    """

    def __init__(
        self, client_url: str, urls_list: Union[str, List[str]], etl_name=None
    ):
        AISBaseClass.__init__(self, client_url, urls_list)
        self.etl_name = etl_name

    def __len__(self):
        return len(self._object_info)

    def __getitem__(self, index: int):
        object_name = self._object_info[index]["object"]
        obj = (
            self.client.bucket(
                bck_name=self._object_info[index]["bck_name"],
                provider=self._object_info[index]["provider"],
            )
            .object(obj_name=object_name)
            .get(etl_name=self.etl_name)
            .read_all()
        )
        return object_name, obj
