modules/core_logic.py

import yaml
import random
import os
import json
import tempfile
import gradio as gr
from .config import FEATURES_FILE, FEATURE_SEQUENCE, SECTIONS, EXAMPLES_DIR

def load_features():
    with open(FEATURES_FILE, "r") as f:
        return yaml.safe_load(f)

features_data = load_features()

def get_detail(category, subcategory, key):
    """Retrieves the detailed description for a given key in a category/subcategory."""
    return features_data.get(category, {}).get(subcategory, {}).get(key, key)

def generate_prompt(*args):
    """
    Assembles the prompt based on dropdown selections and extra text info.
    """
    character_name = args[0]
    num_features = len(FEATURE_SEQUENCE)
    feature_keys = args[1:num_features+1]
    # Checkboxes are at args[num_features+1 : num_features+1 + num_features]
    # Extra info starts after the name, dropdowns, and checkboxes
    extra_infos = args[1 + num_features*2 : 1 + num_features*2 + len(SECTIONS)]

    template = features_data.get("templates", {}).get("default", "")
    
    context = {}
    for i, (cat, subcat, t_key) in enumerate(FEATURE_SEQUENCE):
        key = feature_keys[i]
        context[t_key] = get_detail(cat, subcat, key)
    
    # Handle multiple accessories dynamically
    acc_list = []
    # Identify accessory keys in context
    for i, (cat, subcat, t_key) in enumerate(FEATURE_SEQUENCE):
        if subcat == 'accessory' and feature_keys[i] != "None":
            acc_list.append(context[t_key])
    
    if acc_list:
        context['accessories'] = ", and has " + " as well as ".join(acc_list)
    else:
        context['accessories'] = ""
    
    # Inject extra info into the context
    if extra_infos[0]: # Identity
        context['age'] += f", {extra_infos[0]}"
    if extra_infos[1]: # Appearance
        context['distinguishing_feature'] += f", also {extra_infos[1]}"
    if extra_infos[2]: # Equipment
        if context['accessories']:
            context['accessories'] += f", further complemented by {extra_infos[2]}"
        else:
            context['accessories'] = f", complemented by {extra_infos[2]}"
    if extra_infos[3]: # Environment
        context['atmosphere'] += f", additionally {extra_infos[3]}"
    if extra_infos[4]: # Style
        context['camera'] += f", art style notes: {extra_infos[4]}"
    
    try:
        base_prompt = template.format(**context)
        if character_name and character_name.strip() and character_name != "Unnamed Hero":
            return f"A portrait of {character_name}, {base_prompt}"
        return base_prompt
    except Exception as e:
        return f"Error building prompt: {e}"

def handle_regeneration(*args):
    """Randomizes checkboxes and returns new values for dropdowns."""
    num_features = len(FEATURE_SEQUENCE)
    current_values = list(args[:num_features])
    checkboxes = args[num_features : num_features*2]
    
    new_values = []
    for i, (is_random, (cat, subcat, t_key)) in enumerate(zip(checkboxes, FEATURE_SEQUENCE)):
        if is_random:
            choices = list(features_data[cat][subcat].keys())
            new_values.append(random.choice(choices))
        else:
            new_values.append(current_values[i])
            
    return new_values

def save_character(*args):
    """Saves all current UI values to a JSON file."""
    character_name = args[0]
    num_features = len(FEATURE_SEQUENCE)
    feature_keys = args[1:num_features+1]
    checkboxes = args[num_features+1 : num_features+1 + num_features]
    extra_infos = args[num_features+1 + num_features : num_features+1 + num_features + len(SECTIONS)]
    
    data = {
        "name": character_name,
        "features": {FEATURE_SEQUENCE[i][2]: feature_keys[i] for i in range(num_features)},
        "randomization": {FEATURE_SEQUENCE[i][2]: checkboxes[i] for i in range(num_features)},
        "extra_info": {SECTIONS[i].lower(): extra_infos[i] for i in range(len(SECTIONS))}
    }
    
    # Save to a file with a friendly name in a temp directory
    safe_name = "".join([c if c.isalnum() else "_" for c in character_name]).strip("_")
    filename = f"{safe_name}_data.json" if safe_name else "rpg_character_data.json"
    
    temp_dir = tempfile.mkdtemp()
    path = os.path.join(temp_dir, filename)
    with open(path, 'w') as f:
        json.dump(data, f, indent=4)
    return path

def load_character(file):
    """Loads UI values from a JSON file."""
    if file is None:
        return [gr.update()] * (1 + len(FEATURE_SEQUENCE) * 2 + len(SECTIONS))
    
    try:
        file_path = file.name if hasattr(file, 'name') else file
        with open(file_path, 'r') as f:
            data = json.load(f)
        
        updates = []
        # Update name
        updates.append(data.get("name", "Unnamed Hero"))
        
        # Update dropdowns
        for _, _, t_key in FEATURE_SEQUENCE:
            updates.append(data.get("features", {}).get(t_key, gr.update()))
        
        # Update checkboxes (randomization flags)
        for _, _, t_key in FEATURE_SEQUENCE:
            updates.append(data.get("randomization", {}).get(t_key, gr.update()))
        
        # Update extra info textboxes
        for section in SECTIONS:
            updates.append(data.get("extra_info", {}).get(section.lower(), ""))
            
        return updates
    except Exception as e:
        print(f"Error loading character: {e}")
        return [gr.update()] * (1 + len(FEATURE_SEQUENCE) * 2 + len(SECTIONS))

def get_example_list():
    """Returns a list of JSON filenames in the examples directory."""
    if not os.path.exists(EXAMPLES_DIR):
        return []
    return [f for f in os.listdir(EXAMPLES_DIR) if f.endswith('.json')]

def load_example_character(filename):
    """Loads a character config from the examples directory."""
    if not filename:
        return [gr.update()] * (1 + len(FEATURE_SEQUENCE) * 2 + len(SECTIONS))
    
    path = os.path.join(EXAMPLES_DIR, filename)
    return load_character(path)