import logging
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

try:
    from ngcsdk import Client  # type: ignore
except ImportError:  # pragma: no cover - SDK might not be available locally
    LOGGER.warning("NVIDIA NGC SDK not found. Please install it with 'pip install ngcsdk'.")
    Client = None


LOGGER = logging.getLogger(__name__)


def _load_team_choices(image_list, pattern):
    """
    Extracts unique team names from a list of NGC images.

    Args:
        image_list: A list of image batches from the NGC SDK.
        pattern: The pattern used to query for images, for logging.

    Returns:
        A list of dictionaries for dropdown choices, with 'label' and 'value'.
    """
    if not image_list:
        LOGGER.warning("No images found for pattern: %s", pattern)
        return []
         
    team_choices = {
        image.teamName
        for image_batch in image_list
        for image in image_batch
    }
    
    sorted_team_choices = sorted(list(team_choices))

    return_choices = [{
        "label": team.upper(),
        "value": team
    } for team in sorted_team_choices]

    LOGGER.info(f"Found {len(return_choices)} teams.")
    return return_choices

def _load_model_choices(image_list, pattern):
    """
    Extracts model choices for a given team from a list of NGC images.

    Args:
        image_list: A list of image batches from the NGC SDK.
        pattern: The pattern used to generate the full image paths.

    Returns:
        A list of dictionaries for dropdown choices, with 'label' and 'value'.
    """
    if not image_list:
        LOGGER.warning("No images found for pattern: %s", pattern)
        return []
         
    return_choices = []

    # Process the image data
    for image_batch in image_list:
        for image in image_batch:
            model_w_tag = f"{image.name}:{image.latestTag}"
            full_path = pattern.replace('*',model_w_tag)
            return_choices.append({
                "label": image.displayName,
                "value": full_path
            })

    LOGGER.info(f"Found {len(return_choices)} models for pattern '{pattern}'.")
    return return_choices

def do(payload, config, plugin_config, inputs):
    """
    Provide dependent drop-down options for the NIM team, model and tag selectors.

    This function is called by Dataiku to populate the choices for parameters
    that use 'getChoicesFromPython'. It fetches available NIM models from the
    NVIDIA NGC registry.
    """
    parameter_name = payload.get("parameterName")
    LOGGER.info(f"Getting choices for parameter: {parameter_name}")

    if not Client:
        return {"choices": [{"label": "NGC SDK not installed", "value": ""}]}

    api_key = config.get("nim_services_authentication", {}).get("nim_container_registry_api_key")
    client = Client()
    client.configure(api_key=api_key)
    image_api = client.registry.image
    choices = []

    if parameter_name == "nim_model_team":
        pattern = "nvcr.io/nim/*/*"
        LOGGER.info(f"Listing images with pattern: {pattern}")
        nim_images = image_api.list(pattern=pattern)
        return {"choices": _load_team_choices(nim_images, pattern)}


    elif parameter_name == "nim_model_name":
        team_name = config.get("nim_model_team", "")
        if team_name:
            pattern = f"nvcr.io/nim/{team_name}/*"
            LOGGER.info(f"Listing images with pattern: {pattern}")
            nim_images = image_api.list(pattern=pattern)
            return {"choices": _load_model_choices(nim_images, pattern)}
        else:
            LOGGER.warning("Team name is not selected, cannot load model choices.")
    else:
        LOGGER.warning(f"Unhandled parameter name: {parameter_name}")
        return {"choices": {}}
