from smolagents import CodeAgent, tool, OpenAIServerModel, FinalAnswerTool,ToolCallingAgent
import datetime
import pytz
import yfinance as yf
import os
from Gradio_UI import GradioUI

# Fetch the API key from the environment variable (if needed)
API_KEY = os.getenv("KEY")

@tool
def get_currency_rate(currency_pair: str) -> str:
    """Fetches the current exchange rate for a given currency pair (e.g., 'USD-EUR').
    Args:
        currency_pair: The currency pair to query, formatted as 'BASE-CURRENCY'.
    """
    try:
        base, target = currency_pair.split("-")
        # Use yfinance to get the exchange rate
        ticker = f"{base}{target}=X"  # Format for yfinance (e.g., "USDEUR=X")
        data = yf.Ticker(ticker)
        rate = data.history(period="1d")["Close"].iloc[-1].astype(str)
        return f"Current {currency_pair} rate: {rate}"
    except Exception as e:
        return f"Error fetching rate: {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)}"

# Set up the model
model = OpenAIServerModel(
    model_id="deepseek-chat",
    api_base="https://api.deepseek.com",
    api_key=API_KEY
)

# Load tools
final_answer = FinalAnswerTool()
currency_tool = get_currency_rate
time_tool = get_current_time_in_timezone

# Create the agent
agent = ToolCallingAgent(
    model=model,
    tools=[currency_tool, time_tool, final_answer],
    max_steps=2,
    verbosity_level=1
)

example_usage = """
### Example Usage
1. **Currency Rate**: Ask for the exchange rate of a currency pair.
   - Example: "What is the current exchange rate for USD to JPY?"
   - Example: "Get the EUR to GBP rate."

2. **Current Time**: Ask for the current time in a specific timezone.
   - Example: "What is the current time in America/New_York?"
   - Example: "Get the time in Asia/Tokyo."
"""

# Extend the GradioUI class to include the example usage
class CustomGradioUI(GradioUI):
    def __init__(self, agent, description: str = "", **kwargs):
        super().__init__(agent, **kwargs)
        self.description = description

    def launch(self, **kwargs):
        import gradio as gr

        with gr.Blocks(fill_height=True) as demo:
            # Add the description as a Markdown component
            if self.description:
                gr.Markdown(self.description)

            # Add the existing chatbot and input components
            stored_messages = gr.State([])
            file_uploads_log = gr.State([])
            chatbot = gr.Chatbot(
                label="Agent",
                type="messages",
                avatar_images=(
                    None,
                    "https://huggingface.co/datasets/agents-course/course-images/resolve/main/en/communication/Alfred.png",
                ),
                resizeable=True,
                scale=1,
            )
            # If an upload folder is provided, enable the upload feature
            if self.file_upload_folder is not None:
                upload_file = gr.File(label="Upload a file")
                upload_status = gr.Textbox(label="Upload Status", interactive=False, visible=False)
                upload_file.change(
                    self.upload_file,
                    [upload_file, file_uploads_log],
                    [upload_status, file_uploads_log],
                )
            text_input = gr.Textbox(lines=1, label="Chat Message")
            text_input.submit(
                self.log_user_message,
                [text_input, file_uploads_log],
                [stored_messages, text_input],
            ).then(self.interact_with_agent, [stored_messages, chatbot], [chatbot])

        demo.launch(debug=True, share=True, **kwargs)

# Launch the Gradio UI with the example usage
CustomGradioUI(agent, description=example_usage).launch()