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

# 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 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)}"


@tool
def calculate_age(birth_date: str, reference_date: str = "today") -> str:
    """Calculates precise age in years, months, and days from birth date to reference date.
    Handles multiple date formats (YYYY-MM-DD, MM/DD/YYYY, DD.MM.YYYY, etc.).
    
    Args:
        birth_date: The birth date in any common format (e.g., "1990-05-15", "05/15/1990", "15.05.1990")
        reference_date: The date to calculate age at (defaults to today). Use "today" or specific date in same formats.
    """

    from datetime import datetime
    import re
    
    def parse_date(date_str):
        """Parse various date formats into datetime object"""
        if date_str.lower() == "today":
            return datetime.today()
            
        # Try different date patterns
        patterns = [
            (r"(\d{4})-(\d{1,2})-(\d{1,2})", "%Y-%m-%d"),  # YYYY-MM-DD
            (r"(\d{1,2})/(\d{1,2})/(\d{4})", "%m/%d/%Y"),    # MM/DD/YYYY
            (r"(\d{1,2})\.(\d{1,2})\.(\d{4})", "%d.%m.%Y"),  # DD.MM.YYYY
            (r"(\d{1,2}) (\w{3,}) (\d{4})", "%d %B %Y"),     # 15 May 1990
            (r"(\w{3,}) (\d{1,2}), (\d{4})", "%B %d, %Y"),   # May 15, 1990
        ]
        
        for pattern, fmt in patterns:
            if re.fullmatch(pattern, date_str):
                try:
                    return datetime.strptime(date_str, fmt)
                except ValueError:
                    continue
                    
        raise ValueError(f"Could not parse date: {date_str}")

    try:
        # Parse dates
        birth_dt = parse_date(birth_date)
        ref_dt = parse_date(reference_date)
        
        # Calculate age
        if ref_dt < birth_dt:
            return "Error: Reference date is before birth date"
            
        years = ref_dt.year - birth_dt.year
        months = ref_dt.month - birth_dt.month
        days = ref_dt.day - birth_dt.day
        
        # Handle negative months/days
        if days < 0:
            months -= 1
            # Get days in previous month
            if ref_dt.month == 1:
                prev_month = 12
                prev_year = ref_dt.year - 1
            else:
                prev_month = ref_dt.month - 1
                prev_year = ref_dt.year
            days_in_prev_month = (datetime(prev_year, prev_month + 1, 1) - datetime(prev_year, prev_month, 1)).days
            days += days_in_prev_month
            
        if months < 0:
            years -= 1
            months += 12
            
        # Format output
        parts = []
        if years > 0:
            parts.append(f"{years} year{'s' if years != 1 else ''}")
        if months > 0:
            parts.append(f"{months} month{'s' if months != 1 else ''}")
        if days > 0:
            parts.append(f"{days} day{'s' if days != 1 else ''}")
            
        if not parts:
            return "0 days old (born today)"
            
        return ", ".join(parts) + " old"
        
    except Exception as e:
        return f"Error calculating age: {str(e)}. Please check date formats."


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, calculate_age], ## 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()