import os
import gc
import torch
import gradio as gr

from diffusers import WanAnimatePipeline
from diffusers.utils import load_image, load_video, export_to_video


# ============================================================
# CONFIG
# ============================================================

MODEL_ID = "Wan-AI/Wan2.2-Animate-14B-Diffusers"

OUTPUT_DIR = "/tmp/outputs"
os.makedirs(OUTPUT_DIR, exist_ok=True)


# ============================================================
# PRESETS
# ============================================================

PRESETS = {
    "Realistic": {
        "prompt": (
            "A highly realistic video of the reference character performing "
            "the movements from the input motion video. Preserve identity, "
            "facial appearance, hairstyle and clothing. Natural body motion, "
            "realistic lighting and physics."
        ),
        "steps": 20,
        "guidance": 1.0,
    },

    "Cinematic": {
        "prompt": (
            "A cinematic photorealistic video of the character performing "
            "the exact movements and actions from the reference motion video. "
            "Natural facial expressions, realistic skin, detailed clothing, "
            "cinematic lighting, shallow depth of field, professional camera."
        ),
        "steps": 20,
        "guidance": 1.0,
    },

    "Portrait": {
        "prompt": (
            "A photorealistic portrait video of the reference character. "
            "Preserve the character's identity and facial features while "
            "accurately following the body and facial motion from the input video. "
            "Natural expression, realistic skin and cinematic portrait lighting."
        ),
        "steps": 20,
        "guidance": 1.0,
    },

    "Anime": {
        "prompt": (
            "An anime-style cinematic video featuring the reference character, "
            "accurately reproducing the movement and performance from the input "
            "motion video. Consistent character identity, expressive animation, "
            "detailed anime background."
        ),
        "steps": 20,
        "guidance": 1.0,
    },
}


# ============================================================
# LOAD MODEL
# ============================================================

print("Loading Wan2.2 Animate...")

if not torch.cuda.is_available():
    raise RuntimeError(
        "CUDA GPU is required. "
        "Use a GPU-enabled Hugging Face Space."
    )

print("GPU:", torch.cuda.get_device_name(0))

pipe = WanAnimatePipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
)

# Important for limited VRAM
pipe.enable_model_cpu_offload()

print("Model loaded.")


# ============================================================
# GENERATION
# ============================================================

def generate_video(
    prompt,
    reference_image,
    motion_video,
    preset,
    seed,
    inference_steps,
    guidance_scale,
    progress=gr.Progress(),
):

    if reference_image is None:
        raise gr.Error("Please upload a reference image.")

    if motion_video is None:
        raise gr.Error("Please upload a motion video.")

    # --------------------------------------------------------
    # PRESET
    # --------------------------------------------------------

    preset_config = PRESETS[preset]

    if not prompt or not prompt.strip():
        prompt = preset_config["prompt"]

    # --------------------------------------------------------
    # LOAD INPUTS
    # --------------------------------------------------------

    progress(0.1, desc="Loading reference image...")

    image = load_image(reference_image)

    progress(0.2, desc="Loading motion video...")

    motion = load_video(motion_video)

    # --------------------------------------------------------
    # CONDITIONING
    #
    # For the basic workflow we use the motion video as
    # both pose and face conditioning.
    #
    # For best results, replace this with Wan Animate's
    # official preprocessing pipeline.
    # --------------------------------------------------------

    pose_video = motion
    face_video = motion

    # --------------------------------------------------------
    # GENERATION
    # --------------------------------------------------------

    progress(0.3, desc="Generating video...")

    generator = torch.Generator(
        device="cuda"
    ).manual_seed(int(seed))

    with torch.inference_mode():

        result = pipe(
            image=image,
            pose_video=pose_video,
            face_video=face_video,
            prompt=prompt,
            mode="animate",
            segment_frame_length=77,
            prev_segment_conditioning_frames=1,
            guidance_scale=float(guidance_scale),
            num_inference_steps=int(inference_steps),
            generator=generator,
        ).frames[0]

    progress(0.9, desc="Encoding output video...")

    # --------------------------------------------------------
    # SAVE
    # --------------------------------------------------------

    output_path = os.path.join(
        OUTPUT_DIR,
        f"generated_{int(seed)}.mp4"
    )

    export_to_video(
        result,
        output_path,
        fps=30,
    )

    # --------------------------------------------------------
    # CLEANUP
    # --------------------------------------------------------

    gc.collect()
    torch.cuda.empty_cache()

    progress(1.0, desc="Done")

    return output_path


# ============================================================
# PRESET HANDLER
# ============================================================

def update_preset(preset):

    config = PRESETS[preset]

    return (
        config["prompt"],
        config["steps"],
        config["guidance"],
    )


# ============================================================
# GRADIO UI
# ============================================================

with gr.Blocks(
    title="Wan2.2 Character Animation"
) as demo:

    gr.Markdown(
        """
        # 🎬 Wan2.2 Character Animation

        Generate a video using:

        **Reference Photo + Motion Video + Prompt → Generated Video**

        The reference image provides the character identity.
        The motion video provides the movement.
        """
    )

    with gr.Row():

        # ----------------------------------------------------
        # LEFT
        # ----------------------------------------------------

        with gr.Column():

            prompt = gr.Textbox(
                label="Prompt",
                placeholder=(
                    "Describe the generated video..."
                ),
                lines=5,
            )

            preset = gr.Dropdown(
                choices=list(PRESETS.keys()),
                value="Realistic",
                label="Style Preset",
            )

            reference_image = gr.Image(
                label="Reference Photo / Character",
                type="filepath",
            )

            motion_video = gr.Video(
                label="Motion Video",
                sources=["upload"],
            )

        # ----------------------------------------------------
        # RIGHT
        # ----------------------------------------------------

        with gr.Column():

            output_video = gr.Video(
                label="Generated Video",
                autoplay=True,
            )

            generate_button = gr.Button(
                "🎬 Generate Video",
                variant="primary",
            )

    # ========================================================
    # ADVANCED SETTINGS
    # ========================================================

    with gr.Accordion(
        "Advanced Settings",
        open=False,
    ):

        seed = gr.Number(
            label="Seed",
            value=42,
            precision=0,
        )

        inference_steps = gr.Slider(
            minimum=5,
            maximum=50,
            value=20,
            step=1,
            label="Inference Steps",
        )

        guidance_scale = gr.Slider(
            minimum=0.5,
            maximum=5.0,
            value=1.0,
            step=0.1,
            label="Guidance Scale",
        )

    # ========================================================
    # PRESET EVENT
    # ========================================================

    preset.change(
        fn=update_preset,
        inputs=[preset],
        outputs=[
            prompt,
            inference_steps,
            guidance_scale,
        ],
    )

    # ========================================================
    # GENERATE EVENT
    # ========================================================

    generate_button.click(
        fn=generate_video,
        inputs=[
            prompt,
            reference_image,
            motion_video,
            preset,
            seed,
            inference_steps,
            guidance_scale,
        ],
        outputs=output_video,
    )


# ============================================================
# LAUNCH
# ============================================================

if __name__ == "__main__":

    demo.queue(
        max_size=10,
        default_concurrency_limit=1,
    )

    demo.launch()