import json
import os
import random
import secrets
from datetime import datetime, timezone

import gradio as gr
import spaces
import torch

# Private staging: directedbykobyperez/Qwen-Image-2.1-Create
# Bucket: directedbykobyperez/Qwen-Image-2.1-community
COMMUNITY_BUCKET = "directedbykobyperez/Qwen-Image-2.1-community"
COMMUNITY_URL = "https://huggingface.co/buckets/directedbykobyperez/Qwen-Image-2.1-community"

MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen-Image-2.1")
OUT = "/tmp/qwen21_out"
os.makedirs(OUT, exist_ok=True)

# Official 2K aspect presets from Qwen/Qwen-Image-2.1 model card.
SIZE_PRESETS = {
    "1:1 (2048x2048)": (2048, 2048),
    "16:9 (2752x1536)": (2752, 1536),
    "9:16 (1536x2752)": (1536, 2752),
    "4:3 (2400x1792)": (2400, 1792),
    "3:4 (1792x2400)": (1792, 2400),
    "3:2 (2528x1696)": (2528, 1696),
    "2:3 (1696x2528)": (1696, 2528),
    "1:1 fast (1024x1024)": (1024, 1024),
    "16:9 fast (1344x768)": (1344, 768),
    "9:16 fast (768x1344)": (768, 1344),
}

pipe = None

# Flagged words: generation proceeds normally, but flagged results are NOT
# uploaded to the public community bucket. Silent, no user-facing warning.
FLAGGED = (
    "naked", "nude", "nsfw", "porn", "pornographic", "hentai", "erotic",
    "sex", "sexual", "undress", "unclothed", "topless", "bottomless",
    "nak3d", "nudes", "bathing", "panties", "lingerie", "underwear",
    "orgasm", "nipple", "nipples", "genital", "penis", "vagina", "boobs",
    "titties", "incest", "rape",
    "child", "kid ", "kids", "minor", "teen", "teenager", "underage",
    "schoolgirl", "schoolboy", "loli", "shota", "toddler", "infant",
    "小女孩", "小男孩", "少女", "儿童", "裸体", "裸足", "脱下", "脱掉",
    "裸", "色情", "性爱", "幼女",
)


def is_flagged(*texts):
    joined = " " + " ".join(t or "" for t in texts).lower() + " "
    return any(w in joined for w in FLAGGED)


def get_pipeline():
    global pipe
    if pipe is None:
        from diffusers import QwenImage21Pipeline
        print(f"Loading {MODEL_ID} (first run downloads weights)...")
        pipe = QwenImage21Pipeline.from_pretrained(
            MODEL_ID, torch_dtype=torch.bfloat16,
        ).to("cuda" if torch.cuda.is_available() else "cpu")
        print("Pipeline loaded.")
    return pipe


def _upload_to_community(png_path, meta_dict, image_id):
    """Silently share images/<id>/<id>.png + meta.json to the community bucket. Never raises."""
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
    if not token:
        return
    try:
        from huggingface_hub import HfFileSystem
        fs = HfFileSystem(token=token)
        meta_local = os.path.join(OUT, f"{image_id}_meta.json")
        with open(meta_local, "w", encoding="utf-8") as f:
            json.dump(meta_dict, f, ensure_ascii=False, indent=2)
        base = f"buckets/{COMMUNITY_BUCKET}/images/{image_id}"
        fs.put_file(png_path, f"{base}/{image_id}.png")
        fs.put_file(meta_local, f"{base}/meta.json")
        print(f"[community] uploaded {image_id} to bucket")
    except Exception as e:
        print(f"[community] upload failed: {e}")


@spaces.GPU(duration=180)
def generate(prompt, negative_prompt, size_preset, steps, seed, edit_image):
    pipeline = get_pipeline()
    if seed is None or seed < 1:
        try:
            seed = int(seed) if seed else random.randint(1, 10**6)
        except Exception:
            seed = random.randint(1, 10**6)
    seed = int(seed)
    width, height = SIZE_PRESETS.get(size_preset, (1024, 1024))
    device = "cuda" if torch.cuda.is_available() else "cpu"
    generator = torch.Generator(device).manual_seed(seed)

    kwargs = dict(
        prompt=prompt,
        width=width, height=height,
        num_inference_steps=int(steps),
        generator=generator,
    )
    if negative_prompt and negative_prompt.strip():
        kwargs["negative_prompt"] = negative_prompt.strip()
    mode = "t2i"
    if edit_image is not None:
        from PIL import Image as PILImage
        kwargs["image"] = PILImage.open(edit_image).convert("RGB")
        mode = "edit"

    image = pipeline(**kwargs).images[0]

    image_id = secrets.token_hex(6)
    png_path = os.path.join(OUT, f"{image_id}.png")
    image.save(png_path)

    flagged = is_flagged(prompt, negative_prompt) or mode == "edit"

    words = (prompt or "").strip().replace("\n", " ").split()[:8]
    title = " ".join(words).title()[:80] if words else f"Qwen Image {image_id[:6]}"
    if not flagged:
        _upload_to_community(png_path, {
            "id": image_id,
            "title": title,
            "prompt": prompt or "",
            "negative_prompt": negative_prompt or "",
            "caption": (prompt or "").strip(),
            "width": width,
            "height": height,
            "steps": int(steps),
            "seed": seed,
            "mode": mode,
            "model": MODEL_ID,
            "image_file": f"{image_id}.png",
            "created_at": datetime.now(timezone.utc).isoformat(),
        }, image_id)

    return png_path, png_path, f"seed={seed} | {width}x{height} | {steps} steps | {mode}"


with gr.Blocks(title="Qwen Image 2.1 Create") as demo:
    gr.Markdown("# Qwen Image 2.1 Create\nText-to-image + image editing. Every generation auto-shares (png + meta.json) to the community bucket.")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", value='A neon shop sign that reads "QWEN IMAGE 2.1", rainy night, reflections on wet pavement', lines=4)
            negative = gr.Textbox(label="Negative prompt (optional)", value="", lines=2)
            edit_image = gr.Image(label="Edit image (optional — leave empty for text-to-image)", type="filepath")
            with gr.Row():
                size_preset = gr.Dropdown(label="Size", choices=list(SIZE_PRESETS.keys()), value="1:1 fast (1024x1024)")
                steps = gr.Number(label="Steps (40 = official default)", value=30, precision=0)
                seed = gr.Number(label="Seed (0 = random)", value=0, precision=0)
            btn = gr.Button("Generate image", variant="primary")
            gr.HTML("<div style='text-align:center;margin:-12px 0 -8px;font-size:0.9rem;'>Don't forget to<a style='margin-left:5px;padding:0;' href='https://huggingface.co/Qwen/Qwen-Image-2.1' target='_blank'>like the model ❤️</a></div>")
        with gr.Column():
            out_img = gr.Image(label="Result", type="filepath")
            out_file = gr.File(label="Download PNG")
            info = gr.Textbox(label="Info")
    btn.click(fn=generate, inputs=[prompt, negative, size_preset, steps, seed, edit_image], outputs=[out_img, out_file, info], queue=True)
    gr.Markdown(
        "---\nPowered by [Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1) · "
        "🎨 **Community images:** every generation is auto-shared (png + meta.json) to the "
        f"[Qwen-Image-2.1-community bucket]({COMMUNITY_URL})"
    )

demo.queue().launch()