app.py

import gradio as gr
import json
import re
from typing import Dict, List, Tuple
import os

# Important: Make sure you have python-dotenv installed (`pip install python-dotenv`)
# and a .env file if you run locally. On HF Spaces, use secrets.
from dotenv import load_dotenv

# Try to import CrewAI libraries. If they fail, it means they are not installed.
try:
    from crewai import Agent, Task, Crew, Process
    from langchain_openai import ChatOpenAI
except ImportError:
    raise ImportError(
        "CrewAI or LangChain not installed. Please ensure your requirements.txt is correct and dependencies are installed."
    )

load_dotenv()

# --- Configuration ---
# You can change the model here. Claude 3 Haiku is fast, cheap, and very capable.
DEFAULT_MODEL = "minimax/minimax-m2:free"

# --- Agent Definitions ---

# 1. Color Specialist Agent
color_specialist = Agent(
    role='Expert Color Palette Designer',
    goal="""Select a visually appealing and contextually appropriate color palette 
    based on the user's website description. The palette must include primary, 
    secondary, background, surface, and text colors.""",
    backstory="""You are a world-renowned graphic designer with a deep understanding of 
    color theory and its psychological impact. You create palettes that are not 
    only beautiful but also functional and accessible.""",
    verbose=True,
    allow_delegation=False
)

# 2. Typography Specialist Agent
typography_specialist = Agent(
    role='Master Typographer for Digital Interfaces',
    goal="""Define a complete and harmonious typography system, including a web-safe 
    font family and a responsive size scale (h1, h2, body, small).""",
    backstory="""With decades of experience, you excel at choosing fonts that are readable, 
    stylish, and appropriate for a brand's voice. You create clear typographic 
    hierarchies that work perfectly on the web.""",
    verbose=True,
    allow_delegation=False
)

# 3. Design System Architect Agent (The "Stitcher")
design_system_architect = Agent(
    role='Senior UI Design System Architect',
    goal="""Consolidate the color and typography schemes into a single, comprehensive,
    and developer-friendly JSON design system. Also add standard 'spacing' and 
    'borderRadius' systems.""",
    backstory="""You are a meticulous senior UI designer who creates robust and scalable 
    design systems. Your work ensures absolute consistency and serves as the single 
    source of truth for the entire development team.""",
    verbose=True,
    allow_delegation=False
)

# 4. Web UI Developer Agent
web_developer = Agent(
    role='Expert Frontend Developer specializing in Web UI',
    goal="""Generate a complete, single HTML file for a given page description, 
    strictly adhering to the provided design system. The code must be clean, 
    responsive, and modern.""",
    backstory="""You are a pixel-perfect frontend developer who translates design systems 
    into clean, responsive, and maintainable HTML and CSS. You never deviate from 
    the design specifications.""",
    verbose=True,
    allow_delegation=False
)

# --- Task Definitions ---

def create_design_tasks(prompt):
    task_colors = Task(
        description=f"""Analyze the user's prompt: '{prompt}'. 
        Generate a JSON object for the color palette. 
        The JSON must contain keys: 'primary', 'secondary', 'background', 'surface', 'text'.""",
        expected_output="A single, valid JSON object containing the color palette.",
        agent=color_specialist
    )

    task_typography = Task(
        description=f"""Analyze the user's prompt: '{prompt}'. 
        Generate a JSON object for the typography system. 
        It must include 'fontFamily', and sizes for 'h1', 'h2', 'body', 'small'.""",
        expected_output="A single, valid JSON object for the typography system.",
        agent=typography_specialist
    )

    task_architect = Task(
        description="""Take the color palette and typography system from the specialists. 
        Combine them into a single, final JSON object. Also, add standard 'spacing' 
        (sm, md, lg) and 'borderRadius' (md, lg) systems. 
        Your final output must be ONLY this complete JSON object and nothing else.""",
        expected_output="A single, valid JSON object representing the complete design system.",
        agent=design_system_architect,
        context=[task_colors, task_typography]
    )
    return task_colors, task_typography, task_architect

def create_developer_task(design_system_json, page_description):
    return Task(
        description=f"""Using the final design system provided below, create a complete HTML file 
        for the following page: '{page_description}'. All CSS must be in a single `<style>` tag 
        in the `<head>`.

        FINAL DESIGN SYSTEM:
        {design_system_json}""",
        expected_output="The complete, raw HTML code for the webpage, starting with `<!DOCTYPE html>` and nothing else.",
        agent=web_developer
    )

# --- Main Gradio Function ---

def create_ui_design(prompt: str, api_key: str, progress=gr.Progress(track_tqdm=True)):
    if not api_key:
        error_msg = "API Key is missing. Please enter your OpenRouter API key."
        return (error_msg,) * 5
    
    os.environ["OPENROUTER_API_KEY"] = api_key
    
    # Configure the LLM to use OpenRouter
    openrouter_llm = ChatOpenAI(
        model=DEFAULT_MODEL,
        api_key=api_key,
        base_url="https://openrouter.ai/api/v1"
    )

    try:
        progress(0.1, desc="Briefing the design team...")
        design_tasks = create_design_tasks(prompt)
        design_crew = Crew(
            agents=[color_specialist, typography_specialist, design_system_architect],
            tasks=list(design_tasks),
            process=Process.sequential,
            verbose=2,
            llm=openrouter_llm
        )
        
        progress(0.3, desc="Architecting the Design System...")
        design_system_json_string = design_crew.kickoff(inputs={'prompt': prompt})
        
        design_system = json.loads(design_system_json_string)
        
        progress(0.6, desc="Briefing the Web Developer...")
        page_descriptions = re.findall(r'\d+\)([^0-9]+?)(?=\d+\)|$)', prompt) or [prompt]
        while len(page_descriptions) < 3:
            page_descriptions.append(f"Page {len(page_descriptions)+1} placeholder")

        generated_pages = []
        developer_crew = Crew(agents=[web_developer], tasks=[], process=Process.sequential, llm=openrouter_llm)

        for i, page_desc in enumerate(page_descriptions[:3]):
            progress(0.6 + (i * 0.1), desc=f"Developer is building Page {i+1}...")
            dev_task = create_developer_task(design_system_json_string, page_desc)
            developer_crew.tasks = [dev_task]
            html_output = developer_crew.kickoff()
            generated_pages.append(html_output)
            
        progress(0.9, desc="Creating Visualizer...")
        visualizer_html = visualize_design_system(design_system)

        progress(1.0, desc="Done!")
        return generated_pages[0], generated_pages[1], generated_pages[2], design_system_json_string, visualizer_html

    except Exception as e:
        error_html = f"An error occurred with the CrewAI agents: {str(e)}. Check the logs for details."
        print(error_html)
        return (error_html,) * 5

# --- Visualizer Function ---
def visualize_design_system(design_system: Dict) -> str:
    if not design_system or not isinstance(design_system, dict): return "<div class='error-box'>No valid design system to display.</div>"
    colors = design_system.get('colors', {})
    typography = design_system.get('typography', {})
    html = """<!DOCTYPE html><html lang="en"><head><meta charset="UTF-8"><title>Design System</title><style>body{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',Roboto,sans-serif;background-color:#f9fafb;color:#374151;padding:1.5rem;}.section{margin-bottom:2rem;}h2{font-size:1.5rem;font-weight:600;border-bottom:1px solid #e5e7eb;padding-bottom:.5rem;margin-bottom:1rem;}.swatch-container{display:flex;flex-wrap:wrap;gap:1rem;}.swatch{width:100px;height:100px;border-radius:.5rem;display:flex;flex-direction:column;justify-content:flex-end;padding:.5rem;font-size:.8rem;box-shadow:0 2px 4px rgba(0,0,0,0.1);}.swatch .name{font-weight:600;}.type-container .sample{padding:1rem;background-color:#fff;border:1px solid #e5e7eb;border-radius:.5rem;margin-bottom:1rem;}</style></head><body>"""
    html += "<div class='section'><h2>Color Palette</h2><div class='swatch-container'>"
    for name, hex_val in colors.items():
        try: r,g,b=int(hex_val[1:3],16),int(hex_val[3:5],16),int(hex_val[5:7],16); text_color='#fff' if (0.299*r+0.587*g+0.114*b)/255<0.5 else '#000'
        except: text_color='#000'
        html+=f'<div class="swatch" style="background-color:{hex_val};color:{text_color};"><span class="name">{name.capitalize()}</span><span>{hex_val}</span></div>'
    html += "</div></div>"
    font_family = typography.get('fontFamily', 'sans-serif')
    html += f"<div class='section'><h2>Typography</h2><p>Font: <code>{font_family}</code></p>"
    for name, sample_text in {'h1':'Heading 1','h2':'Heading 2','body':'This is a paragraph of body text.'}.items():
        html+=f'<div class="sample"><div style="font-family:{font_family};font-size:{typography.get(name,"1rem")};">{sample_text}</div></div>'
    html += "</div></body></html>"
    return html

# --- Gradio Interface ---
def create_interface():
    with gr.Blocks(title="CrewAI UI Designer", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🤖 CrewAI Powered UI Designer")
        with gr.Row():
            with gr.Column(scale=1):
                api_key_input = gr.Textbox(label="OpenRouter API Key", placeholder="sk-or-...", type="password")
                prompt_input = gr.Textbox(label="Website Description", placeholder="e.g., A dark-mode, futuristic site for a synthwave musician: 1) landing page, 2) music, 3) tour dates", lines=4)
                generate_btn = gr.Button("🚀 Assemble the Crew!", variant="primary")
                with gr.Tabs():
                    with gr.TabItem("Visualizer"): design_system_visualizer = gr.HTML()
                    with gr.TabItem("JSON"): design_system_output = gr.Code(language="json", lines=15)
            with gr.Column(scale=2):
                with gr.Tabs():
                    with gr.TabItem("📄 Page 1"): page1_output = gr.HTML()
                    with gr.TabItem("📄 Page 2"): page2_output = gr.HTML()
                    with gr.TabItem("📄 Page 3"): page3_output = gr.HTML()
        generate_btn.click(fn=create_ui_design, inputs=[prompt_input, api_key_input], outputs=[page1_output, page2_output, page3_output, design_system_output, design_system_visualizer])
    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch()