from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import datetime
import requests
import pytz
import yaml
import tempfile
import os
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

@tool
def my_custom_tool(arg1: str, arg2: int) -> str:
    """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:
        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)}"

@tool
def generate_and_save_image(prompt: str, filename: str = "generated_image.png") -> str:
    """Generate an image and save it to file for better display
    Args:
        prompt: Description of the image to generate
        filename: Name of the file to save (optional)
    """
    try:
        # Verwende das geladene image_generation_tool
        image = image_generation_tool(prompt=prompt)
        
        # Speichere in temporärem Verzeichnis für bessere Organisation
        temp_dir = tempfile.gettempdir()
        full_path = os.path.join(temp_dir, filename)
        
        # Speichere das Bild
        image.save(full_path)
        
        # Erstelle auch eine Kopie im aktuellen Verzeichnis für einfachen Zugriff
        image.save(filename)
            
        return f"Image generated and saved as {filename} (also available at {full_path}). You can download it from the current directory."
    except Exception as e:
        return f"Error generating image: {str(e)}"

# Load tools
final_answer = FinalAnswerTool()
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

# Debug: Prüfe den Namen des geladenen Tools
print(f"Loaded image tool name: {image_generation_tool.name}")

# Erstelle einen Enhanced Image Generator mit Download-Funktionalität
@tool
def enhanced_image_generator(prompt: str, filename: str = None) -> str:
    """Enhanced image generator that creates images and provides download information
    Args:
        prompt: Text description of the image to generate
        filename: Optional custom filename (without extension)
    """
    try:
        image = image_generation_tool(prompt=prompt)
        
        # Generiere Dateinamen falls nicht angegeben
        if filename is None:
            timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
            filename = f"generated_image_{timestamp}"
        
        # Stelle sicher dass Dateiname .png Extension hat
        if not filename.endswith('.png'):
            filename += '.png'
        
        # Speichere das Bild für Download
        image.save(filename)
        
        # Gib eine klare Download-Anweisung zurück
        return f"✅ Image generated successfully!\n📁 File saved as: {filename}\n💡 You can find and download the image file '{filename}' from your current directory.\n🎨 Prompt used: '{prompt}'"
    except Exception as e:
        return f"❌ Error generating image: {str(e)}"

@tool  
def quick_image_generator(prompt: str) -> str:
    """Quick image generator with automatic filename
    Args:
        prompt: Text description of the image to generate
    """
    try:
        image = image_generation_tool(prompt=prompt)
        
        # Automatischer Dateiname basierend auf Prompt (erste 3 Wörter)
        words = prompt.split()[:3]
        safe_name = "_".join(word.lower().replace(",", "").replace(".", "") for word in words)
        timestamp = datetime.datetime.now().strftime('%H%M%S')
        filename = f"{safe_name}_{timestamp}.png"
        
        # Speichere das Bild
        image.save(filename)
        
        return f"🎨 Image created and saved as '{filename}'! You can download it from your file directory. Generated from: '{prompt}'"
    except Exception as e:
        return f"❌ Error generating image: {str(e)}"

# Test das Tool beim Start (optional - kann entfernt werden)
print("Testing image generation tool...")
try:
    test_result = image_generation_tool(prompt="test")
    print(f"Image tool loaded successfully: {type(test_result)}")
except Exception as e:
    print(f"Warning: Image tool test failed: {e}")

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,
)

# Load system prompt from prompt.yaml file
with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[
        final_answer, 
        image_generation_tool,  # Das Original-Tool
        enhanced_image_generator,  # Tool mit besseren Download-Infos
        quick_image_generator,  # Schnelle Bildgenerierung mit automatischen Namen
        get_current_time_in_timezone,
        my_custom_tool
    ],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)

# Debugging Info
print("Agent initialized with tools:")
for tool in agent.tools:
    print(f"- {tool.name if hasattr(tool, 'name') else str(tool)}")

print("\nStarting Gradio UI...")
GradioUI(agent).launch()

# Beispiel für direkten Test (kann auskommentiert werden):
"""
# Direkter Test ohne UI:
if __name__ == "__main__":
    print("Testing agent directly...")
    response = agent.run("Generiere ein Bild eines Fisches")
    print(f"Agent response: {response}")
"""