code_kit/agent_style_generator.py

import json
import gradio as gr
from model_handler import ModelHandler
from config import LING_FLASH_2_0

def format_palette_html(palette_data):
    """
    Generates an HTML string to display color boxes with roles.
    palette_data: List of dicts {'role': str, 'hex': str} or list of strings (fallback)
    """
    html = '<div style="display: flex; gap: 10px; margin-top: 5px; flex-wrap: wrap;">'
    
    for item in palette_data:
        if isinstance(item, dict):
            color = item.get("hex", "#000000")
            role = item.get("role", "")
        else:
            color = item
            role = ""

        html += f'''
        <div style="display: flex; flex-direction: column; align-items: center; width: 60px;">
            <div style="width: 40px; height: 40px; background-color: {color}; border-radius: 8px; border: 1px solid #ccc; box-shadow: 0 2px 4px rgba(0,0,0,0.1);" title="{role}: {color}"></div>
            <span style="font-size: 9px; color: #666; margin-top: 4px; text-align: center; line-height: 1.1; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; width: 100%;">{role}</span>
            <span style="font-size: 9px; color: #999; text-align: center;">{color}</span>
        </div>
        '''
    html += '</div>'
    return html

def generate_random_style(current_model_display_name):
    """
    Generates a random web design style using the LLM.
    """
    model_choice = LING_FLASH_2_0 
    
    system_prompt = """
    You are a creative Design Director. Your task is to generate a unique, cohesive visual style for a web interface.
    
    Return ONLY a valid JSON object with the following keys:
    1. "palette": A list of 6 to 8 color objects, where each object has:
        - "role": A descriptive name for the color usage (e.g., "Page Background", "Card/Surface", "Primary Text", "Secondary Text", "Brand/Action Color", "Accent/Highlight", "Border/Divider").
        - "hex": The hexadecimal color code.
    2. "decoration": A concise description (10-20 words) of the UI decoration style (e.g., "Neo-brutalism with thick black borders and hard shadows", "Glassmorphism with high blur and white transparency").
    3. "theme": A creative, metaphorical description (5-10 words) of the overall vibe (e.g., "Cyberpunk Neon City", "Minimalist Zen Garden").
    
    Ensure the colors contrast well (especially text on background) and the roles make sense for a standard web layout.
    Do not include any markdown formatting (like ```json). Just the raw JSON string.
    """
    
    user_prompt = "Generate a new random design style now."
    
    model_handler = ModelHandler()
    
    full_response = ""
    for chunk in model_handler.generate_code(system_prompt, user_prompt, model_choice):
        full_response += chunk
    
    clean_response = full_response.strip()
    if clean_response.startswith("```json"):
        clean_response = clean_response[7:]
    if clean_response.startswith("```"):
        clean_response = clean_response[3:]
    if clean_response.endswith("```"):
        clean_response = clean_response[:-3]
    clean_response = clean_response.strip()
    
    try:
        style_data = json.loads(clean_response)
        palette = style_data.get("palette", [])
        
        # Validate palette structure, fallback if it's just a list of strings (old format)
        if palette and isinstance(palette[0], str):
            palette = [{"role": f"Color {i+1}", "hex": c} for i, c in enumerate(palette)]
            
        if not palette:
             palette = [
                {"role": "Background", "hex": "#F0F0F0"},
                {"role": "Text", "hex": "#333333"},
                {"role": "Accent", "hex": "#007BFF"}
            ]

        decoration = style_data.get("decoration", "Standard modern web style.")
        theme = style_data.get("theme", "Default Theme")
        
        palette_html = format_palette_html(palette)
        
        # Create a structured string for the prompt
        # e.g., "Page Background: #FFF, Primary Text: #000, ..."
        palette_str = ", ".join([f"{p['role']}: {p['hex']}" for p in palette])
        
        return palette_html, palette_str, decoration, theme
        
    except json.JSONDecodeError as e:
        print(f"Error parsing style JSON: {e}. Response: {full_response}")
        fallback_palette = [
            {"role": "Background", "hex": "#FFFFFF"},
            {"role": "Primary Text", "hex": "#000000"},
            {"role": "Action", "hex": "#007BFF"},
            {"role": "Error", "hex": "#DC3545"}
        ]
        return format_palette_html(fallback_palette), ", ".join([f"{p['role']}: {p['hex']}" for p in fallback_palette]), "Simple and clean.", "Fallback Default"