import os
import gradio as gr
from smolagents import CodeAgent, DuckDuckGoSearchTool, FinalAnswerTool, InferenceClientModel, Tool, tool, VisitWebpageTool


# ── Tools ──────────────────────────────────────────────────────────────────────

@tool
def suggest_menu(occasion: str) -> str:
    """
    Suggests a menu based on the occasion.
    Args:
        occasion: The type of occasion for the party (casual, formal, superhero).
    """
    menus = {
        "casual": "Pizza, sliders, nachos, and a selection of craft beers & soft drinks.",
        "formal": "3-course dinner: shrimp cocktail starter, filet mignon entrée, crème brûlée dessert with wine pairing.",
        "superhero": "High-energy buffet: protein bowls, power smoothies, and hero-themed cupcakes.",
        "villain masquerade": "Dark elegance spread: black caviar, smoked meats, red-velvet cake, and 'Poison Ivy' mocktails.",
    }
    return menus.get(occasion.lower(), "Custom gourmet menu crafted by the Wayne Manor chef.")


@tool
def catering_service_tool(query: str) -> str:
    """
    Returns the highest-rated catering service in Gotham City.
    Args:
        query: A search term for finding catering services.
    """
    services = {
        "Gotham Catering Co.": 4.9,
        "Wayne Manor Catering": 4.8,
        "Gotham City Events": 4.7,
    }
    best_service = max(services, key=services.get)
    return f"{best_service} (rated {services[best_service]}/5)"


@tool
def party_checklist_tool(theme: str) -> str:
    """
    Generates a party planning checklist for a given theme.
    Args:
        theme: The party theme (e.g., villain masquerade, superhero gala).
    """
    base = [
        "✅ Book the venue (Wayne Manor Grand Hall)",
        "✅ Send invitations 2 weeks in advance",
        "✅ Arrange catering & bar service",
        "✅ Set up a photo booth with themed props",
        "✅ Hire a DJ or live band",
        "✅ Arrange valet parking",
    ]
    theme_extras = {
        "villain masquerade": [
            "🎭 Provide masquerade masks at the entrance",
            "🕯️ Use dramatic candlelit décor",
            "🌹 Add rose & smoke machine effects",
        ],
        "classic heroes": [
            "🦸 Set up a cape & costume station",
            "🏆 Run a 'Best Costume' contest",
            "🎨 Commission hero-themed artwork",
        ],
        "futuristic gotham": [
            "🔵 Install neon LED lighting throughout",
            "🤖 Feature tech/gadget displays",
            "🎮 Set up a VR experience zone",
        ],
    }
    extras = theme_extras.get(theme.lower(), ["🎉 Custom décor to match your unique theme"])
    return "\n".join(base + extras)


class SuperheroPartyThemeTool(Tool):
    name = "superhero_party_theme_generator"
    description = (
        "Suggests creative superhero-themed party ideas based on a category. "
        "Returns a unique party theme concept with décor and activity ideas."
    )
    inputs = {
        "category": {
            "type": "string",
            "description": "Type of superhero party: 'classic heroes', 'villain masquerade', or 'futuristic Gotham'.",
        }
    }
    output_type = "string"

    def forward(self, category: str) -> str:
        themes = {
            "classic heroes": (
                "⚡ Justice League Gala — Guests dress as DC heroes. "
                "Themed cocktails like 'The Kryptonite Punch'. Giant hero murals as photo backdrops."
            ),
            "villain masquerade": (
                "🎭 Gotham Rogues' Ball — Mysterious masquerade where guests dress as Batman villains. "
                "Dark floral décor, candlelight, and a 'Who's Behind the Mask?' contest."
            ),
            "futuristic gotham": (
                "🌆 Neo-Gotham Night — Cyberpunk party inspired by Batman Beyond. "
                "Neon decorations, holographic displays, and futuristic gadget showcases."
            ),
        }
        return themes.get(
            category.lower(),
            "Theme not found. Try 'classic heroes', 'villain masquerade', or 'futuristic Gotham'.",
        )


# ── Agent setup ────────────────────────────────────────────────────────────────

def build_agent() -> CodeAgent:
    return CodeAgent(
        tools=[
            DuckDuckGoSearchTool(),
            VisitWebpageTool(),
            suggest_menu,
            catering_service_tool,
            party_checklist_tool,
            SuperheroPartyThemeTool(),
            FinalAnswerTool(),
        ],
        model=InferenceClientModel(token=os.environ.get("HF_TOKEN")),
        max_steps=10,
        verbosity_level=1,
    )


# ── Gradio UI ──────────────────────────────────────────────────────────────────

SYSTEM_NOTE = (
    "🦇 **Alfred at your service.** I am the Wayne Manor party planning assistant. "
    "Ask me about themes, menus, playlists, catering, or full party plans!"
)

EXAMPLE_QUERIES = [
    "Plan a villain masquerade party at Wayne Manor",
    "Suggest a menu for a superhero gala",
    "Give me a party checklist for a futuristic Gotham theme",
    "Find the best catering service in Gotham City",
    "What playlist fits a classic heroes party?",
]


def respond(message: str, history: list) -> str:
    agent = build_agent()
    try:
        answer = agent.run(message)
        return str(answer)
    except Exception as e:
        return f"Alfred encountered an issue: {e}"


with gr.Blocks(title="Wayne Manor Party Planner") as demo:
    gr.HTML(
        """
        <div style="text-align:center; padding: 1rem 0 0.5rem;">
            <h1 style="font-size:2rem;">🦇 Wayne Manor Party Planner</h1>
            <p style="color:#888;">Powered by <strong>Alfred</strong> &mdash; your AI butler &amp; event coordinator</p>
        </div>
        """
    )

    gr.ChatInterface(
        fn=respond,
        chatbot=gr.Chatbot(height=480, placeholder=SYSTEM_NOTE),
        textbox=gr.Textbox(
            placeholder="Ask Alfred to plan your party...",
            container=False,
            scale=7,
        ),
        examples=EXAMPLE_QUERIES,
        cache_examples=False,
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Base(primary_hue="slate", neutral_hue="zinc"))