import gradio as gr
import requests
import random
from io import BytesIO
from PIL import Image
import os

# Constants
BASE_URL = "https://image.pollinations.ai/prompt/"
WATERMARK_PATH = "CodeHubb.png"

# Load and prepare watermark
watermark = Image.open(WATERMARK_PATH).convert("RGBA")
watermark = watermark.resize((300, 100))
watermark.putalpha(128)  # 50% opacity

def generate_image(prompt):
    if not prompt:
        raise gr.Error("Please enter a prompt.")

    seed = random.randint(1, 999999)
    url = f"{BASE_URL}{prompt}?width=2048&height=2048&seed={seed}&nologo=true&model=flux"

    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()
        img = Image.open(BytesIO(response.content)).convert("RGBA")

        # Paste watermark in top-left
        img.paste(watermark, (20, 20), watermark)

        # Ensure /mnt/data directory
        os.makedirs("/mnt/data", exist_ok=True)

        # Save as JPG
        rgb_img = img.convert("RGB")
        download_path = "/mnt/data/generated_image.jpg"
        rgb_img.save(download_path, "JPEG", quality=95)

        return img, download_path

    except Exception as e:
        raise gr.Error(f"❌ Unexpected error: {str(e)}")

# Interface
with gr.Blocks(css=".gr-button-lg {width: 100% !important;}") as demo:
    gr.Markdown("## 🚀 CodeHubb Image Generator\nYour Creative Partner – **Design. Develop. Dominate**")

    prompt = gr.Textbox(label="Enter Prompt", placeholder="e.g. cinematic mountain road", lines=2)
    generate_btn = gr.Button("🎨 Generate Image", elem_classes="gr-button-lg")
    image_output = gr.Image(label="Generated Image", type="pil")
    download_button = gr.File(label="⬇️ Download JPG", elem_classes="gr-button-lg")

    generate_btn.click(fn=generate_image, inputs=[prompt], outputs=[image_output, download_button])

# ✅ Fix: Allow /mnt/data for download
demo.launch(allowed_paths=["/mnt/data"])