"""
Utils for AIS PyTorch Plugin

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

from typing import List, Mapping, Tuple
from urllib.parse import urlparse, urlunparse
from aistore.sdk import Client


def parse_url(url: str) -> Tuple[str, str, str]:
    """
    Parse AIS urls for bucket and object names
    Args:
        url (str): Complete URL of the object (eg. "ais://bucket1/file.txt")
    Returns:
        provider (str): AIS Backend
        bck_name (str): Bucket name identifier
        obj_name (str):  Object name with extension
    """
    parsed_url = urlparse(url)
    path = parsed_url.path
    if len(path) > 0 and path.startswith("/"):
        path = path[1:]

    # returns provider, bck_name, path
    return parsed_url.scheme, parsed_url.netloc, path


# pylint: disable=unused-variable
def list_objects_info(client: Client, urls_list: List[str]) -> List[Mapping[str, str]]:
    """
    Create list of [bucket_name, object_name] from all the object urls
    Args:
        client (Client): AIStore client object of the calling method
        urls_list (List[str]): list of urls
    Returns:
        List[samples](List[Mapping[str, str]]): list of {provider, bucket, path to the object}
    """
    samples = []
    for url in urls_list:
        provider, bck_name, path = parse_url(url)
        objects = client.bucket(bck_name=bck_name, provider=provider).list_objects(
            prefix=path
        )
        for obj_info in objects.get_entries():
            samples.append(
                {"provider": provider, "bck_name": bck_name, "object": obj_info.name}
            )
    return samples


def unparse_url(provider: str, bck_name: str, obj_name: str) -> str:
    """
    To generate URL based on provider, bck_name and object name
    Args:
        provider(str): Provider name ('ais', 'gcp', etc)
        bck_name(str): Bucket name
        obj_name(str): Object name with extension.
    Returns:
        unparsed_url(str): Unparsed url (complete url)
    """
    return urlunparse([provider, bck_name, obj_name, "", "", ""])
