"""Paw Talk — ZeroGPU Space entry. Both MiniCPM models run in-Space on ZeroGPU
(no external GPU). The heavy work is wrapped in a single @spaces.GPU call per request;
frame extraction, subtitle render and ffmpeg muxing stay on CPU."""
import io
import sys
import tempfile
import uuid
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent / "src"))

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

from pawscar.compose import compose
from pawscar.frames import downscale, extract_frames, probe_duration
from pawscar.personas import PERSONAS, get_persona
from pawscar.scripter import (
    build_describe_prompt,
    build_rewrite_prompt,
    build_voice_desc,
    build_voicehint_prompt,
    clean_script,
    pitch_factor,
    target_chars,
)

VISION_MODEL = "openbmb/MiniCPM-V-4.6"
VOICE_MODEL = "openbmb/VoxCPM2"
MAX_SECONDS = 30.0

# ── Load both models at module level. On ZeroGPU, .cuda() at import works via the
# CUDA emulation mode; real CUDA is attached only inside @spaces.GPU functions. ──
_processor = AutoProcessor.from_pretrained(VISION_MODEL)
_model = AutoModelForImageTextToText.from_pretrained(
    VISION_MODEL, torch_dtype=torch.float16).eval().cuda()
_voice = VoxCPM.from_pretrained(VOICE_MODEL, load_denoiser=False, optimize=False)
_SR = int(_voice.tts_model.sample_rate)


def _caption(imgs, prompt: str) -> str:
    content = [{"type": "image", "image": im} for im in imgs]
    content.append({"type": "text", "text": prompt})
    inputs = _processor.apply_chat_template(
        [{"role": "user", "content": content}], tokenize=True, add_generation_prompt=True,
        return_dict=True, return_tensors="pt").to(_model.device)
    for k, v in list(inputs.items()):
        if torch.is_tensor(v) and torch.is_floating_point(v):
            inputs[k] = v.to(torch.float16)
    with torch.no_grad():
        gen = _model.generate(**inputs, do_sample=True, temperature=0.7, max_new_tokens=320)
    out = _processor.batch_decode(
        gen[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True,
        clean_up_tokenization_spaces=False)
    return (out[0] if out else "").strip()


def _speak(text: str, voice_desc: str, steps: int = 16) -> bytes:
    import numpy as np
    import soundfile as sf

    wav = _voice.generate(text=f"({voice_desc}){text}", cfg_value=2.0, inference_timesteps=steps)
    arr = np.asarray(wav, dtype="float32").reshape(-1)
    buf = io.BytesIO()
    sf.write(buf, arr, _SR, format="WAV")
    return buf.getvalue()


@spaces.GPU(duration=120)
def _generate(frames, persona_key, dur):
    """All GPU work for one request: describe → voice timbre → inner monologue → speech."""
    persona = get_persona(persona_key)
    desc = _caption(frames, build_describe_prompt(dur))
    timbre = _caption([], build_voicehint_prompt(desc))
    voice_desc = build_voice_desc(timbre, persona)
    raw = _caption([], build_rewrite_prompt(desc, persona, dur))
    script = clean_script(raw, persona, max_chars=int(target_chars(dur) * 1.8))
    wav = _speak(script, voice_desc)
    return script, voice_desc, wav


def run(video, persona_key):
    if not video:
        raise gr.Error("Please upload a pet video first 🐾")
    try:
        dur = min(float(probe_duration(video)), MAX_SECONDS)
    except Exception:
        dur = 12.0
    frames = [downscale(f) for f in extract_frames(video)]
    script, voice_desc, wav = _generate(frames, persona_key, dur)
    out = str(Path(tempfile.gettempdir()) / f"pawscar_{uuid.uuid4().hex[:8]}.mp4")
    return compose(video, wav, out, pitch=pitch_factor(voice_desc), subtitle=script)


CHOICES = [(p.name, key) for key, p in PERSONAS.items()]

with gr.Blocks(title="Paw Talk · pet voiceover") as demo:
    gr.Markdown(
        "# 🐾 Paw Talk\n"
        "Upload a pet video and hear its inner monologue. (Clips over 30s use the first 30s.)"
    )
    with gr.Row():
        with gr.Column():
            vid = gr.Video(label="Pet video (≤30s; long clips use the first 30s)", sources=["upload"])
            style = gr.Radio(CHOICES, value="funny", label="Voice style")
            btn = gr.Button("🎬 Voice it", variant="primary")
        with gr.Column():
            out_vid = gr.Video(label="Voiced video")
    btn.click(run, [vid, style], out_vid)

demo.queue(default_concurrency_limit=4)

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