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

# --- FERRAMENTAS PERSONALIZADAS ---

@tool
def get_time_diff_between_timezones(timezone_a: str, timezone_b: str) -> str:
    """Calcula a diferença de horas entre duas timezones quaisquer.
    Args:
        timezone_a: A primeira timezone (ex: 'Europe/Lisbon').
        timezone_b: A segunda timezone para comparar (ex: 'Asia/Tokyo').
    """
    try:
        # Definir as timezones
        tz_a = pytz.timezone(timezone_a)
        tz_b = pytz.timezone(timezone_b)
        
        # Obter a data atual
        now = datetime.datetime.now()
        
        # Calcular os offsets em relação ao UTC
        offset_a = tz_a.localize(now).utcoffset().total_seconds() / 3600
        offset_b = tz_b.localize(now).utcoffset().total_seconds() / 3600
        
        diff = offset_a - offset_b
        
        if diff == 0:
            return f"Não há diferença horária entre {timezone_a} e {timezone_b}."
        
        # Lógica para explicar a diferença de forma humana
        if diff > 0:
            return f"{timezone_a} está {abs(diff)} horas à frente de {timezone_b}."
        else:
            return f"{timezone_a} está {abs(diff)} horas atrás de {timezone_b}."

    except Exception as e:
        return f"Erro ao calcular a diferença entre {timezone_a} e {timezone_b}: {str(e)}"

@tool
def get_current_time_in_timezone(timezone: str) -> str:
    """Busca a hora local atual numa timezone específica.
    Args:
        timezone: Uma string de timezone válida (ex: 'Europe/Lisbon').
    """
    try:
        tz = pytz.timezone(timezone)
        local_time = datetime.datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"A hora atual em {timezone} é: {local_time}"
    except Exception as e:
        return f"Erro ao obter a hora para '{timezone}': {str(e)}"

# --- INICIALIZAÇÃO ---

final_answer = FinalAnswerTool()
search_tool = DuckDuckGoSearchTool() 

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

# RECOMENDAÇÃO: Se o ficheiro prompts.yaml estiver a causar o KeyError 'final_answer', 
# o melhor é não o carregar até que o ficheiro esteja corrigido.
prompt_templates = None
try:
    with open("prompts.yaml", 'r') as stream:
        prompt_templates = yaml.safe_load(stream)
except Exception:
    print("Aviso: prompts.yaml não encontrado ou inválido. Usando templates padrão.")

# --- CRIAÇÃO DO AGENTE ---

agent = CodeAgent(
    model=model,
    tools=[
        final_answer, 
        search_tool,      
        get_time_diff_between_timezones, # Substitui a antiga por esta
        get_current_time_in_timezone
    ], 
    max_steps=10, 
    verbosity_level=1,
    prompt_templates=prompt_templates if prompt_templates else None 
)

if __name__ == "__main__":
    GradioUI(agent).launch()