from smolagents import CodeAgent, HfApiModel, tool
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

@tool
def generate_coloring_page(theme: str, complexity: str) -> str:
    """Generates a description for a coloring page based on theme and complexity.
    
    Args:
        theme: The theme for the coloring page (animals, nature, space, vehicles, fantasy)
        complexity: The complexity level (simple, medium, detailed)
    
    Returns:
        A description of the coloring page
    """
    # Validate inputs
    valid_themes = ["animals", "nature", "space", "vehicles", "fantasy"]
    valid_complexity = ["simple", "medium", "detailed"]
    
    theme = theme.lower()
    complexity = complexity.lower()
    
    if theme not in valid_themes:
        return f"Error: Theme must be one of {', '.join(valid_themes)}."
    
    if complexity not in valid_complexity:
        return f"Error: Complexity must be one of {', '.join(valid_complexity)}."
    
    # Database of coloring pages by theme and complexity
    coloring_pages = {
        "animals": {
            "simple": ["Simple outline of a happy dog", "Basic cat sitting", "Round fish with bubbles", "Smiling turtle outline"],
            "medium": ["Duck swimming in a pond", "Dog playing with a ball", "Fox in the forest", "Rabbit among carrots"],
            "detailed": ["Detailed elephant with patterns", "Lion with flowing mane", "Peacock with decorative tail", "Owl on a branch with feather details"]
        },
        "nature": {
            "simple": ["Basic flower with petals", "Simple tree outline", "Sun with rays", "Round cloud shapes"],
            "medium": ["Flower garden path", "Tree with leaves and fruits", "Mountain landscape", "Beach with waves"],
            "detailed": ["Forest scene with various trees", "Detailed garden with flowers and butterflies", "Mountain range with details", "Underwater scene with plants"]
        },
        "space": {
            "simple": ["Basic rocket outline", "Simple planet with rings", "Smiling star shape", "Moon with craters"],
            "medium": ["Rocket launching with flames", "Solar system with planets", "Astronaut floating", "Space station design"],
            "detailed": ["Detailed spaceship with components", "Galaxy with stars and nebula", "Space mission scene", "Alien landscape with plants"]
        },
        "vehicles": {
            "simple": ["Simple car outline", "Basic boat shape", "Train with few cars", "Round helicopter"],
            "medium": ["Car with windows and wheels", "Sailboat on waves", "Train with details", "Airplane flying"],
            "detailed": ["Racing car with details", "Ship with many components", "Detailed train with landscape", "Helicopter with mechanical parts"]
        },
        "fantasy": {
            "simple": ["Simple unicorn outline", "Basic dragon shape", "Castle with towers", "Smiling fairy outline"],
            "medium": ["Unicorn with flowing mane", "Dragon with wings", "Castle with details", "Fairy with wand and dress"],
            "detailed": ["Detailed unicorn in magical forest", "Dragon with scales and fire", "Castle kingdom with background", "Fairy scene with flowers and magic"]
        }
    }
    
    # Select a few coloring page descriptions
    selected_pages = coloring_pages[theme][complexity]
    
    # Build response
    response = f"## {theme.capitalize()} Coloring Pages ({complexity.capitalize()} Level)\n\n"
    
    for page in selected_pages:
        response += f"- {page}\n"
    
    response += "\nThese coloring pages feature simple black outlines on white backgrounds, perfect for printing and coloring with crayons, colored pencils, or markers."
    
    return response

# Setup agent
final_answer = FinalAnswerTool()

# Use FLUX.1-schnell as specified in the second code
model = HfApiModel(
    max_tokens=1024,
    temperature=0.5,
    model_id='FLUX.1-schnell'  # Using the model specified in the second code
)

# Define ALL required prompt templates including the missing "task" key
prompt_templates = {
    "react-system": """You are a helpful coloring page generator that suggests coloring pages for children.
You have access to the following tools:

{tools_string}

You should help the user find appropriate coloring pages based on:
1. The theme they're interested in (animals, nature, space, vehicles, fantasy)
2. The complexity level they want (simple, medium, detailed)

Use the generate_coloring_page tool to provide suggestions.

To use a tool, output an action in this format:
```py
action = tool_name(arg1=value1, arg2=value2, ...)
```

Always use Python syntax for calling tools. After using a tool, analyze the result to determine what to do next. Use final_answer to return your final answer.
""",
    "react-human": "{input_string}",
    "react-assistant": "{agent_response}",
    "final_answer": {
        "pre_messages": "Here are some coloring page suggestions for you:",
        "template": "Final Answer: {answer}",
        "post_messages": "Would you like different themes or complexity levels?"
    },
    "system_prompt": "You are a helpful coloring page generator that suggests coloring pages for children.",
    "planning": {
        "pre_messages": "I'll plan how to help you find coloring pages.",
        "template": "Plan: {planning}",
        "post_messages": "Now I'll execute this plan to help you.",
        "initial_plan": "I'll identify the theme and complexity level the user wants, then use the generate_coloring_page tool to find appropriate suggestions.",
        "update_plan_pre_messages": "I need to update my plan based on what I've learned.",
        "update_plan_template": "Updated Plan: {planning}",
        "update_plan_post_messages": "Now I'll execute this updated plan."
    },
    "managed_agent": {
        "pre_messages": "Let me think about the best coloring pages for you.",
        "template": "{agent_response}",
        "post_messages": "I hope these suggestions are helpful!",
        "task": "Generate appropriate coloring page suggestions based on the user's preferred theme and complexity level.",
        "report": "Here are the coloring page suggestions based on the selected theme and complexity level."
    }
}

# Create the agent with all required templates
agent = CodeAgent(
    model=model,
    tools=[final_answer, generate_coloring_page],
    max_steps=3,
    verbosity_level=1,
    prompt_templates=prompt_templates
)

# Create a simple Gradio UI
class ColoringPageUI(GradioUI):
    def launch(self, **kwargs):
        import gradio as gr
        
        with gr.Blocks(theme=gr.themes.Soft()) as demo:
            gr.Markdown("# Simple Coloring Page Generator")
            gr.Markdown("Get coloring page ideas for children")
            
            with gr.Row():
                with gr.Column():
                    theme_input = gr.Radio(
                        ["Animals", "Nature", "Space", "Vehicles", "Fantasy"], 
                        label="Theme"
                    )
                    complexity_input = gr.Radio(
                        ["Simple", "Medium", "Detailed"], 
                        label="Complexity Level"
                    )
                    submit_btn = gr.Button("Get Coloring Pages", variant="primary")
            
            with gr.Row():
                with gr.Column():
                    output = gr.Markdown(label="Coloring Page Suggestions")
            
            def get_suggestions(theme, complexity):
                if not theme:
                    return "Please select a theme."
                if not complexity:
                    return "Please select a complexity level."
                
                try:
                    # Format the query for the agent
                    query = f"I want {complexity.lower()} coloring pages with a {theme.lower()} theme"
                    # Process through the agent
                    result = agent.run(query)
                    return result
                except Exception as e:
                    return f"Error generating suggestions: {str(e)}"
            
            submit_btn.click(
                get_suggestions, 
                inputs=[theme_input, complexity_input], 
                outputs=output
            )
        
        demo.launch(**kwargs)

# Launch the interface
if __name__ == "__main__":
    ColoringPageUI(agent).launch()