from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
from bs4 import BeautifulSoup
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
import re
from typing import Optional

from Gradio_UI import GradioUI

@tool
def get_pets_available_for_adoption(animal_type: Optional[str] = None, breed: Optional[str] = None) -> str:
    """A tool that finds the specified animal and breed from the Brampton Animal Shelter website.
    
    Args:
        animal_type: The type of animal to search for (e.g., "dog", "cat", "bird", "other"). If None, returns all types.
        breed: The specific breed or species name to filter by. If None, returns all breeds.
    """
    
    base_url = "https://www.brampton.ca/EN/residents/Animal-Services/Pages/Adoption.aspx"
    
    matching_animals = []
    processed_ids = {}
    
    session = requests.Session()
    response = session.get(base_url)
    if response.status_code != 200:
        return f"Failed to access the website. Status code: {response.status_code}"
    
    # Only process the main page, no pagination
    process_page(response.text, matching_animals, processed_ids, animal_type, breed)
    
    animal_sentences = []
    for animal in matching_animals:
        if (animal.get("name") in [None, "", "Unknown"] or
            animal.get("age") in [None, "", "unknown age"] or 
            animal.get("sex") in [None, "", "unknown gender"] or
            animal.get("breed") in [None, "", "unknown breed"] or
            animal.get("type") in [None, "", "animal"]):
            continue

        name = animal.get("name", "Unknown").title()
        age = animal.get("age", "unknown age").lower().strip()
        sex = animal.get("sex", "unknown gender").lower()
        breed_name = animal.get("breed", "unknown breed").title()
        animal_type_value = animal.get("type", "animal").lower()
        
        sentence = f"A {age} old {sex} {breed_name} named {name}, type: {animal_type_value}"
        animal_sentences.append(sentence)
    
    if animal_sentences:
        return "\n".join(animal_sentences)
    else:
        return "No matching animals found."

def process_page(html_content, matching_animals, processed_ids, animal_type=None, breed=None):
    """Process a single page of animal listings"""
    soup = BeautifulSoup(html_content, 'html.parser')
    
    animal_cards = soup.find_all("div", class_="dvanimalitem")
    
    for card in animal_cards:
        animal_info = {}
        
        name_element = card.find("span", class_="animalname")
        if name_element:
            animal_info["name"] = name_element.text.strip()
        
        for field in ["id", "type", "sex", "breed", "age", "size", "color", "location"]:
            field_element = card.find("span", id=lambda x: x and x.startswith(f"sp{field}"))
            if field_element:
                animal_info[field] = field_element.text.strip()
        
        if animal_info.get("id") in processed_ids:
            continue
        
        if "id" in animal_info:
            processed_ids[animal_info["id"]] = True
        
        # Matching logic
        type_match = True
        breed_match = True
        
        if animal_type:
            animal_type_lower = animal_type.lower()
            current_type_lower = animal_info.get("type", "").lower()
            if animal_type_lower == "other":
                type_match = current_type_lower not in ["dog", "cat", "bird"]
            else:
                type_match = animal_type_lower in current_type_lower
        
        if breed:
            breed_match = breed.lower() in animal_info.get("breed", "").lower()
        
        if type_match and breed_match:
            matching_animals.append(animal_info)

@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()

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
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_pets_available_for_adoption, get_current_time_in_timezone], ## 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()