import uuid
import os
from PIL import Image
import gradio as gr


COMMON_RESOLUTIONS = {
    "16:9": (1920, 1080),
    "9:16": (1080, 1920),
    "4:3": (1440, 1080),
    "3:4": (1080, 1440),
    "1:1": (1080, 1080),
    "21:9": (2560, 1080),
    "2:1": (2160, 1080),
    "3:2": (1620, 1080),
    "5:4": (1350, 1080),
    "4:5": (864, 1080),
}


os.makedirs("./download", exist_ok=True)

def prompt_maker(aspect_ratio):
  prompt=f"""Take the subject from the first image and place it inside the second image canvas. Expand the background of the first image naturally to fill the extra space of the second image. Keep the subject exactly the same, without cropping or distortion. Extend the background seamlessly in a photorealistic way, matching the original style, details, colors, lighting, and textures. Ensure smooth blending with no visible edges, borders, or repetition. Final result should be in perfect {aspect_ratio}."""
  return prompt
def generate_aspect_image(aspect_ratio):
    """Generate image from preset aspect ratio"""
    width, height = COMMON_RESOLUTIONS[aspect_ratio]
    img = Image.new("RGBA", (width, height), (0, 0, 0, 0))
    # random_str=str(uuid.uuid4())[:5]
    width,height=aspect_ratio.split(":")
    random_str=f"{width}_{height}"
    filename = f"./download/{random_str}.png"
    img.save(filename, "PNG")
    aspect_ratio=f"[{aspect_ratio}] ratio"
    prompt=prompt_maker(aspect_ratio)
    return prompt,filename

def generate_custom_image(width, height):
    """Generate image from custom width/height"""
    try:
        width, height = int(width), int(height)
    except ValueError:
        return None
    if width <= 0 or height <= 0:
        return None
    img = Image.new("RGBA", (width, height), (0, 0, 0, 0))
    # random_str=str(uuid.uuid4())[:5]
    random_str=f"{width}_{height}"
    filename = f"./download/{random_str}.png"
    img.save(filename, "PNG")
    aspect_ratio=f"[{width}x{height}] pixels"
    prompt=prompt_maker(aspect_ratio)
    return prompt,filename

with gr.Blocks() as demo:
    
    gr.HTML("""
            <div style="text-align: center; margin: 20px auto; max-width: 800px;">
                <h1 style="font-size: 2.5em; margin-bottom: 5px;">🍌 Nano Banana Aspect Ratio Change</h1>
            </div>""")
    with gr.Tab("Aspect Ratio Presets"):
        with gr.Row():
            with gr.Column(scale=1):
                ratio_input = gr.Dropdown(choices=list(COMMON_RESOLUTIONS.keys()), label="Aspect Ratio")
                ratio_btn = gr.Button("Generate")
            with gr.Column(scale=2):
                text=gr.Textbox(label="Prompt",lines=4,show_copy_button=True)
                ratio_output = gr.Image(type="filepath", label="Generated Image")
        ratio_btn.click(fn=generate_aspect_image, inputs=ratio_input, outputs=[text,ratio_output])

    with gr.Tab("Custom Size"):
        with gr.Row():
            with gr.Column(scale=1):
                width_input = gr.Number(label="Width", value=1920)
                height_input = gr.Number(label="Height", value=1080)
                custom_btn = gr.Button("Generate")
            with gr.Column(scale=2):
                text=gr.Textbox(label="Prompt",lines=4,show_copy_button=True)
                custom_output = gr.Image(type="filepath", label="Generated Image")
        custom_btn.click(fn=generate_custom_image, inputs=[width_input, height_input], outputs=[text,custom_output])
    gr.Markdown("""
    ### ⚠️ Problem  
    If you upload your photo and ask Nano Banana to change the **aspect ratio** of the image, it won’t work directly.  

    ### 📝 Trick to Change the Aspect Ratio of an Uploaded Image  
    1. Open **Google AI Studio** — [Go to AI Studio](https://aistudio.google.com/) and choose **Try Nano Banana**.  
    2. Upload your **original photo** as the **first image**.  
    3. Upload the **blank transparent PNG** from this tool as the **second image**.  
    4. Copy & paste the provided **prompt**.  
    5. You’ll get your photo in the **desired aspect ratio** with the background expanded.  
    """)

if __name__ == "__main__":
    demo.launch(share=False,debug=False)