import spaces
import torch
import gradio as gr
from transformers import AutoModelForImageTextToText, AutoProcessor

MODEL_ID = "Qwen/Qwen3.8-27B"

print(f"Loading {MODEL_ID}...")
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID, torch_dtype=torch.bfloat16, device_map="cuda",
)
print("Model loaded!")

THINK_OPEN = "<think>"
THINK_CLOSE = "</think>"
IM_END = chr(60) + chr(124) + "im_end" + chr(124) + chr(62)


def _estimate_duration(prompt, system_prompt, enable_thinking, reasoning_effort,
                       temperature, max_new_tokens, image=None, video=None):
    """Rough wall-clock estimate (seconds) for one chat call. Requesting less
    than the 60s default raises queue priority and frees the GPU sooner.
    Media inputs add fixed overhead for vision encoding."""
    per_token = 0.03 if enable_thinking else 0.02
    seconds = int(max_new_tokens) * per_token
    if image:
        seconds += 5
    if video:
        seconds += 30
    return max(30, min(int(seconds) + 20, 240))


def _friendly_gpu_error(err):
    msg = (str(err) or "").lower()
    hints = ("gpu limit", "quota", "no gpu", "could not allocate",
             "gpu is busy", "too many", "concurrent")
    if any(h in msg for h in hints):
        return ("⛔ This demo's shared GPU is at capacity right now — it's not "
                "a problem with your prompt or account. Please wait a minute "
                "and retry.")
    if "out of memory" in msg or "oom" in msg:
        return ("💥 Ran out of GPU memory. Try a shorter prompt, fewer "
                "Max New Tokens, or a smaller/shorter media file, then retry.")
    return "⚠️ Generation failed. Please try again in a moment."


def _media_path(value):
    """Canvas media ports arrive as {'path': ..., 'url': ...} dicts (or None).
    Return a usable local path/URL string or None."""
    if not value:
        return None
    if isinstance(value, str):
        return value or None
    if isinstance(value, dict):
        return value.get("path") or value.get("url") or None
    return None


@spaces.GPU(duration=_estimate_duration, size="xlarge")
def _chat_gpu(prompt, system_prompt, enable_thinking, reasoning_effort,
              temperature, max_new_tokens, image=None, video=None):
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    image_src = _media_path(image)
    video_src = _media_path(video)

    user_content = []
    if image_src:
        user_content.append({"type": "image", "image": image_src})
    if video_src:
        user_content.append({"type": "video", "video": video_src})
    user_content.append({"type": "text", "text": prompt.strip()})

    messages = []
    if system_prompt and system_prompt.strip():
        messages.append({"role": "system",
                         "content": [{"type": "text", "text": system_prompt.strip()}]})
    messages.append({"role": "user", "content": user_content})

    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True,
        enable_thinking=bool(enable_thinking),
    )

    images = [image_src] if image_src else None
    videos = [video_src] if video_src else None
    inputs = processor(text=[text], images=images, videos=videos,
                       return_tensors="pt").to("cuda")

    # Recommended sampling params from the model card.
    if enable_thinking:
        gen_kwargs = dict(temperature=1.0, top_p=0.95, top_k=20)
    else:
        gen_kwargs = dict(temperature=float(temperature), top_p=0.80,
                          top_k=20, presence_penalty=1.5)
    gen_kwargs["max_new_tokens"] = int(max_new_tokens)

    out = model.generate(**inputs, do_sample=True, **gen_kwargs)
    new_tokens = out[0][inputs["input_ids"].shape[1]:]
    full = processor.decode(new_tokens, skip_special_tokens=False)

    reasoning, answer = "", full
    if THINK_OPEN in full and THINK_CLOSE in full:
        reasoning = full.split(THINK_OPEN, 1)[1].split(THINK_CLOSE, 1)[0].strip()
        answer = full.split(THINK_CLOSE, 1)[1].strip()
    answer = answer.replace(IM_END, "").strip()
    return answer, reasoning


def chat(prompt, system_prompt, enable_thinking, reasoning_effort,
         temperature, max_new_tokens, image=None, video=None):
    """Workflow-facing wrapper around the GPU worker. Catches ZeroGPU
    allocator rejections and rewords them into user-facing messages."""
    try:
        return _chat_gpu(prompt, system_prompt, enable_thinking,
                         reasoning_effort, temperature, max_new_tokens,
                         image=image, video=video)
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(_friendly_gpu_error(e)) from e


demo = gr.Workflow(graph="workflow.json", bind={"chat": chat})

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