import os
import gc
import torch
import gradio as gr
import spaces
import ftfy
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)

pipe = None


# ============================================================
# 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,
    },
}


# ============================================================
# MODEL LOADING
# ============================================================

def get_pipeline():

    global pipe

    if pipe is None:

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

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

        pipe.enable_model_cpu_offload()

        print("Wan2.2 Animate loaded.")

    return pipe


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

@spaces.GPU
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."
        )

    # --------------------------------------------------------
    # LOAD MODEL AFTER ZERO GPU ALLOCATION
    # --------------------------------------------------------

    progress(
        0.05,
        desc="Loading Wan2.2 Animate..."
    )

    pipe = get_pipeline()

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

    config = PRESETS[preset]

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

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

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

    image = load_image(
        reference_image
    )

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

    motion = load_video(
        motion_video
    )

    # --------------------------------------------------------
    # TEMPORARY CONDITIONING
    #
    # NOTE:
    # For the real Wan2.2 Animate workflow, the motion video
    # should be processed into dedicated pose/face inputs.
    # This is the basic prototype.
    # --------------------------------------------------------

    pose_video = motion
    face_video = motion

    # --------------------------------------------------------
    # GENERATE
    # --------------------------------------------------------

    progress(
        0.35,
        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]

    # --------------------------------------------------------
    # EXPORT
    # --------------------------------------------------------

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

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

    export_to_video(
        result,
        output_path,
        fps=30,
    )

    progress(
        1.0,
        desc="Done"
    )

    gc.collect()

    return output_path


# ============================================================
# PRESET UPDATE
# ============================================================

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

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

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

    with gr.Row():

        with gr.Column():

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

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

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

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

        with gr.Column():

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

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

    # ========================================================
    # ADVANCED
    # ========================================================

    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",
        )

    # ========================================================
    # EVENTS
    # ========================================================

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

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


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

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

demo.launch()