from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool

from Gradio_UI import GradioUI
import gradio as gr

       

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def get_birth_stone(month:str)-> str: #it's import to specify the return type
    # Keep this format for the description / args / args description but feel free to modify the tool
    """A tool that takes a month string as input. It outputs the birthstone of that month.
    Args:
        month: A string representing a month of the year
    """

    birthstones = {
        "January": "Garnet",
        "February": "Amethyst",
        "March": "Aquamarine",
        "April": "Diamond",
        "May": "Emerald",
        "June": "Pearl",
        "July": "Ruby",
        "August": "Peridot",
        "September": "Sapphire",
        "October": "Opal",
        "November": "Topaz",
        "December": "Turquoise"
    }
        
    month = month
    birthstone = birthstones.get(month, "Unknown")
  
    return birthstone

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """A tool that fetches the current local time in a specified timezone.
    Args:
        timezone: A string representing a valid timezone (e.g., 'America/New_York').
    """
    try:
        # Create timezone object
        tz = pytz.timezone(timezone)
        # Get current time in that timezone
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"The current local time in {timezone} is: {local_time}"
    except Exception as e:
        return f"Error fetching time for timezone '{timezone}': {str(e)}"


final_answer = FinalAnswerTool()

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

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)


# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, get_current_time_in_timezone,image_generation_tool,get_birth_stone,DuckDuckGoSearchTool(),], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


# GradioUI(agent).launch()


css = """
.message-row {
    justify-content: space-evenly !important;
}
.message-bubble-border {
    border-radius: 6px !important;
}
.message-buttons-bot, .message-buttons-user {
    right: 10px !important;
    left: auto !important;
    bottom: 2px !important;
}
.dark.message-bubble-border {
    border-color: #15172c !important;
}
.dark.user {
    background: #10132c !important;
}
.dark.assistant.dark, .dark.pending.dark {
    background: #020417 !important;
}
"""

# Create Gradio interface
gr.ChatInterface(
    predict,
    title=EMOJI + " " + MODEL_NAME,
    description=DESCRIPTION,
    examples=[
        ["Can you solve the equation 2x + 3 = 11 for x?"],
        ["Write an epic poem about Ancient Rome."],
        ["Who was the first person to walk on the Moon?"],
        ["Use a list comprehension to create a list of squares for numbers from 1 to 10."],
        ["Recommend some popular science fiction books."],
        ["Can you write a short story about a time-traveling detective?"]
    ],
    additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False),
    additional_inputs=[
        gr.Textbox("Perform the task to the best of your ability.", label="System prompt"),
        gr.Slider(0, 1, 0.8, label="Temperature"),
        gr.Slider(128, 4096, 512, label="Max new tokens"),
        gr.Slider(1, 80, 40, label="Top K sampling"),
        gr.Slider(0, 2, 1.1, label="Repetition penalty"),
        gr.Slider(0, 1, 0.95, label="Top P sampling"),
    ],
    theme=gr.themes.Soft(primary_hue=COLOR).set(
        background_fill_primary_dark="#020417",
        background_fill_secondary_dark="#020417",
        body_background_fill_dark="#020417",
        block_background_fill_dark="#020417",
        block_border_width="1px",
        block_title_background_fill_dark="#15172c",
        input_background_fill_dark="#15172c",
        button_secondary_background_fill_dark="#15172c",
        border_color_accent_dark="#15172c",
        border_color_primary_dark="#15172c",
        color_accent_soft_dark="#10132c",
        code_background_fill_dark="#15172c",
    ),
    css=css,
    retry_btn="Retry",
    undo_btn="Undo",
    clear_btn="Clear",
    submit_btn="Send",
    chatbot=gr.Chatbot(
        scale=1,
        show_copy_button=True
    )
).queue().launch()