from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool
# Importante: Asumiendo que GradioUI permite un argumento de título.
from Gradio_UI import GradioUI 

# Inicialización de la herramienta de búsqueda
duckduckgo_tool = DuckDuckGoSearchTool()

# --- Herramientas personalizadas (Solo queda la de compositores) ---

@tool
def get_composer_lifespan_from_web(composer_lastname: str) -> str:
    """Busca el nombre completo, el año de nacimiento y el año de fallecimiento de un compositor de música clásica.
    
    Args:
        composer_lastname: El apellido del compositor (p. ej., 'Beethoven', 'Mozart').
    
    INSTRUCCIÓN ESTRICTA: El Agente debe utilizar los resultados de la búsqueda en bruto para extraer el nombre, el año de nacimiento (AAAA) y el año de fallecimiento (AAAA) y generar la respuesta FINAL y ÚNICA utilizando FinalAnswerTool con el siguiente formato, incluyendo el nombre completo y los saltos de línea (\\n) entre campos:
    
    Nombre: [Nombre Completo del Compositor]
    Fecha de nacimiento: [AAAA]
    Fecha de fallecimiento: [AAAA]
    """
    global duckduckgo_tool
    
    query = f"classical composer {composer_lastname} full name birth and death year"
    
    search_result = duckduckgo_tool(query=query)
    
    # La clave está en la interpretación del Agente de la instrucción estricta en el docstring.
    return f"Consulta de búsqueda: '{query}'. \n\nResultado de búsqueda en bruto:\n\n{search_result}"

# --- Inicialización del Agente y Gradio ---

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,
        # La lista de herramientas ahora solo incluye las esenciales y la de compositores
        tools=[final_answer, duckduckgo_tool, get_composer_lifespan_from_web],
        max_steps=6,
        verbosity_level=1,
        grammar=None,
        planning_interval=None,
        name=None,
        description=None,
        prompt_templates=prompt_templates
    )

# 2. Ajuste del GradioUI para el título:
custom_title = "Nacimiento y Muerte de Compositores. Introduce el apellido del compositor:"

# Si GradioUI no soporta directamente el argumento 'title', tendrás que modificar el archivo Gradio_UI.py.
GradioUI(agent, title=custom_title).launch()