import folder_paths
from typing import Optional
from .logger import get_logger
from .model_cache import cache
from .prompt_utils import clean_keywords

logger = get_logger('utils.lora_stack')


def norm_lora_stack(lora_stack: Optional[list]) -> Optional[list]:
    """
    Normalize the lora_stack input to always be a list of lists/tuples.
    
    Args:
        lora_stack (list, optional): The input lora_stack which can be None, a single lora tuple/list,
                                     or a list of lora tuples/lists.
                                     
    Returns:
        list or None: The normalized lora_stack as a list of tuples/lists, or None if input was None.
    """
    if lora_stack is None:
        return None

    if isinstance(lora_stack, (tuple, list)):
        if len(lora_stack) == 0:
            return None
        if not isinstance(lora_stack[0], (list, tuple)):
            lora_stack = [lora_stack]
    else:
        lora_stack = None

    return lora_stack
    
def get_lora_keywords(lora_name):
    # Local import minimizes module initialization coupling.
    from .model_metadata import pull_metadata

    lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
    if cache.by_path(lora_path).get("trainedWords") is None:
        pull_metadata(lora_path, timestamp=True)
    return cache.by_path(lora_path).get("trainedWords", [])

def get_lora_stack_keywords(lora_stack=None):
    # Local import minimizes module initialization coupling.
    from .model_metadata import pull_metadata

    lora_stack = norm_lora_stack(lora_stack)
    if lora_stack is None:
        return []

    # Collect unique lora names
    lora_names = {lora[0] for lora in lora_stack}
    lora_paths = [folder_paths.get_full_path_or_raise("loras", name) for name in lora_names]
    pull_metadata(lora_paths)
    # Gather all keywords into a set for uniqueness
    all_keywords = set()
    for name in lora_names:
        try:
            keywords = cache.by_path(folder_paths.get_full_path_or_raise("loras", name)).get("trainedWords", [])
            if keywords:
                all_keywords.update(map(str.strip, keywords))
        except Exception as e:
            logger.debug(f"Exception getting keywords for {name}: {e}")
            continue
    return clean_keywords(all_keywords)

def add_lora_to_stack(lora_name, model_weight, clip_weight, lora_stack=None):
    lora = (lora_name, model_weight, clip_weight)
    if lora_stack is None:
        return [lora]
    return [*lora_stack, lora]