from transformers import HfAgent
from PIL import Image, ImageDraw, ImageFont
import gradio as gr
import os

#%% Utility Functions

def add_label_to_image(image, label):
    draw = ImageDraw.Draw(image)
    font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
    font_size = 30
    try:
        font = ImageFont.truetype(font_path, font_size)
    except:
        font = ImageFont.load_default()
    
    text_bbox = draw.textbbox((0, 0), label, font=font)
    text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
    position = (image.width - text_width - 20, image.height - text_height - 20)
    
    rect_margin = 10
    rect_position = [
        position[0] - rect_margin,
        position[1] - rect_margin,
        position[0] + text_width + rect_margin,
        position[1] + text_height + rect_margin,
    ]
    draw.rectangle(rect_position, fill=(0, 0, 0, 128))
    draw.text(position, label, fill="white", font=font)
    return image

def plot_and_save_agent_image(image, label, save_path=None):
    labeled_image = add_label_to_image(image, label)
    labeled_image.show()
    if save_path:
        labeled_image.save(save_path)
        print(f"Image saved to {save_path}")

def generate_prompts_for_object(object_name):
    return {
        "past": f"Show an old version of a {object_name} from its early days.",
        "present": f"Show a {object_name} with current features/design/technology.",
        "future": f"Show a futuristic version of a {object_name}, predicting advanced features and design."
    }

#%% HF Agent Initialization

agent = HfAgent("https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta")

#%% Core Generation Function

def generate_object_history(object_name):
    images = []
    gif_frames = []

    prompts = generate_prompts_for_object(object_name)
    labels = {
        "past": f"{object_name} - Past",
        "present": f"{object_name} - Present",
        "future": f"{object_name} - Future"
    }

    output_image_paths = []

    for period, prompt in prompts.items():
        print(f"Generating image for: {prompt}")
        result = agent.run(prompt)
        
        # Find first image in result (agent returns a dict)
        image_path = next((v for v in result.values() if isinstance(v, str) and v.endswith((".png", ".jpg"))), None)
        if not image_path or not os.path.exists(image_path):
            raise RuntimeError(f"No valid image generated for {prompt}")

        image = Image.open(image_path).convert("RGB")
        labeled_image = add_label_to_image(image, labels[period])

        filename = f"{object_name}_{period}.png"
        labeled_image.save(filename)
        output_image_paths.append((filename, labels[period]))
        gif_frames.append(labeled_image)

    # Save animated GIF
    gif_path = f"{object_name}_evolution.gif"
    gif_frames[0].save(gif_path, save_all=True, append_images=gif_frames[1:], duration=1000, loop=0)

    return output_image_paths, gif_path

#%% Gradio Interface

def create_gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# TimeMetamorphy: An Object Evolution Generator")

        gr.Markdown("""
        Explore how objects change over time — from past, to present, to future.
        Enter any object below and let AI visualize its transformation through the ages.
        """)

        default_images = [
            ("car_past.png", "Car - Past"),
            ("car_present.png", "Car - Present"),
            ("car_future.png", "Car - Future")
        ]
        default_gif_path = "car_evolution.gif"

        with gr.Row():
            with gr.Column():
                object_name_input = gr.Textbox(label="Enter an object name", placeholder="e.g., bicycle, phone")
                generate_button = gr.Button("Generate Evolution")
                image_gallery = gr.Gallery(label="Generated Images", show_label=True, columns=3, rows=1, value=default_images)
                gif_output = gr.Image(label="Generated GIF", show_label=True, value=default_gif_path)

        generate_button.click(fn=generate_object_history, inputs=[object_name_input], outputs=[image_gallery, gif_output])

    return demo

# Launch
demo = create_gradio_interface()
demo.launch(share=True)