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

from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager

@tool
def get_zh_top_news() -> tuple[str, str]:
    """A tool that retrieves the current top news article's title and URL from www.zerohedge.com.
    
    Returns:
        tuple[str, str]: A tuple containing the article title (str) and its URL (str).

    Raises:
        Exception: If the page fails to load or the expected element is not found.
    """
    # Set up Chrome options for headless browsing
    chrome_options = Options()
    chrome_options.add_argument("--headless")  # Run without opening a browser window
    chrome_options.add_argument("--disable-gpu")  # Disable GPU for headless mode
    chrome_options.add_argument("--no-sandbox")  # Required for Docker environments
    chrome_options.add_argument("--disable-dev-shm-usage")  # Avoid memory issues in containers

    # Specify ChromeDriver path (installed via Dockerfile)
    service = Service(executable_path="/usr/local/bin/chromedriver")

    # Initialize the WebDriver
    driver = webdriver.Chrome(service=service, options=chrome_options)

    try:
        # Navigate to ZeroHedge homepage
        driver.get("https://www.zerohedge.com")

        # Find the first <h2> with class starting with 'Article_title___'
        top_article = driver.find_element(By.CSS_SELECTOR, "h2[class^='Article_title___']")

        print(f" article object: {top_article}")
        
        # Extract the title from the <a> tag inside the <h2>
        article_title = top_article.find_element(By.TAG_NAME, "a").text.strip()

        # Extract the URL from the href attribute of the <a> tag
        article_link = top_article.find_element(By.TAG_NAME, "a").get_attribute("href")
        
        # Ensure the link is absolute
        if not article_link.startswith("http"):
            article_link = f"https://www.zerohedge.com{article_link}"

        return article_title, article_link

    except Exception as e:
        print(f"Error retrieving top headline: {e}")
        return "Error: Headline not found", "https://www.zerohedge.com"

    finally:
        driver.quit()

@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:
        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)}"

final_answer = FinalAnswerTool()

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

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=[get_zh_top_news, get_current_time_in_timezone, final_answer],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()