# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

from ._session import (
    get_boto3_client,
    get_user_and_identity_store_id,
)
from .. import api


def _call_paginated_deadline_list_api(list_api, list_property_name, **kwargs):
    """
    Calls a deadline:List* API repeatedly to concatenate all pages.

    Example:
        deadline = get_boto3_client("deadline")
        return _call_paginated_deadline_list_api(deadline.list_farms, "farms", **kwargs)

    Args:
      list_api (callable): The List* API function to call, from the boto3 client.
      list_property_name (str): The name of the property in the response that contains
                                the list.
    """
    response = list_api(**kwargs)
    result = {list_property_name: response[list_property_name]}

    while "nextToken" in response:
        response = list_api(nextToken=response["nextToken"], **kwargs)
        result[list_property_name].extend(response[list_property_name])

    return result


@api.record_function_latency_telemetry_event()
def list_farms(config=None, **kwargs):
    """
    Calls the [deadline:ListFarms] API call, applying the filter for user membership
    depending on the configuration. If the response is paginated, it repeated
    calls the API to get all the farms.

    [deadline:ListFarms]: https://docs.aws.amazon.com/deadline-cloud/latest/APIReference/API_ListFarms.html
    """
    if "principalId" not in kwargs:
        user_id, _ = get_user_and_identity_store_id(config=config)
        if user_id:
            kwargs["principalId"] = user_id

    deadline = get_boto3_client("deadline", config=config)
    return _call_paginated_deadline_list_api(deadline.list_farms, "farms", **kwargs)


@api.record_function_latency_telemetry_event()
def list_queues(config=None, **kwargs):
    """
    Calls the [deadline:ListQueues] API call, applying the filter for user membership
    depending on the configuration. If the response is paginated, it repeated
    calls the API to get all the queues.

    [deadline:ListQueues]: https://docs.aws.amazon.com/deadline-cloud/latest/APIReference/API_ListQueues.html
    """
    if "principalId" not in kwargs:
        user_id, _ = get_user_and_identity_store_id(config=config)
        if user_id:
            kwargs["principalId"] = user_id

    deadline = get_boto3_client("deadline", config=config)
    return _call_paginated_deadline_list_api(deadline.list_queues, "queues", **kwargs)


@api.record_function_latency_telemetry_event()
def list_jobs(config=None, **kwargs):
    """
    Calls the [deadline:ListJobs] API call, applying the filter for user membership
    depending on the configuration. If the response is paginated, it repeated
    calls the API to get all the jobs.

    [deadline:ListJobs]: https://docs.aws.amazon.com/deadline-cloud/latest/APIReference/API_ListJobs.html
    """
    if "principalId" not in kwargs:
        user_id, _ = get_user_and_identity_store_id(config=config)
        if user_id:
            kwargs["principalId"] = user_id

    deadline = get_boto3_client("deadline", config=config)
    return _call_paginated_deadline_list_api(deadline.list_jobs, "jobs", **kwargs)


@api.record_function_latency_telemetry_event()
def list_fleets(config=None, **kwargs):
    """
    Calls the [deadline:ListFleets] API call, applying the filter for user membership
    depending on the configuration. If the response is paginated, it repeated
    calls the API to get all the fleets.

    [deadline:ListFleets]: https://docs.aws.amazon.com/deadline-cloud/latest/APIReference/API_ListFleets.html
    """
    if "principalId" not in kwargs:
        user_id, _ = get_user_and_identity_store_id(config=config)
        if user_id:
            kwargs["principalId"] = user_id

    deadline = get_boto3_client("deadline", config=config)
    return _call_paginated_deadline_list_api(deadline.list_fleets, "fleets", **kwargs)


@api.record_function_latency_telemetry_event()
def list_storage_profiles_for_queue(config=None, **kwargs):
    """
    Calls the [deadline:ListStorageProfilesForQueue] API call. If the response is paginated, it repeated
    calls the API to get all the storage profiles.

    [deadline:ListStorageProfilesForQueue]: https://docs.aws.amazon.com/deadline-cloud/latest/APIReference/API_ListStorageProfilesForQueue.html
    """
    deadline = get_boto3_client("deadline", config=config)

    return _call_paginated_deadline_list_api(
        deadline.list_storage_profiles_for_queue, "storageProfiles", **kwargs
    )
