import gc
import json
import os
import tempfile
from pathlib import Path
from threading import Lock

import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import AutoencoderKLWan
from diffusers.utils import export_to_video
from lingbot_video import (
    FlowUniPCMultistepScheduler,
    LingBotVideoPipeline,
    LingBotVideoTransformer3DModel,
)
from lingbot_video.pipeline_lingbot_video import DEFAULT_NEGATIVE_PROMPT
from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor


MODEL_ID = "robbyant/lingbot-video-dense-1.3b"
MAX_SEED = 2**32 - 1

RESOLUTIONS = {
    "Landscape 16:9 (832 x 480)": (480, 832),
    "Portrait 9:16 (480 x 832)": (832, 480),
    "Landscape 4:3 (640 x 480)": (480, 640),
    "Portrait 3:4 (480 x 640)": (640, 480),
    "Square 1:1 (480 x 480)": (480, 480),
}

CAMERA_DESCRIPTIONS = {
    "Static": "The camera remains stable and fixed while the subject moves naturally.",
    "Slow push in": "The camera slowly pushes toward the main subject with smooth cinematic motion.",
    "Slow pull out": "The camera slowly pulls away to reveal more of the environment.",
    "Pan left": "The camera performs a smooth pan toward the left while keeping the subject in view.",
    "Pan right": "The camera performs a smooth pan toward the right while keeping the subject in view.",
    "Handheld": "The camera uses subtle, natural handheld motion without abrupt shaking.",
    "Tracking": "The camera smoothly tracks the main subject through the scene.",
}


pipe = None
pipe_lock = Lock()


def get_pipe():
    global pipe
    if pipe is not None:
        return pipe

    with pipe_lock:
        if pipe is None:
            print(f"Loading {MODEL_ID}...", flush=True)
            common = {
                "local_files_only": True,
                "low_cpu_mem_usage": True,
            }

            print("Loading LingBot transformer...", flush=True)
            transformer = LingBotVideoTransformer3DModel.from_pretrained(
                MODEL_ID,
                subfolder="transformer",
                torch_dtype=torch.bfloat16,
                **common,
            ).to("cuda")

            print("Loading Qwen3-VL text encoder...", flush=True)
            text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
                MODEL_ID,
                subfolder="text_encoder",
                torch_dtype=torch.bfloat16,
                **common,
            ).to("cuda")

            print("Loading Wan VAE...", flush=True)
            vae = AutoencoderKLWan.from_pretrained(
                MODEL_ID,
                subfolder="vae",
                torch_dtype=torch.float32,
                **common,
            ).to("cuda")

            processor = Qwen3VLProcessor.from_pretrained(
                MODEL_ID,
                subfolder="processor",
                local_files_only=True,
            )
            scheduler = FlowUniPCMultistepScheduler.from_pretrained(
                MODEL_ID,
                subfolder="scheduler",
                local_files_only=True,
            )

            pipe = LingBotVideoPipeline(
                transformer=transformer,
                vae=vae,
                text_encoder=text_encoder,
                processor=processor,
                scheduler=scheduler,
            )
            pipe.set_progress_bar_config(disable=False)
            print("LingBot-Video pipeline loaded.", flush=True)
    return pipe


def build_structured_caption(prompt, camera_motion, duration):
    return {
        "comprehensive_description": {
            "scene_content_description": prompt.strip(),
            "camera_movement_description": CAMERA_DESCRIPTIONS[camera_motion],
        },
        "camera_info": {
            "color": "Natural",
            "frame_size": "Medium",
            "shot_type_angle": "Eye level",
            "lens_size": "Medium",
            "composition": "Center",
            "lighting": "Soft light",
            "lighting_type": "Natural light",
        },
        "world_knowledge": [],
        "prominent_elements": [],
        "target_duration_seconds": int(duration),
    }


def resolve_caption(prompt, camera_motion, duration, structured_json):
    if structured_json and structured_json.strip():
        try:
            parsed = json.loads(structured_json)
        except json.JSONDecodeError as exc:
            raise gr.Error(f"Structured caption JSON is invalid: {exc}") from exc

        caption = parsed.get("caption", parsed) if isinstance(parsed, dict) else parsed
        if not isinstance(caption, dict):
            raise gr.Error("Structured caption JSON must be an object or contain a caption object.")
    else:
        if not prompt or not prompt.strip():
            raise gr.Error("Please enter a prompt.")
        caption = build_structured_caption(prompt, camera_motion, duration)

    return caption, json.dumps(caption, ensure_ascii=False, indent=2)


def valid_frame_count(duration, fps):
    requested = max(1, int(round(float(duration) * int(fps))))
    # LingBot-Video requires one frame or 4n+1 frames.
    return max(1, ((requested - 1 + 3) // 4) * 4 + 1)


@spaces.GPU(size="large", duration=300)
def generate_video(
    prompt,
    negative_prompt,
    structured_json,
    aspect_ratio,
    camera_motion,
    duration,
    fps,
    num_inference_steps,
    guidance_scale,
    seed,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True),
):
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is not available inside the ZeroGPU worker.")

    progress(0.02, desc="Preparing LingBot-Video")
    model = get_pipe()

    if randomize_seed:
        seed = torch.randint(0, MAX_SEED, (1,)).item()
    seed = int(seed) % MAX_SEED

    height, width = RESOLUTIONS[aspect_ratio]
    fps = int(fps)
    num_frames = valid_frame_count(duration, fps)
    caption, caption_json = resolve_caption(
        prompt, camera_motion, duration, structured_json
    )
    caption_text = json.dumps(caption, ensure_ascii=False, separators=(",", ":"))
    negative_text = (
        negative_prompt.strip()
        if isinstance(negative_prompt, str) and negative_prompt.strip()
        else DEFAULT_NEGATIVE_PROMPT
    )

    generator = torch.Generator(device="cuda").manual_seed(seed)
    progress(0.08, desc=f"Generating {num_frames} frames")

    with torch.inference_mode():
        result = model(
            prompt=caption_text,
            negative_prompt=negative_text,
            height=height,
            width=width,
            num_frames=num_frames,
            num_inference_steps=int(num_inference_steps),
            guidance_scale=float(guidance_scale),
            shift=3.0,
            generator=generator,
            output_type="np",
            batch_cfg=False,
        )

    frames = np.asarray(result.frames[0])
    output_dir = Path(tempfile.gettempdir()) / "lingbot_video_outputs"
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / f"lingbot_{seed}_{height}x{width}_{num_frames}f.mp4"
    export_to_video(frames, str(output_path), fps=fps)

    del result, frames
    gc.collect()
    torch.cuda.empty_cache()

    status = (
        f"Seed: {seed} | Resolution: {width} x {height} | "
        f"Frames: {num_frames} | FPS: {fps} | Steps: {int(num_inference_steps)}"
    )
    progress(1.0, desc="Video ready")
    return str(output_path), status, caption_json


EXAMPLES = [
    [
        "A humanoid robot walks steadily along a gravel path in a formal garden, realistic materials, cinematic daylight",
        "Tracking",
    ],
    [
        "A glass bottle falls from a wooden table and shatters naturally on a concrete floor, realistic physics, high-speed detail",
        "Static",
    ],
    [
        "A young woman in a cream cardigan turns toward the camera and adjusts her collar in a bright modern apartment",
        "Slow push in",
    ],
]


with gr.Blocks(title="LingBot-Video Dense 1.3B") as demo:
    gr.Markdown(
        """
        <div class="hero">
          <h1>🎬 LingBot-Video Dense 1.3B</h1>
          <p>Generate short physical-world videos with
          <a href="https://huggingface.co/robbyant/lingbot-video-dense-1.3b">Robbyant LingBot-Video</a>
          on Hugging Face ZeroGPU.</p>
          <p class="hint">The app turns a normal prompt into a compact structured caption. Start with 2 seconds / 16 FPS for a faster first run.</p>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the scene, subject, actions and physical motion...",
                lines=5,
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt (optional)",
                placeholder="Leave empty to use LingBot-Video's official default negative prompt.",
                lines=2,
            )

            with gr.Accordion("Advanced structured caption", open=False):
                structured_json = gr.Code(
                    label="Structured Caption JSON (optional)",
                    language="json",
                    value="",
                    lines=10,
                )

            with gr.Row():
                aspect_ratio = gr.Dropdown(
                    choices=list(RESOLUTIONS),
                    value="Landscape 16:9 (832 x 480)",
                    label="Aspect Ratio",
                )
                camera_motion = gr.Dropdown(
                    choices=list(CAMERA_DESCRIPTIONS),
                    value="Static",
                    label="Camera Motion",
                )

            with gr.Row():
                duration = gr.Slider(
                    minimum=1,
                    maximum=5,
                    value=2,
                    step=1,
                    label="Duration (seconds)",
                )
                fps = gr.Dropdown(
                    choices=[16, 24],
                    value=16,
                    label="FPS",
                )

            with gr.Row():
                num_inference_steps = gr.Slider(
                    minimum=20,
                    maximum=40,
                    value=30,
                    step=1,
                    label="Inference Steps",
                )
                guidance_scale = gr.Slider(
                    minimum=1.0,
                    maximum=5.0,
                    value=3.0,
                    step=0.1,
                    label="CFG Guidance Scale",
                )

            with gr.Row():
                seed = gr.Number(label="Seed", value=42, precision=0)
                randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)

            generate_btn = gr.Button(
                "🚀 Generate Video",
                variant="primary",
                size="lg",
                elem_classes="generate-btn",
            )

        with gr.Column(scale=1):
            output_video = gr.Video(
                label="Generated Video",
                autoplay=True,
                loop=True,
                include_audio=False,
            )
            generation_info = gr.Textbox(label="Generation Info", interactive=False)
            caption_preview = gr.Code(
                label="Structured Caption Used",
                language="json",
                interactive=False,
                lines=12,
            )

    gr.Markdown("### 💡 Example Prompts")
    gr.Examples(
        examples=EXAMPLES,
        inputs=[prompt, camera_motion],
        cache_examples=False,
    )

    gr.Markdown(
        "Model and inference code by [Robbyant](https://github.com/Robbyant/lingbot-video). "
        "Model license: Apache-2.0. This community demo uses the Dense 1.3B checkpoint without the separate 27B prompt rewriter."
    )

    inputs = [
        prompt,
        negative_prompt,
        structured_json,
        aspect_ratio,
        camera_motion,
        duration,
        fps,
        num_inference_steps,
        guidance_scale,
        seed,
        randomize_seed,
    ]
    outputs = [output_video, generation_info, caption_preview]

    generate_btn.click(
        fn=generate_video,
        inputs=inputs,
        outputs=outputs,
        api_name="generate_video",
        concurrency_limit=1,
    )
    prompt.submit(
        fn=generate_video,
        inputs=inputs,
        outputs=outputs,
        concurrency_limit=1,
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=8).launch(
        mcp_server=True,
        show_error=True,
    )