import yaml
import os
import gradio as gr

# ---------------------------------------------------
# 🔧 Patch 1: Blocks ignores deprecated kwargs
# ---------------------------------------------------
_old_blocks = gr.Blocks

def PatchedBlocks(*args, **kwargs):
    # Remove smolagents old kwargs
    for k in ["theme", "fill_height"]:
        kwargs.pop(k, None)
    return _old_blocks(*args, **kwargs)

gr.Blocks = PatchedBlocks


# ---------------------------------------------------
# 🔧 Patch 2: Chatbot ignores ALL deprecated kwargs
# ---------------------------------------------------
_old_chatbot = gr.Chatbot

def PatchedChatbot(*args, **kwargs):
    # Allowed kwargs in Gradio 6 Chatbot:
    allowed = {
        "value", "height", "scale", "min_height", "max_height",
        "rtl", "show_copy_button", "avatar_images",
        "layout","render","visible","interactive"
    }

    # Strip unsupported ones
    kwargs = {k: v for k, v in kwargs.items() if k in allowed}

    return _old_chatbot(*args, **kwargs)

gr.Chatbot = PatchedChatbot
# ---------------------------------------------------


from smolagents import GradioUI, CodeAgent, InferenceClientModel

# Get current directory path
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))

from tools.web_search import WebSearchTool as WebSearch
from tools.visit_webpage import VisitWebpageTool as VisitWebpage
from tools.suggest_menu import SimpleTool as SuggestMenu
from tools.catering_service_tool import SimpleTool as CateringServiceTool
from tools.superhero_party_theme_generator import SuperheroPartyThemeTool as SuperheroPartyThemeGenerator
from tools.final_answer import FinalAnswerTool as FinalAnswer


# Model
model = InferenceClientModel(
    model_id='Qwen/Qwen3-Next-80B-A3B-Thinking',
)

# Tools
web_search = WebSearch()
visit_webpage = VisitWebpage()
suggest_menu = SuggestMenu()
catering_service_tool = CateringServiceTool()
superhero_party_theme_generator = SuperheroPartyThemeGenerator()
final_answer = FinalAnswer()

# Load Prompts
with open(os.path.join(CURRENT_DIR, "prompts.yaml"), 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

# Agent
agent = CodeAgent(
    model=model,
    tools=[
        web_search,
        visit_webpage,
        suggest_menu,
        catering_service_tool,
        superhero_party_theme_generator,
    ],
    managed_agents=[],
    max_steps=10,
    verbosity_level=2,
    planning_interval=None,
    executor_type='local',
    prompt_templates=prompt_templates
)

# Launch UI
if __name__ == "__main__":
    GradioUI(agent).launch()