from smolagents import CodeAgent, HfApiModel, DuckDuckGoSearchTool, VisitWebpageTool, load_tool, tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
import xml.etree.ElementTree as ET

@tool
def pull_ecb_eur_usd_rate(requested_date: str = None) -> str:
    """
    A tool that fetches the EUR-to-USD exchange rate from the European Central Bank for today.
    
    Args:
        requested_date: (Optional) A string in "YYYY-MM-DD" format. If provided and it is not equal to today's date,
                        the function returns an error.
                        
    Returns:
        A string indicating the exchange rate as of the ECB data's date, or an error message if:
          - the requested date is not today, or
          - the ECB data is not updated for today.
    """
    try:
        # Determine today's date in the ECB's (Central European) timezone.
        brussels_tz = pytz.timezone("Europe/Brussels")
        today = datetime.datetime.now(brussels_tz).date()

        # If a date is provided, check that it matches today's date.
        if requested_date:
            try:
                req_date = datetime.datetime.strptime(requested_date, "%Y-%m-%d").date()
            except Exception:
                return f"Error: The provided date '{requested_date}' is not in the correct format (YYYY-MM-DD)."
            if req_date != today:
                return f"Error: The requested date {req_date} is not available. Only today's data ({today}) is supported."

        # Fetch the ECB daily XML data.
        url = "https://www.ecb.europa.eu/stats/eurofxref/eurofxref-daily.xml"
        response = requests.get(url)
        if response.status_code != 200:
            return f"Error: Received status code {response.status_code} from ECB."

        xml_content = response.content
        tree = ET.ElementTree(ET.fromstring(xml_content))
        root = tree.getroot()

        # Define namespaces used in the ECB XML.
        ns = {
            'gesmes': 'http://www.gesmes.org/xml/2002-08-01',
            'def': 'http://www.ecb.int/vocabulary/2002-08-01/eurofxref'
        }

        # Locate the Cube element with the time attribute.
        time_cube = root.find('.//def:Cube/def:Cube', ns)
        if time_cube is not None:
            data_date = time_cube.get('time')
            # If ECB's data date is not today, return an error.
            if data_date != str(today):
                return f"Error: ECB data is not updated for today. Latest available data is for {data_date}."
            
            # Find the Cube element for USD.
            usd_cube = time_cube.find("def:Cube[@currency='USD']", ns)
            if usd_cube is not None:
                rate = usd_cube.get('rate')
                return f"As of {data_date}, EUR 1 is equivalent to {rate} USD."
        
        return "Error: USD rate not found in ECB data."
    except Exception as e:
        return f"Error fetching ECB 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').
    
    Returns:
        A string with the current local time in the specified timezone.
    """
    try:
        tz = pytz.timezone(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)}"

# Tools for final answer and additional capabilities.
final_answer = FinalAnswerTool()
search_result = DuckDuckGoSearchTool()
visit_webpage = VisitWebpageTool()

# Initialize the model.
model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

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

# Load prompt templates from the YAML file.
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
# Create the CodeAgent and register the tools (including the updated ECB rate tool).
agent = CodeAgent(
    model=model,
    tools=[final_answer, get_current_time_in_timezone, pull_ecb_eur_usd_rate],
    max_steps=10,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

# Launch the Gradio UI for interacting with the agent.
GradioUI(agent).launch()