# app.py
# VisPrompt Studio — Image/Text → Long, Detailed Prompts
# - Runs on Hugging Face Spaces (free CPU)
# - Stronger scene descriptions (esp. people/adult) using Qwen2-VL-2B-Instruct
# - BLIP base as a fast fallback
# - Model selector is non-typable (Radio) + server-side guard + automatic fallback
# - Uses Textboxes (not Code) for copy-friendly outputs (Gradio 5.x safe)

import os
import math
import random
from typing import Optional

import gradio as gr
from PIL import Image, ExifTags
import torch
from transformers import (
    AutoProcessor,
    AutoModelForVision2Seq,
    BlipProcessor,
    BlipForConditionalGeneration,
)

# -------------------- Models --------------------
MODEL_QWEN2_VL = "Qwen/Qwen2-VL-2B-Instruct"
MODEL_BLIP = "Salesforce/blip-image-captioning-base"
VALID_MODELS = ("Qwen2-VL-2B-Instruct", "BLIP base")

# Lazy globals
_q_proc = _q_model = None
_b_proc = _b_model = None


def device() -> str:
    """
    Return runtime device. Keep 'cpu' for free Spaces.
    If you later use ZeroGPU, you can switch to:
        return "cuda" if torch.cuda.is_available() else "cpu"
    """
    return "cpu"


# -------------------- Helpers --------------------
def safe_open_image(inp) -> Optional[Image.Image]:
    if inp is None:
        return None
    if isinstance(inp, Image.Image):
        return inp.convert("RGB")
    return Image.open(inp).convert("RGB")


def exif_info(img: Image.Image) -> dict:
    out = {}
    try:
        raw = img.getexif()
        if not raw:
            return out
        tag_map = {ExifTags.TAGS.get(k, k): v for k, v in raw.items()}

        fl = tag_map.get("FocalLength")
        if isinstance(fl, tuple) and len(fl) == 2 and fl[1] != 0:
            out["focal_length_mm"] = round(fl[0] / fl[1], 1)
        elif isinstance(fl, (int, float)):
            out["focal_length_mm"] = float(fl)

        if "Model" in tag_map:
            out["camera_model"] = str(tag_map["Model"]).strip()
        if "Make" in tag_map:
            out["camera_make"] = str(tag_map["Make"]).strip()
    except Exception:
        pass
    return out


# Lexicons for enrichment
CAMERA_ANGLES = ["eye-level", "low-angle", "high-angle", "aerial/top-down", "worm's-eye", "Dutch tilt"]
SHOT_SIZES = ["extreme close-up", "close-up", "medium shot", "wide shot", "extreme wide shot"]
LENSES = ["14mm ultra-wide", "24mm wide-angle", "35mm street", "50mm standard", "85mm portrait", "135mm telephoto", "200mm telephoto"]
LIGHTING = [
    "soft diffused light", "golden hour rim light", "overcast softbox feel", "hard noon sun with crisp shadows",
    "blue hour ambient", "tungsten practicals", "neon key with colored rim", "studio softbox left + kicker right",
    "volumetric god rays", "moody chiaroscuro"
]
COMPOSITION = [
    "rule of thirds", "leading lines", "central composition", "frame within a frame", "negative space",
    "foreground bokeh", "shallow depth of field", "symmetry", "diagonal dynamics", "layered depth"
]
POST = [
    "high micro-contrast", "subtle film grain", "cinematic color grade", "Kodak Portra palette",
    "teal–orange split tone", "matte shadows", "HDR tonality restrained", "8k detail", "photorealistic finish"
]
STYLES = [
    "documentary realism", "editorial fashion", "street photography", "nature landscape", "architectural minimalism",
    "cyberpunk neon", "sci-fi cinematic", "fantasy epic", "moody noir", "retro 35mm film"
]
VIDEO_MOTION = [
    "slow dolly-in", "gentle push-in", "handheld micro-jitters", "locked-off tripod", "gimbal glide",
    "orbit right-to-left", "crane rise", "parallax truck shot", "rack focus mid-shot"
]
VIDEO_META = ["24 fps", "30 fps", "48 fps", "cinematic motion blur", "shutter 1/48", "global illumination feel"]
VIDEO_TRANSITIONS = ["straight cut", "cinematic fade-in/out", "cross dissolve", "match cut"]
NEGATIVE = ["no distortion", "no chromatic aberration", "no blown highlights", "no excessive sharpening", "no artifacts", "no watermark", "no text overlay"]


def smart_choices(caption: str, w: int, h: int, exif: dict, seed: int):
    rnd = random.Random(seed)
    aspect = (w / h) if h > 0 else 1.0

    shot = "wide shot" if aspect >= 1.5 else "medium shot"
    t = caption.lower()
    if any(k in t for k in ["portrait", "selfie", "face", "headshot"]):
        shot = "close-up"

    if any(k in t for k in ["aerial", "drone", "top-down", "overhead", "bird's-eye", "bird’s-eye"]):
        angle = "aerial/top-down"
    elif any(k in t for k in ["towering", "imposing", "hero", "statue"]):
        angle = "low-angle"
    else:
        angle = rnd.choice(CAMERA_ANGLES[:3])  # eye/high/low

    lens = f'{exif["focal_length_mm"]}mm' if "focal_length_mm" in exif else rnd.choice(LENSES)
    light = ", ".join(rnd.sample(LIGHTING, k=2))
    comp = ", ".join(rnd.sample(COMPOSITION, k=3))
    style = rnd.choice(STYLES)
    post = ", ".join(rnd.sample(POST, k=3))
    neg = ", ".join(NEGATIVE[:3] + rnd.sample(NEGATIVE[3:], k=2))
    return angle, shot, lens, light, comp, style, post, neg, aspect


def build_image_prompt(desc, angle, shot, lens, light, comp, style, post, aspect):
    ar = "16:9" if aspect >= 1.6 else ("4:3" if aspect >= 1.3 else "1:1")
    return (
        f"{desc}\n"
        f"Camera: {angle}; {shot}; lens {lens}.\n"
        f"Lighting: {light}; realistic reflections and shadow behavior.\n"
        f"Composition: {comp}; clean subject separation; coherent scale and perspective.\n"
        f"Style: {style}; grounded materials and textures.\n"
        f"Post: {post}; minimal artifacts; print-ready clarity.\n"
        f"Aspect ratio hint: {ar}."
    )


def build_video_prompt(img_prompt: str, motion: str, meta: list, transition: str, seconds: int):
    meta_str = ", ".join(meta)
    return (
        f"{img_prompt}\n"
        f"Video rendering: {seconds}s; camera motion: {motion}; {meta_str}; transitions: {transition}.\n"
        f"Maintain temporal consistency, realistic parallax, and stable lighting throughout."
    )


# -------------------- Caption backends --------------------
@torch.inference_mode()
def _ensure_qwen():
    global _q_proc, _q_model
    if _q_proc is None or _q_model is None:
        _q_proc = AutoProcessor.from_pretrained(MODEL_QWEN2_VL, trust_remote_code=True)
        _q_model = AutoModelForVision2Seq.from_pretrained(
            MODEL_QWEN2_VL, torch_dtype=torch.float32, trust_remote_code=True
        )
        _q_model.to(device()).eval()


@torch.inference_mode()
def caption_qwen(img: Image.Image, max_new_tokens: int = 220, temperature: float = 0.3, top_p: float = 0.9) -> str:
    _ensure_qwen()
    # Force long, explicit visual descriptions
    inst = (
        "Describe this image in exhaustive visual detail. Include: subjects, appearance, clothing, pose, "
        "body details, accessories, objects, background, textures, colors, time of day, mood, context. "
        "Be concrete and specific. Minimum ~80 words."
    )
    inputs = _q_proc(text=inst, images=img, return_tensors="pt")
    out = _q_model.generate(
        **inputs,
        max_new_tokens=int(max_new_tokens),
        do_sample=True,
        temperature=float(temperature),
        top_p=float(top_p),
        repetition_penalty=1.06,
    )
    return _q_proc.batch_decode(out, skip_special_tokens=True)[0].strip()


@torch.inference_mode()
def caption_qwen_retry_if_short(img: Image.Image, min_words: int, **gen_kwargs) -> str:
    """Retry once with a stricter structured instruction if too short."""
    first = caption_qwen(img, **gen_kwargs)
    if len(first.split()) >= min_words:
        return first

    _ensure_qwen()
    retry_inst = (
        "Rewrite a much longer, structured description of the image. Sections: "
        "1) Subjects & actions; 2) Appearance & clothing; 3) Scene & background; "
        "4) Lighting & colors; 5) Context & mood. Use 120–180 words. Be explicit and visual."
    )
    inputs = _q_proc(text=retry_inst, images=img, return_tensors="pt")
    out = _q_model.generate(
        **inputs,
        max_new_tokens=max(220, gen_kwargs.get("max_new_tokens", 220)),
        do_sample=True,
        temperature=gen_kwargs.get("temperature", 0.3),
        top_p=gen_kwargs.get("top_p", 0.9),
        repetition_penalty=1.08,
    )
    second = _q_proc.batch_decode(out, skip_special_tokens=True)[0].strip()
    return second if len(second.split()) >= min_words else first


@torch.inference_mode()
def _ensure_blip():
    global _b_proc, _b_model
    if _b_proc is None or _b_model is None:
        _b_proc = BlipProcessor.from_pretrained(MODEL_BLIP)
        _b_model = BlipForConditionalGeneration.from_pretrained(MODEL_BLIP)
        _b_model.to(device()).eval()


@torch.inference_mode()
def caption_blip(img: Image.Image, beam_search: bool = True, max_len: int = 140) -> str:
    _ensure_blip()
    inputs = _b_proc(images=img, return_tensors="pt")
    max_len = int(max(32, min(200, max_len)))
    kwargs = dict(max_length=max_len, num_beams=4) if beam_search else dict(max_length=max_len)
    out = _b_model.generate(**inputs, **kwargs)
    return _b_proc.decode(out[0], skip_special_tokens=True).strip()


def caption_image(model_choice: str, img: Image.Image, min_words: int, max_new_tokens: int) -> str:
    # Guard invalid values coming from the UI
    if model_choice not in VALID_MODELS:
        model_choice = "Qwen2-VL-2B-Instruct"

    # Try requested model, fall back to BLIP on any exception so UX doesn't hard-crash
    try:
        if model_choice == "Qwen2-VL-2B-Instruct":
            return caption_qwen_retry_if_short(
                img,
                min_words=min_words,
                max_new_tokens=max_new_tokens,
                temperature=0.35,
                top_p=0.92,
            )
        else:
            # BLIP tends to be briefer; push max_len up
            return caption_blip(img, beam_search=True, max_len=min(200, max_new_tokens))
    except Exception:
        # Fallback
        try:
            return caption_blip(img, beam_search=True, max_len=min(200, max_new_tokens))
        except Exception as e:
            # Last-resort: return readable error text in the caption box
            return f"[Error] Captioning failed: {type(e).__name__}: {e}"


# -------------------- Main generate --------------------
def generate(
    image: Optional[Image.Image],
    text_scene: str,
    mode: str,
    min_words: int,
    max_new_tokens: int,
    seconds: int,
    seed: int,
    model_choice: str,
):
    rnd = random.Random(int(seed))

    if image is not None:
        img = safe_open_image(image)
        caption = caption_image(model_choice, img, min_words=min_words, max_new_tokens=max_new_tokens)
        w, h = img.size
        exif = exif_info(img)
    else:
        # Text-only path: use user text as the scene description directly
        if not text_scene or not text_scene.strip():
            # Return empty strings so Gradio boxes clear but no crash
            return "", "", ""
        caption = text_scene.strip()
        w, h, exif = 1920, 1080, {}

    angle, shot, lens, light, comp, style, post, neg, aspect = smart_choices(caption, w, h, exif, seed)
    image_prompt = build_image_prompt(caption, angle, shot, lens, light, comp, style, post, aspect)
    negative_line = f"\nNegative prompt hints: {neg}"

    video_prompt = build_video_prompt(
        image_prompt,
        motion=rnd.choice(VIDEO_MOTION),
        meta=rnd.sample(VIDEO_META, k=2),
        transition=rnd.choice(VIDEO_TRANSITIONS),
        seconds=int(max(2, min(20, seconds))),
    )

    if mode == "Image prompt":
        return caption, image_prompt + negative_line, ""
    elif mode == "Video prompt":
        return caption, "", video_prompt
    else:
        return caption, image_prompt + negative_line, video_prompt


# -------------------- UI --------------------
TITLE = "VisPrompt Studio — Long-Form Scene Descriptions"
DESC = """
Upload an image **or** describe a scene.  
This build forces **long, explicit scene descriptions** first, then adds detailed camera/lighting/comp info for image and video prompts.

**Models**
- Qwen2-VL-2B-Instruct (default): richer descriptions on people/adult photos (slower on CPU).
- BLIP base: faster on CPU; simpler captions.
"""

with gr.Blocks(title=TITLE) as demo:
    gr.Markdown(f"# {TITLE}\n{DESC}")

    with gr.Row():
        with gr.Column(scale=1):
            model_choice = gr.Radio(
                choices=["Qwen2-VL-2B-Instruct", "BLIP base"],
                value="Qwen2-VL-2B-Instruct",
                label="Captioning model (drives the base scene description)"
            )

            with gr.Tab("Image input"):
                image = gr.Image(type="pil", label="Upload an image", height=360)
            with gr.Tab("Text input"):
                text_scene = gr.Textbox(lines=6, placeholder="Describe the scene...", label="Scene description")

            mode = gr.Radio(choices=["Image prompt", "Video prompt", "Both"], value="Both", label="Output type")

            with gr.Accordion("Description depth", open=True):
                min_words = gr.Slider(40, 200, value=90, step=10, label="Minimum words for scene description (Qwen)")
                max_new_tokens = gr.Slider(96, 320, value=220, step=4, label="Max tokens for caption generation")

            with gr.Accordion("Video settings", open=False):
                seconds = gr.Slider(3, 12, value=6, step=1, label="Video duration for prompt (seconds)")

            seed = gr.Slider(0, 9999, value=1234, step=1, label="Style seed")
            run = gr.Button("Generate prompts", variant="primary")

        with gr.Column(scale=1):
            caption_box = gr.Textbox(
                label="Long scene description (from image or your text)",
                interactive=False,
                lines=10,
                show_copy_button=True,
            )
            img_prompt_box = gr.Textbox(
                label="Image Prompt (copy to image generators)",
                interactive=False,
                lines=14,
                show_copy_button=True,
            )
            vid_prompt_box = gr.Textbox(
                label="Video Prompt (copy to video generators)",
                interactive=False,
                lines=14,
                show_copy_button=True,
            )

    run.click(
        fn=generate,
        inputs=[image, text_scene, mode, min_words, max_new_tokens, seconds, seed, model_choice],
        outputs=[caption_box, img_prompt_box, vid_prompt_box],
    )

if __name__ == "__main__":
    demo.launch()