import os
from models import Settings, Character, Look
from utils import _LORA_DEFAULTS


def get_available_loras(category):
    """Return sorted list of LoRA paths for the given category.
    category: one of 'characters','outfits','actions','styles','scenes','detailers'
    """
    settings = Settings.query.first()
    lora_dir = (getattr(settings, f'lora_dir_{category}', None) if settings else None) or _LORA_DEFAULTS.get(category, '')
    if not lora_dir or not os.path.isdir(lora_dir):
        return []
    subfolder = os.path.basename(lora_dir.rstrip('/'))
    return sorted(f"Illustrious/{subfolder}/{f}" for f in os.listdir(lora_dir) if f.endswith('.safetensors'))


def get_available_checkpoints():
    settings = Settings.query.first()
    checkpoint_dirs_str = (settings.checkpoint_dirs if settings else None) or \
        '/ImageModels/Stable-diffusion/Illustrious,/ImageModels/Stable-diffusion/Noob'
    checkpoints = []
    for ckpt_dir in checkpoint_dirs_str.split(','):
        ckpt_dir = ckpt_dir.strip()
        if not ckpt_dir or not os.path.isdir(ckpt_dir):
            continue
        prefix = os.path.basename(ckpt_dir.rstrip('/'))
        for f in os.listdir(ckpt_dir):
            if f.endswith('.safetensors') or f.endswith('.ckpt'):
                checkpoints.append(f"{prefix}/{f}")
    return sorted(checkpoints)


def _count_look_assignments():
    """Return a dict mapping look_id to the count of characters it's assigned to."""
    assignment_counts = {}
    looks = Look.query.all()
    for look in looks:
        if look.character_id:
            assignment_counts[look.look_id] = 1
        else:
            assignment_counts[look.look_id] = 0
    return assignment_counts


def _count_outfit_lora_assignments():
    """Return a dict mapping outfit LoRA filename to the count of characters using it."""
    assignment_counts = {}
    characters = Character.query.all()

    for character in characters:
        char_lora = character.data.get('lora', {}).get('lora_name', '')
        if char_lora and 'Clothing' in char_lora:
            assignment_counts[char_lora] = assignment_counts.get(char_lora, 0) + 1

        wardrobe = character.data.get('wardrobe', {})
        if 'default' in wardrobe and isinstance(wardrobe.get('default'), dict):
            for outfit_name, outfit_data in wardrobe.items():
                if isinstance(outfit_data, dict):
                    outfit_lora = outfit_data.get('lora', {})
                    if isinstance(outfit_lora, dict):
                        lora_name = outfit_lora.get('lora_name', '')
                        if lora_name:
                            assignment_counts[lora_name] = assignment_counts.get(lora_name, 0) + 1

    return assignment_counts