import torch
import gradio as gr
from diffusers import StableDiffusionPipeline
from PIL import Image, ImageDraw, ImageFont
import os

# Function to Generate Image
def generate_image(prompt, height=1024, width=1024):
    """
    Generate an image using Stable Diffusion from Hugging Face.

    Args:
        prompt (str): The text prompt to guide image generation.
        height (int): Height of the generated image.
        width (int): Width of the generated image.

    Returns:
        PIL Image: Generated image.
    """
    # Load Hugging Face access token from environment variables
    HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACE_API_KEY")
    if not HUGGINGFACE_API_KEY:
        raise ValueError("Hugging Face API key is not set. Export it as HUGGINGFACE_API_KEY.")
    
    # Load the Stable Diffusion model
    print("Loading Stable Diffusion model...")
    model_id = "CompVis/stable-diffusion-v1-4"  # You can change to v2 models if needed
    pipeline = StableDiffusionPipeline.from_pretrained(
        model_id, 
        use_auth_token=HUGGINGFACE_API_KEY,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
    )
    
    # Move to GPU if available
    device = "cuda" if torch.cuda.is_available() else "cpu"
    pipeline = pipeline.to(device)
    
    # Generate the image
    print(f"Generating image for prompt: '{prompt}'...")
    image = pipeline(prompt, height=height, width=width).images[0]
    
    return image

# Function to Add Emphasized Text to the Image (Product Name, Tagline, CTA)
def add_text_to_image(image, product_name, tagline, cta_text, font_size=50):
    """
    Add clean and sharp text (product name, tagline, and CTA) to the generated image with emphasis.
    
    Args:
        image (PIL Image): Generated image to add text to.
        product_name (str): Product name to be emphasized.
        tagline (str): Tagline to be emphasized.
        cta_text (str): Call to action text to be emphasized.
        font_size (int): The font size.
    
    Returns:
        PIL Image: Image with added text.
    """
    draw = ImageDraw.Draw(image)
    
    try:
        # Load font for product name, tagline, and CTA
        product_font = ImageFont.truetype("arial.ttf", font_size + 20)  # Larger for product name
        tagline_font = ImageFont.truetype("arial.ttf", font_size)       # Slightly smaller for tagline
        cta_font = ImageFont.truetype("arial.ttf", font_size - 10)      # Smaller for CTA
    except IOError:
        # Fallback font in case "arial.ttf" is not available
        product_font = ImageFont.load_default()
        tagline_font = ImageFont.load_default()
        cta_font = ImageFont.load_default()

    # Define positions for text
    product_name_position = (50, 50)
    tagline_position = (50, 150)
    cta_position = (50, 250)

    # Add product name in large font
    draw.text(product_name_position, product_name, font=product_font, fill="white")
    
    # Add tagline below the product name
    draw.text(tagline_position, tagline, font=tagline_font, fill="white")
    
    # Add CTA below tagline
    draw.text(cta_position, cta_text, font=cta_font, fill="gold")  # CTA in gold to stand out
    
    return image

# Main function to generate advertisement
def generate_ad(brand_title, tagline, cta, brand_logo=None, product_image=None, custom_prompt=None):
    """
    Generate an advertisement image with the provided details.
    
    Args:
        brand_title (str): Brand title for the advertisement.
        tagline (str): Tagline for the advertisement.
        cta (str): Call to action text.
        brand_logo (file, optional): Brand logo image (optional).
        product_image (file, optional): Product image (optional).
        custom_prompt (str, optional): Custom prompt for image generation.

    Returns:
        PIL Image: Final advertisement image.
    """
    # Prepare the final prompt
    prompt = custom_prompt if custom_prompt else (
        f"An elegant advertisement for {brand_title}, featuring gold and white tones, "
        f"with a radiant and premium look. Product focus and beautiful typography for '{tagline}'."
    )

    # Generate the image
    generated_image = generate_image(prompt)

    # Add product and CTA text to the image
    final_image = add_text_to_image(generated_image, brand_title, tagline, cta)

    # Optionally, add brand logo and product image if provided
    if brand_logo:
        logo = Image.open(brand_logo).resize((150, 150))
        final_image.paste(logo, (50, 350), logo.convert('RGBA'))

    if product_image:
        product = Image.open(product_image).resize((300, 300))
        final_image.paste(product, (250, 350), product.convert('RGBA'))
    
    return final_image

# Gradio Interface
iface = gr.Interface(
    fn=generate_ad,
    inputs=[
        gr.Textbox(label="Brand Title", placeholder="Enter brand title"),
        gr.Textbox(label="Tagline", placeholder="Enter tagline"),
        gr.Textbox(label="CTA", placeholder="Enter Call to Action"),
        gr.Image(type="file", label="Brand Logo (optional)", optional=True),
        gr.Image(type="file", label="Product Image (optional)", optional=True),
        gr.Textbox(label="Custom Prompt (optional)", placeholder="Enter a custom prompt (or leave empty)"),
    ],
    outputs=gr.Image(label="Generated Advertisement"),
)

# Launch the Gradio Interface
iface.launch()