from smolagents import CodeAgent, HfApiModel, tool
import datetime
import pytz
import yaml
import os
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI
from transformers import Tool  # Import the Tool class
from PIL import Image

hf_token = os.getenv("HF_ACCESS_TOKEN")


@tool
def generate_landscape_image(timezone: str) -> Image.Image:
    """
    A tool that generates a landscape image based on the local time in a given timezone.
    If it's daytime, the image will feature the sun in a position corresponding to the current time.
    If it's nighttime, the image will feature the moon without a preferred position.
    
    Args:
        timezone: A valid timezone string (e.g., 'America/New_York').
        
    Returns:
        A PIL.Image object of the generated image or an error message as a string.
    """
    try:
        # Create timezone object and get current time in that zone
        tz = pytz.timezone(timezone)
        now = datetime.datetime.now(tz)
        hour = now.hour
        
        # Build prompt based on day or night
        if 6 <= hour < 18:
            if hour < 12:
                position = "rising in the east"
            elif hour == 12:
                position = "directly overhead"
            else:
                position = "setting in the west"
            prompt = (
                f"A breathtaking daytime landscape with the sun {position} at {now.strftime('%H:%M')}, "
                "illuminating the scene with warm, golden light."
            )
        else:
            prompt = "A serene nighttime landscape featuring a luminous moon shining in the dark sky."
        
        # Call the image generation tool loaded from Space
        result = image_generation_tool(prompt)
        
        # If the result is an AgentImage, extract the underlying image.
        if hasattr(result, "to_raw"):
            image = result.to_raw()
        else:
            # Otherwise assume it's a file path and open it.
            image = Image.open(result)
        return image
    except Exception as e:
        return f"Error generating image for timezone '{timezone}': {str(e)}"

@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 valid timezone string (e.g., 'America/New_York')
    
    Returns:
        A string with the current local time or an error message.
    """
    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)}"

final_answer = FinalAnswerTool()

model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
    custom_role_conversions=None,
    token=hf_token
)

# Load the image generation tool directly from the public FLUX Space
image_generation_tool = Tool.from_space(
    space_id="black-forest-labs/FLUX.1-dev",
    name="image_generator",
    description="Generate an image from a prompt"
)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer, generate_landscape_image, get_current_time_in_timezone],
    max_steps=6,
    verbosity_level=1,
    prompt_templates=prompt_templates
)

GradioUI(agent).launch()