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

@tool
def get_a_joke(subject: str) -> str:
    """This tool searches for jokes about a specific subject and returns them.
    Args:
        subject: The subject to find a joke about.
    """
    
    search_tool = DuckDuckGoSearchTool()
    joke = None
    
    try:
        # Search for jokes about the subject
        results = search_tool(f"funny {subject} jokes one liner")
        
        # If we have results, try to find a good joke
        if results and len(results) > 0:
            for result in results[:3]:
                # Look for joke-like patterns (question marks, etc.)
                if '?' in result and len(result) < 200:
                    joke = result
                    break
            
            # If no ideal joke found, use the first result
            if not joke and len(results[0]) < 150:
                joke = results[0]
        
        # If we still don't have a joke, use a default one
        if not joke:
            default_jokes = {
                "chicken": "Why did the chicken cross the road? To get to the other side!",
                "dog": "What do you call a dog magician? A labracadabrador!",
                "cat": "What do you call a cat that gets anything it wants? Purrsuasive!",
                "default": "Why don't scientists trust atoms? Because they make up everything!"
            }
            
            if subject.lower() in default_jokes:
                joke = default_jokes[subject.lower()]
            else:
                joke = default_jokes["default"]
        
        return f"😄 Here's a joke about {subject}: {joke}"
        
    except Exception as e:
        return f"Sorry, I couldn't find a joke about {subject}. How about this one instead: Why don't scientists trust atoms? Because they make up everything!"

final_answer = FinalAnswerTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud'
model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud'

model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id=model_id,
    custom_role_conversions=None,
)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, DuckDuckGoSearchTool(), get_a_joke],
    max_steps=10,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="JokeBot",
    description="A bot that finds jokes on any subject",
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()