mock.py

"""Celebrity Deathmatch — mock data + CPU placeholder renderers (DEATHMATCH_MOCK=1).

Lets the whole UX run on CPU in seconds with no backend: a canned fight card,
labelled placeholder keyframes, and a mock "fight video" (animated GIF cycling
the keyframes with commentary captions burned in — proving the captions decision).
Real models swap in behind the same model_runtime signatures.
"""
from __future__ import annotations

import os
import tempfile
import textwrap

from PIL import Image, ImageDraw, ImageFont

# Canned Stage-1 output — the marquee matchup. Same shape as the backend FightCard.
MOCK_FIGHTCARD = {
    "fighter_a": {
        "name": "Britney Spears",
        "appearance": "blonde clay figure with a high ponytail, sparkly red bodysuit, "
                      "headset microphone, confident pop-diva pose",
        "persona": "Unbothered pop royalty who turns every hit into choreography",
        "signature_move": "The Toxic Spin Kick",
        "stats": {"power": 6, "speed": 9, "showmanship": 10},
    },
    "fighter_b": {
        "name": "Eminem",
        "appearance": "lanky clay figure with bleached buzzcut, white tank top, baggy jeans, "
                      "permanent scowl, fists raised",
        "persona": "Rapid-fire trash-talker who fights to a 140 BPM flow",
        "signature_move": "The Rap God Combo",
        "stats": {"power": 8, "speed": 8, "showmanship": 7},
    },
    "arena": "Classic Deathmatch ring",
    "beats": [
        {"id": 1, "title": "The stare-down / entrances",
         "action": "Britney moonwalks into the ring under a glitter cannon while Eminem "
                   "stomps in from the opposite corner, cracking his clay knuckles",
         "commentary": [
             {"speaker": "JOHNNY", "line": "Pop royalty versus the Rap God — "
                                           "and the glitter cannon is FIRING!"},
             {"speaker": "NICK", "line": "There's confetti in Slim's eyes already, Johnny."}],
         "video_prompt": "two clay fighters enter a wrestling ring, glitter falling, slow zoom"},
        {"id": 2, "title": "First blood",
         "action": "Eminem fires off a flurry of word-bullet punches; Britney back-flips "
                   "away and flicks her ponytail like a whip",
         "commentary": [
             {"speaker": "JOHNNY", "line": "Rapid-fire jabs from Eminem!"},
             {"speaker": "NICK", "line": "And the ponytail snaps back like a bullwhip. Cute."}],
         "video_prompt": "fast punches, clay fighter dodges with a backflip, motion blur"},
        {"id": 3, "title": "Signature move",
         "action": "Britney launches the Toxic Spin Kick, a neon-green tornado of legs "
                   "that sends Eminem bouncing off the ropes",
         "commentary": [
             {"speaker": "JOHNNY", "line": "THE TOXIC SPIN KICK! It's a green tornado of doom!"},
             {"speaker": "NICK", "line": "Slim Shady is ping-ponging off the ropes, folks."}],
         "video_prompt": "spinning kick, green energy tornado, opponent flung into ropes"},
        {"id": 4, "title": "Comeback / turning point",
         "action": "Eminem shakes it off and unleashes the Rap God Combo — fists moving "
                   "so fast they blur into a clay smear",
         "commentary": [
             {"speaker": "JOHNNY", "line": "Here comes the Rap God Combo — his hands are a BLUR!"},
             {"speaker": "NICK", "line": "Her tiara just got rearranged."}],
         "video_prompt": "blurred rapid punch combo, clay smear trails, dramatic angle"},
        {"id": 5, "title": "Finishing move + KO",
         "action": "Britney counters mid-combo, drops the beat AND Eminem with one final "
                   "spin, then strikes a victory pose on his flattened clay body",
         "commentary": [
             {"speaker": "JOHNNY", "line": "She drops the beat — and drops EMINEM! IT'S OVER!"},
             {"speaker": "NICK", "line": "Britney Spears wins the Deathmatch. Finally."}],
         "video_prompt": "final spin counter, opponent flattened, victory pose, confetti"},
    ],
    "winner": "A",
    "winner_reason": "Britney's showmanship and speed turned every hit into a dance "
                     "finale — Eminem never found the off-beat.",
}

# Mirror the real backend's validated shape: keep the structured two-announcer
# lines in `commentary_lines`, and expose a flat "SPEAKER: line" caption string in
# `commentary` (what the UI renders and what burns in as subtitles).
for _beat in MOCK_FIGHTCARD["beats"]:
    _lines = _beat["commentary"]
    _beat["commentary_lines"] = _lines
    _beat["commentary"] = "  ".join(f"{ln['speaker']}: {ln['line']}" for ln in _lines)

_PALETTE = [
    (210, 60, 70), (240, 160, 48), (64, 160, 200), (120, 90, 200), (60, 170, 110),
]


def _font(size: int):
    for path in (
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
    ):
        if os.path.exists(path):
            return ImageFont.truetype(path, size)
    try:
        return ImageFont.load_default(size)  # Pillow >= 10
    except TypeError:
        return ImageFont.load_default()


def _wrap(draw, text, font, max_w):
    words, lines, cur = text.split(), [], ""
    for w in words:
        trial = f"{cur} {w}".strip()
        if draw.textlength(trial, font=font) <= max_w:
            cur = trial
        else:
            if cur:
                lines.append(cur)
            cur = w
    if cur:
        lines.append(cur)
    return lines


def placeholder_keyframe(card: dict, beat: dict, size=(1024, 576)) -> Image.Image:
    """A labelled placeholder panel standing in for a FLUX render."""
    w, h = size
    idx = (beat["id"] - 1) % len(_PALETTE)
    base = _PALETTE[idx]
    img = Image.new("RGB", size, base)
    draw = ImageDraw.Draw(img)

    # Vignette-ish darkening band at the bottom for caption legibility.
    draw.rectangle([0, h - 150, w, h], fill=(0, 0, 0))

    title_f = _font(40)
    body_f = _font(26)
    small_f = _font(20)

    a = card["fighter_a"]["name"]
    b = card["fighter_b"]["name"]
    draw.text((36, 28), f"ROUND {beat['id']}", font=_font(24), fill=(255, 255, 255))
    draw.text((36, 60), f"{a}  vs  {b}", font=title_f, fill=(255, 255, 255))

    for i, line in enumerate(_wrap(draw, beat["action"], body_f, w - 72)[:4]):
        draw.text((36, 150 + i * 34), line, font=body_f, fill=(245, 245, 245))

    for i, line in enumerate(_wrap(draw, beat["commentary"], small_f, w - 72)[:3]):
        draw.text((36, h - 140 + i * 26), line, font=small_f, fill=(255, 220, 120))

    draw.text((w - 150, 28), "MOCK", font=_font(28), fill=(255, 255, 255))
    return img


def placeholder_reel(card: dict, size=(1024, 576)) -> list[Image.Image]:
    return [placeholder_keyframe(card, beat, size) for beat in card["beats"]]


def mock_fight_video(frames: list[Image.Image], seconds_per_beat: float = 1.6,
                     fps: int = 24) -> str:
    """Stitch keyframes into an MP4 — the mock stand-in for the LTX fight video.

    Commentary is already drawn onto each placeholder frame, so this mirrors the
    real Stage-3 output type (an MP4 for gr.Video) without a GPU. Real mode swaps
    in the chained, caption-burned clip from the deathmatch-video service.
    """
    import imageio.v2 as imageio
    import numpy as np

    fd, path = tempfile.mkstemp(prefix="deathmatch_mock_", suffix=".mp4")
    os.close(fd)
    reps = max(1, int(seconds_per_beat * fps))
    writer = imageio.get_writer(path, fps=fps, codec="libx264",
                                macro_block_size=8)
    try:
        for f in frames:
            arr = np.asarray(f.convert("RGB"))
            for _ in range(reps):
                writer.append_data(arr)
    finally:
        writer.close()
    return path