"""
Gradio app.py-style script for Hugging Face Spaces (ZeroGPU) using CogVideoX-5B.

Model: https://huggingface.co/zai-org/CogVideoX-5b

Notes:
- CogVideoX-5B is heavy, so this script is optimized for ZeroGPU usage:
  lazy model load, queue, bounded frames/steps, and GPU-decorated inference.
- "Max length" here means model-safe max output in one generation call.
  The model card indicates 49 frames at 8 FPS (~6s), so we cap at 49.
"""

from __future__ import annotations

import os
import tempfile
from pathlib import Path

import gradio as gr
import spaces
import torch
from diffusers import CogVideoXPipeline
from diffusers.utils import export_to_video


MODEL_ID = "zai-org/CogVideoX-5b"
MAX_FRAMES = 49
MAX_STEPS = 50
DEFAULT_FPS = 8

_pipe: CogVideoXPipeline | None = None


def get_pipe() -> CogVideoXPipeline:
    global _pipe
    if _pipe is not None:
        return _pipe

    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
    _pipe = CogVideoXPipeline.from_pretrained(MODEL_ID, torch_dtype=dtype)

    if torch.cuda.is_available():
        _pipe.enable_model_cpu_offload()
        _pipe.vae.enable_tiling()
    else:
        _pipe.to("cpu")

    return _pipe


def _generate(
    prompt: str,
    negative_prompt: str,
    num_inference_steps: int,
    guidance_scale: float,
    seed: int,
    fps: int,
) -> str:
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Prompt is required.")

    pipe = get_pipe()
    steps = max(10, min(int(num_inference_steps), MAX_STEPS))
    fps = max(4, min(int(fps), DEFAULT_FPS))
    generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(int(seed))

    out = pipe(
        prompt=prompt,
        negative_prompt=(negative_prompt or "").strip() or None,
        num_videos_per_prompt=1,
        num_inference_steps=steps,
        num_frames=MAX_FRAMES,
        guidance_scale=float(guidance_scale),
        generator=generator,
    )
    frames = out.frames[0]

    tmp_dir = Path(tempfile.mkdtemp(prefix="cogvideox_"))
    out_path = tmp_dir / "output.mp4"
    export_to_video(frames, str(out_path), fps=fps)
    return str(out_path)


@spaces.GPU(duration=240)
def generate_gpu(
    prompt: str,
    negative_prompt: str,
    num_inference_steps: int,
    guidance_scale: float,
    seed: int,
    fps: int,
) -> str:
    return _generate(prompt, negative_prompt, num_inference_steps, guidance_scale, seed, fps)


with gr.Blocks(title="CogVideoX-5B ZeroGPU Text-to-Video") as demo:
    gr.Markdown(
        """
        # CogVideoX-5B Text-to-Video (ZeroGPU)
        - Model: `zai-org/CogVideoX-5b`
        - Max per-run length is capped to model-safe output (49 frames @ 8 FPS ~ 6s)
        """
    )
    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="A cinematic drone shot over neon-lit city streets at night...",
                lines=4,
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt (optional)",
                placeholder="blurry, distorted, low quality, watermark, text overlay",
                lines=2,
            )
            with gr.Row():
                steps = gr.Slider(10, MAX_STEPS, value=35, step=1, label="Inference Steps")
                guidance = gr.Slider(1.0, 9.0, value=6.0, step=0.1, label="Guidance Scale")
            with gr.Row():
                seed = gr.Number(value=42, precision=0, label="Seed")
                fps = gr.Slider(4, DEFAULT_FPS, value=DEFAULT_FPS, step=1, label="FPS")
            run_btn = gr.Button("Generate Video", variant="primary")
        with gr.Column(scale=3):
            video = gr.Video(label="Generated Video", autoplay=False)

    run_btn.click(
        fn=generate_gpu,
        inputs=[prompt, negative_prompt, steps, guidance, seed, fps],
        outputs=video,
    )

demo.queue(max_size=20, default_concurrency_limit=1).launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))