import os
import yaml
import gradio as gr

from smolagents import CodeAgent, HfApiModel

from tools.party_theme_tool import PartyThemeTool
from tools.search import search
from tools.categorize_occasion import categorize_occasion
from tools.estimate_preparation_time import estimate_preparation_time
from tools.web_search import DuckDuckGoSearchTool as WebSearch

from huggingface_hub import login

login(token=os.getenv('HF_TOKEN'))

model = HfApiModel(
    model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    provider=None,
)


with open("prompts.yaml", "r") as stream:
    prompt_templates = yaml.safe_load(stream)

web_search = WebSearch()

agent = CodeAgent(
    tools=[categorize_occasion, search, web_search, estimate_preparation_time, PartyThemeTool()],
    model=model,
    managed_agents=[],
    max_steps=10,
    verbosity_level=2,
    grammar=None,
    planning_interval=None,
    description=None,
    prompt_templates=prompt_templates,
)


def agent_interface(user_input):
    """
    This function defines the interaction with the agent.

    Args:
        user_input: The input provided by the user in the Gradio interface.

    Returns:
        The agent's response, formatted for readability.
    """
    try:
        response = agent.run(user_input)

        if isinstance(response, dict):
            formatted_response = ""
            for key, value in response.items():
                formatted_response += f"{key.upper()}\n\n"
                if isinstance(value, dict):
                    for nested_key, nested_value in value.items():
                        if isinstance(nested_value, list):
                            nested_value = "\n  - " + "\n  - ".join(nested_value)
                        formatted_response += (
                            f"- {nested_key.capitalize()}:\n  {nested_value}\n\n"
                        )

                elif isinstance(value, list):
                    formatted_response += "\n  - " + "\n  - ".join(value) + "\n\n"

                else:
                    formatted_response += f"{value}\n\n"

            response = formatted_response.strip()

    except Exception as e:
        response = f"Error: {e}"

    return response


interface = gr.Interface(
    fn=agent_interface,
    inputs=gr.Textbox(
        lines=1, placeholder="Enter your request for the agent", label=""
    ),
    outputs=gr.Textbox(label=""),
    title="Party Prep Agent",
    description="Ask the agent for help planning your party!",
    flagging_mode="never",
    show_progress="full",
    examples=[
        ["Suggest a menu and a theme for formal dinner party"],
        ["Suggest a menu, music list for highschool graduation party"],
        ["Sugest a fine dining menu and how long it till take to prpate"],
        ["Teen birthday party theme and menu ideas"],
    ],
)

if __name__ == "__main__":
    interface.launch(debug=True)