import json
import re

import spaces
import torch
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

MODEL_ID = "Qwen/Qwen2.5-VL-7B-Instruct"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

STYLES = {
    "Football (Premier League hype)":
        "an over-the-top English football commentator, dramatic and poetic",
    "MasterChef judge":
        "a brutally dramatic cooking-show commentator treating everything like "
        "a high-stakes elimination",
    "Nature documentary":
        "a hushed, awestruck nature narrator observing these humans like rare "
        "wildlife",
    "Boxing announcer":
        "a booming boxing announcer treating every move as championship-defining",
    "Diva Hour":
        "a fabulous, shady red-carpet diva commentator narrating like everyone "
        "is a celebrity arriving at a star-studded gala, full of glamour and sass",
}
SYSTEM_PROMPT = """You are {persona}, calling live play-by-play over a short clip.

You'll receive {n_frames} keyframes in time order, each tagged "t=<seconds>".
Some people have their NAME burned in a box above their head — use those exact
names. Refer to anyone unnamed by what you see ("the one in the red shirt").
People we already know in this clip: {roster}.

Write EXACTLY one line per keyframe — {n_frames} lines, in the same order.

Make it land:
- React to what is ACTUALLY in each frame (a pose, a movement, who's present) —
  never generic filler that could fit any clip.
- Build an arc across the lines: set the scene, let the tension rise, pay off
  the final beat. Each line should feel like it follows the last.
- Vary your openings and rhythm. Never reuse a phrase or a sentence shape twice.
- Drop names naturally, only when it hits — not in every single line.
- Stay fully in character as {persona} from first line to last.
- Save CAPS and exclamation marks for the BIG beats only; overusing them kills
  the punch (the voice reads punctuation as emotion).

Output rules (strict):
- Output ONLY a JSON array. No markdown, no text outside it.
- Schema: [{{"time": <float seconds>, "text": "<one line>"}}]
- Use the EXACT "t=" value of each keyframe as its "time".
- Keep every line under 16 words — it must be spoken before the next beat."""


# -- eager module-level load (ZeroGPU pattern) --------------------------------
print(f"[info] Loading Qwen2.5-VL-7B (first boot: download ~16GB + load) -> {DEVICE}")
_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID, torch_dtype=torch.bfloat16)
_model.to(DEVICE)
_processor = AutoProcessor.from_pretrained(MODEL_ID)
print("[info] Qwen ready.")


@spaces.GPU(duration=120)
def generate_commentary(event_frames, roster, style_key):
    """event_frames: [(timestamp_sec, PIL.Image), ...] (annotated, key events)
       roster: list of names the user assigned
       -> [{"time": float, "text": str}, ...] aligned to the event times.
    """
    from qwen_vl_utils import process_vision_info

    n = len(event_frames)
    content = []
    for i, (ts, img) in enumerate(event_frames):
        content.append({"type": "text", "text": f"Keyframe {i+1}/{n} — t={ts:.2f}s:"})
        content.append({"type": "image", "image": img})
    content.append({"type": "text",
                    "text": f"Now write the {n}-line commentary JSON."})

    persona = STYLES.get(style_key, list(STYLES.values())[0])
    roster_str = ", ".join(roster) if roster else "nobody named yet"
    messages = [
        {"role": "system",
         "content": SYSTEM_PROMPT.format(persona=persona, roster=roster_str,
                                         n_frames=n)},
        {"role": "user", "content": content},
    ]

    text = _processor.apply_chat_template(messages, tokenize=False,
                                          add_generation_prompt=True)
    image_inputs, _ = process_vision_info(messages)
    inputs = _processor(text=[text], images=image_inputs,
                        return_tensors="pt").to(_model.device)

    out = _model.generate(**inputs, max_new_tokens=500, do_sample=True,
                          temperature=0.8, top_p=0.95)
    new = out[:, inputs.input_ids.shape[1]:]
    raw = _processor.batch_decode(new, skip_special_tokens=True)[0]

    valid_times = [round(ts, 2) for ts, _ in event_frames]
    return _parse(raw, valid_times)


def _parse(raw: str, valid_times: list[float]) -> list[dict]:
    raw = re.sub(r"```(?:json)?", "", raw).strip()
    m = re.search(r"\[.*\]", raw, re.DOTALL)
    try:
        items = json.loads(m.group(0) if m else raw)
    except (json.JSONDecodeError, AttributeError):
        items = []

    script = []
    for i, it in enumerate(items):
        try:
            txt = str(it["text"]).strip()
            t = float(it.get("time", valid_times[min(i, len(valid_times) - 1)]))
        except (KeyError, TypeError, ValueError):
            continue
        t = min(valid_times, key=lambda v: abs(v - t)) if valid_times else t
        if txt:
            script.append({"time": t, "text": txt})

    if not script and valid_times:
        script = [{"time": valid_times[0],
                   "text": "WHAT A MOMENT, LADIES AND GENTLEMEN!"}]
    script.sort(key=lambda x: x["time"])
    return script