import os
from threading import Thread

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  (must be imported before torch / transformers)

import av  # noqa: E402
import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
from transformers import (  # noqa: E402
    AutoModelForImageTextToText,
    AutoProcessor,
    BitsAndBytesConfig,
    TextIteratorStreamer,
)

MODEL_ID = "Qwen/Qwen3-VL-32B-Instruct"
HF_TOKEN = os.environ.get("HF_TOKEN", None)

TITLE = "Qwen3-VL-32B-Instruct — Vision & Video Chat"
SUBTITLE = "Ask questions about images and videos on ZeroGPU (4-bit, NF4)"
DESCRIPTION = (
    "A multimodal chat demo for Qwen3-VL-32B-Instruct running on Hugging Face "
    "ZeroGPU. Upload one or more images, or a short video, and ask questions about them."
)

SYSTEM_PROMPT = (
    "You are Qwen3-VL-32B-Instruct, a highly capable multimodal assistant. "
    "Be clear, accurate, and helpful."
)
PLACEHOLDER = (
    "Upload images or a video and ask anything about them — or just type a question."
)

MAX_NEW_TOKENS = 8192
DEFAULT_MAX_NEW_TOKENS = 2048
MAX_HISTORY_MESSAGES = 10
MAX_IMAGES_PER_MESSAGE = 3
MAX_VIDEO_FRAMES = 32

IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif"}
VIDEO_EXTS = {".mp4", ".webm", ".mov", ".mkv", ".avi", ".m4v", ".ogv"}


# ---------------------------------------------------------------------------
# Model loading (ZeroGPU: module scope, bnb 4-bit, device_map="cuda")
# ---------------------------------------------------------------------------

BNB_CONFIG = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

processor = AutoProcessor.from_pretrained(MODEL_ID, token=HF_TOKEN)
# Cap per-image pixel budget so single images stay fast on shared ZeroGPU queues.
processor.image_processor.size = {"shortest_edge": 65536, "longest_edge": 1_600_000}

model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    token=HF_TOKEN,
    quantization_config=BNB_CONFIG,
    device_map="cuda",
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
    low_cpu_mem_usage=True,
)
model.eval()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def file_kind(path: str):
    ext = os.path.splitext(path)[1].lower()
    if ext in VIDEO_EXTS:
        return "video"
    if ext in IMAGE_EXTS:
        return "image"
    return None


def sample_video_frames(path: str, max_frames: int = MAX_VIDEO_FRAMES) -> np.ndarray:
    """Decode a video into a uniformly-sampled stack of RGB frames (T, H, W, 3)."""
    container = av.open(path)
    stream = container.streams.video[0]
    frames = [frame.to_ndarray(format="rgb24") for frame in container.decode(stream)]
    container.close()
    if not frames:
        raise ValueError("Could not decode any frames from this video.")
    if len(frames) > max_frames:
        indices = np.linspace(0, len(frames) - 1, max_frames, dtype=int)
        frames = [frames[i] for i in indices]
    return np.stack(frames)


def append_files_to_content(paths, items, allow_video, note_video):
    images = 0
    for path in paths:
        if not path or not os.path.exists(path):
            continue
        kind = file_kind(path)
        if kind == "image":
            if images < MAX_IMAGES_PER_MESSAGE:
                items.append({"type": "image", "image": path})
                images += 1
        elif kind == "video":
            if allow_video:
                items.append({"type": "video", "video": path})
            elif note_video:
                items.append({"type": "text", "text": "[a video was shared in this message]"})


def build_messages(history, message, system_prompt):
    """Convert Gradio multimodal history + message into transformers messages."""
    messages = []
    if system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})

    history = history[-MAX_HISTORY_MESSAGES:]
    for i, msg in enumerate(history):
        role = msg.get("role", "user")
        content = msg.get("content", "")
        if isinstance(content, str):
            if content.strip():
                messages.append({"role": role, "content": content.strip()})
            continue
        items = []
        is_last_user = role == "user" and i == len(history) - 1
        for item in content:
            if item.get("type") == "text":
                text = item.get("text", "")
                if text.strip():
                    items.append({"type": "text", "text": text.strip()})
        paths = [item.get("file") for item in content if item.get("type") == "file"]
        append_files_to_content(paths, items, allow_video=is_last_user, note_video=True)
        if not items:
            items = [{"type": "text", "text": "(empty message)"}]
        messages.append({"role": role, "content": items})

    items = []
    paths = [f if isinstance(f, str) else f.get("path") for f in message.get("files") or []]
    append_files_to_content(paths, items, allow_video=True, note_video=False)
    text = (message.get("text") or "").strip()
    if text:
        items.append({"type": "text", "text": text})
    elif not items:
        items = [{"type": "text", "text": "What do you see here?"}]
    messages.append({"role": "user", "content": items})
    return messages


def format_output(text: str) -> str:
    text = text.replace("<|im_end|>", "")
    return text


def estimate_duration(
    message,
    history,
    system_prompt,
    temperature,
    max_new_tokens,
    top_p,
    top_k,
):
    del history, system_prompt, temperature, top_p, top_k
    base = 90 + int(max_new_tokens / 40)
    for f in message.get("files") or []:
        path = f if isinstance(f, str) else f.get("path")
        kind = file_kind(path)
        if kind == "video":
            base += 40
        elif kind == "image":
            base += 15
    return min(360, max(90, base))


# ---------------------------------------------------------------------------
# ZeroGPU chat handler
# ---------------------------------------------------------------------------

@spaces.GPU(duration=estimate_duration, size="large")
def chat(
    message: dict,
    history: list,
    system_prompt: str,
    temperature: float,
    max_new_tokens: int,
    top_p: float,
    top_k: int,
):
    """Chat with Qwen3-VL-32B-Instruct over images and videos. Streams the answer."""
    messages = build_messages(history, message, system_prompt)

    for msg in messages:
        if isinstance(msg["content"], list):
            for item in msg["content"]:
                if item.get("type") == "video":
                    item["video"] = sample_video_frames(item["video"])

    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    streamer = TextIteratorStreamer(
        processor.tokenizer,
        timeout=300.0,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    generation_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=temperature > 0,
        temperature=max(temperature, 1e-5),
        top_p=top_p,
        top_k=top_k,
        use_cache=True,
    )

    worker = Thread(target=model.generate, kwargs=generation_kwargs)
    worker.start()

    output = ""
    for chunk in streamer:
        output += chunk
        yield format_output(output)


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

CSS = """
.gradio-container { max-width: 1200px !important; margin: 0 auto !important; }
.title h1 { text-align: center; margin-bottom: 0.2rem !important; }
.subtitle p, .meta p { text-align: center; }
.meta p { font-size: 0.95rem; color: #6b7280; margin-top: 0.35rem !important; }
.duplicate-button { margin: 0 auto 14px auto !important; }
"""

chatbot = gr.Chatbot(
    height=620,
    placeholder=PLACEHOLDER,
    avatar_images=(None, "https://huggingface.co/front/assets/huggingface_logo-noborder.svg"),
)

with gr.Blocks() as demo:
    gr.Markdown(f"# {TITLE}", elem_classes="title")
    gr.Markdown(SUBTITLE, elem_classes="subtitle")
    gr.Markdown(
        f"{DESCRIPTION} Model: [{MODEL_ID}](https://huggingface.co/{MODEL_ID})",
        elem_classes="meta",
    )
    gr.DuplicateButton("Duplicate Space", elem_classes="duplicate-button")
    gr.ChatInterface(
        fn=chat,
        chatbot=chatbot,
        fill_height=True,
        multimodal=True,
        additional_inputs_accordion=gr.Accordion("⚙️ Parameters", open=False, render=False),
        additional_inputs=[
            gr.Textbox(value=SYSTEM_PROMPT, label="System prompt", lines=3, render=False),
            gr.Slider(minimum=0.0, maximum=1.2, step=0.05, value=0.7, label="Temperature", render=False),
            gr.Slider(
                minimum=256,
                maximum=MAX_NEW_TOKENS,
                step=256,
                value=DEFAULT_MAX_NEW_TOKENS,
                label="Max new tokens",
                render=False,
            ),
            gr.Slider(minimum=0.1, maximum=1.0, step=0.05, value=0.9, label="Top-p", render=False),
            gr.Slider(minimum=1, maximum=100, step=1, value=20, label="Top-k", render=False),
        ],
        examples=[
            [{"text": "What animal is on the candy?", "files": ["examples/candy.jpg"]}],
            [{"text": "Solve the geometry problem in this image.", "files": ["examples/math.jpg"]}],
            [{"text": "How many porcelain jars were discovered in the niches located in the primary chamber of the tomb?", "files": ["examples/sample_video.mp4"]}],
            ["Explain the benefits of multimodal language models."],
        ],
        cache_examples=False,
    )

if __name__ == "__main__":
    demo.launch(theme="soft", css=CSS)