app.py

import gradio as gr
import os
import zipfile
import requests
from PIL import Image
from io import BytesIO
from datetime import datetime

# ---------- CONFIG ----------
BYTEZ_API_KEY = os.getenv("BYTEZ_API_KEY")  # put your real key here
BYTEZ_ENDPOINT = "https://api.bytez.com/models/v2/google/imagen-4.0-fast-generate-001"

BASE_DIR = os.path.abspath(os.getcwd())
HISTORY_ROOT = os.path.join(BASE_DIR, "user_image_history")
os.makedirs(HISTORY_ROOT, exist_ok=True)

# ---------- HELPERS ----------

def get_user_dir(username: str) -> str:
    """Return (and create) per-user history directory."""
    if not username:
        username = "_anonymous"
    safe_username = "".join(c for c in username if c.isalnum() or c in ("_", "-")) or "_anonymous"
    user_dir = os.path.join(HISTORY_ROOT, safe_username)
    os.makedirs(user_dir, exist_ok=True)
    return user_dir

def get_history(username: str):
    """Return sorted list of image paths for this user."""
    user_dir = get_user_dir(username)
    history = sorted([
        os.path.join(user_dir, f)
        for f in os.listdir(user_dir)
        if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
    ])
    # First output for gallery, second output unused but kept to match callbacks
    return history, history

def generate_image_txt2img(username: str, prompt: str):
    """Call Bytez Imagen API with text prompt, download image, store it."""
    user_dir = get_user_dir(username)

    if not prompt.strip():
        # No prompt, just return history
        history_txt, history_img = get_history(username)
        return None, history_txt, None

    headers = {
        "Authorization": BYTEZ_API_KEY,
        "Content-Type": "application/json",
    }
    payload = {
        "text": prompt
    }

    try:
        resp = requests.post(BYTEZ_ENDPOINT, json=payload, headers=headers)
        resp.raise_for_status()
        data = resp.json()
    except Exception as e:
        history_txt, history_img = get_history(username)
        # Show error in label, no image
        return gr.update(label=f"Error calling API: {e}", value=None), history_txt, None

    # Bytez sample response structure (from your Postman test):
    # {
    #   "error": null,
    #   "output": "https://cdn.bytez.com/model/output/google/imagen-4.0-fast-generate-001/....png",
    #   "provider": { "generatedImages": [ { "image": { "imageBytes": "..." } } ] }
    # }

    img = None

    # Prefer using the public URL in "output"
    try:
        img_url = data.get("output")
        if img_url:
            img_resp = requests.get(img_url)
            img_resp.raise_for_status()
            img = Image.open(BytesIO(img_resp.content)).convert("RGB")
        else:
            raise ValueError("No 'output' URL in response.")
    except Exception as e:
        history_txt, history_img = get_history(username)
        return gr.update(label=f"Error downloading image: {e}", value=None), history_txt, None

    # Save image to user's history
    if img:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        filename = os.path.join(user_dir, f"{timestamp}.png")
        img.save(filename)

    history_txt, history_img = get_history(username)
    return img, history_txt, None

def select_image(image_path):
    """Return selected image path for popup display."""
    return image_path

def delete_selected_image(username: str, selected_image_path: str):
    """Delete the selected image from the user's directory."""
    if selected_image_path:
        user_dir = get_user_dir(username)
        # Only delete if path is inside this user's folder
        try:
            if os.path.commonpath(
                [os.path.abspath(selected_image_path), user_dir]
            ) == os.path.abspath(user_dir):
                if os.path.exists(selected_image_path):
                    os.remove(selected_image_path)
        except Exception:
            pass

    history_txt, history_img = get_history(username)
    return None, history_txt, None

def download_history(username: str):
    """Create a ZIP of the user's history folder and return its path."""
    user_dir = get_user_dir(username)
    zip_path = os.path.join(user_dir, "image_history.zip")

    with zipfile.ZipFile(zip_path, "w") as zipf:
        for root, dirs, files in os.walk(user_dir):
            for file in files:
                if file == "image_history.zip":
                    continue
                full_path = os.path.join(root, file)
                arcname = os.path.relpath(full_path, user_dir)
                zipf.write(full_path, arcname=arcname)

    return zip_path

# ---------- GRADIO UI ----------

with gr.Blocks() as demo:
    gr.Markdown("# Bytez Imagen 4.0 Fast – Txt2Img (per‑user history)")

    with gr.Row():
        with gr.Column(scale=2):
            username = gr.Textbox(
                label="Username (for personal image history)",
                placeholder="Enter a username, e.g. alice",
                value=""
            )
            prompt_txt2img = gr.Textbox(
                label="Prompt",
                placeholder="Describe the image...",
                lines=5
            )
            generate_btn_txt2img = gr.Button("Generate Image")
            delete_btn_txt2img = gr.Button("Delete Selected Image")
            download_btn_txt2img = gr.Button("Download My History ZIP")
        with gr.Column(scale=3):
            output_image_txt2img = gr.Image(
                label="Generated Image",
                interactive=False,
                type="pil"
            )

    with gr.Accordion("My History", open=False):
        gallery_txt2img = gr.Gallery(
            label="Image History",
            show_label=False
        )
        popup_image_txt2img = gr.Image(
            label="Full Image View",
            visible=False,
            type="filepath"
        )

    # Reload history when username changes
    username.change(
        fn=get_history,
        inputs=[username],
        outputs=[gallery_txt2img, gallery_txt2img]
    )

    # Load initial history for default username
    demo.load(
        fn=get_history,
        inputs=[username],
        outputs=[gallery_txt2img, gallery_txt2img]
    )

    # Generate image
    generate_btn_txt2img.click(
        fn=generate_image_txt2img,
        inputs=[username, prompt_txt2img],
        outputs=[output_image_txt2img, gallery_txt2img, popup_image_txt2img]
    )

    # Select image from gallery
    gallery_txt2img.select(
        fn=select_image,
        inputs=gallery_txt2img,
        outputs=popup_image_txt2img
    )

    # Delete selected image
    delete_btn_txt2img.click(
        fn=delete_selected_image,
        inputs=[username, popup_image_txt2img],
        outputs=[output_image_txt2img, gallery_txt2img, popup_image_txt2img]
    )

    # Download ZIP
    download_btn_txt2img.click(
        fn=download_history,
        inputs=[username],
        outputs=[gr.File(label="Download ZIP")]
    )

demo.launch()