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

from Gradio_UI import GradioUI

# Below is an example of a tool that does nothing. Amaze us with your creativity !
@tool
def my_custom_tool(arg1:str, arg2:int)-> 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 does nothing yet 
    Args:
        arg1: the first argument
        arg2: the second argument
    """
    return "What magic will you build ?"

@tool
def search_flights(departure: str, destination: str, date: str) -> str:
    """Finds flights from departure to destination on the given date using DuckDuckGo.
    
    Args:
        departure: The city or airport code where the flight starts.
        destination: The city or airport code where the flight ends.
        date: The departure date in YYYY-MM-DD format.

    Returns:
        A string containing flight search results.
    """
    query = f"flights from {departure} to {destination} on {date}"
    search_results = DuckDuckGoSearchTool()(query)  # Calls DuckDuckGo search
    return f"Here are some flight options:\n{search_results}"

API_KEY = os.getenv("FREECURRENCYAPI_KEY")

@tool
def convert_currency(amount: float, from_currency: str, to_currency: str) -> str:
    """Converts currency from one to another using FreeCurrencyAPI.
    
    Args:
        amount: The amount to convert.
        from_currency: The original currency (e.g., "USD").
        to_currency: The target currency (e.g., "EUR").

    Returns:
        The converted amount in the target currency.
    """
    try:
        url = f"https://api.freecurrencyapi.com/v1/latest?apikey={API_KEY}&base_currency={from_currency.upper()}"

        response = requests.get(url).json()

        # ✅ Check if the API returned valid exchange rates
        if "data" in response and to_currency.upper() in response["data"]:
            rate = response["data"][to_currency.upper()]
            converted_amount = amount * rate
            return f"{amount} {from_currency.upper()} is approximately {converted_amount:.2f} {to_currency.upper()}."

        return f"Error: Could not find exchange rate for {to_currency.upper()}."

    except Exception as e:
        return f"Error fetching exchange rates: {str(e)}"

chat_history = []  # Store conversation history

@tool
def chat_with_ai(message: str) -> str:
    """A tool that allows the AI to engage in general conversation with memory.
    Args:
        message: The user's message.
    """
    global chat_history  

    # Keep the last 5 messages for context
    if len(chat_history) > 5:
        chat_history.pop(0)

    # Add user message to history
    chat_history.append({"role": "user", "content": message})

    # Format the history as input for the AI model
    formatted_history = "\n".join(f"{msg['role']}: {msg['content']}" for msg in chat_history)

    # ✅ Ensure AI response is handled correctly
    try:
        response = model(formatted_history)  # Call model
        
        # ✅ Handle both string and dictionary responses
        if isinstance(response, dict):  
            response_text = response.get("text", str(response))  # Extract text if available
        else:
            response_text = str(response)  # Convert non-dict responses to string

        # Add AI response to history
        chat_history.append({"role": "assistant", "content": response_text})
        return response_text

    except Exception as e:
        return f"Error processing chat: {str(e)}"

@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=512,
temperature=0.5,
model_id='deepseek-ai/DeepSeek-R1-Distill-Qwen-32B',# 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,
        my_custom_tool,
        chat_with_ai,           # Regular chat tool
        search_flights,
        #scrape_webpage,
        convert_currency,
        #get_weather,
        #generate_ai_image
    ],
    max_steps=8,
    verbosity_level=5,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()