from __future__ import annotations

import gc
import os
import random
import tempfile
import traceback
import zipfile

# ZeroGPU currently needs this compatible CUDA/PyTorch build before importing torch.
os.system(
    'pip install --quiet --upgrade --pre --extra-index-url '
    'https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces'
)

import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import AutoPipelineForText2Image, FlowMatchEulerDiscreteScheduler
from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
from diffusers.utils.export_utils import export_to_video
from PIL import Image
from torchao.quantization import (
    Float8DynamicActivationFloat8WeightConfig,
    Int8WeightOnlyConfig,
    quantize_,
)
from transformers import CLIPImageProcessor


WAN_MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
WAN_WEIGHTS_ID = "cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers"
IMAGE_MODEL_ID = "stabilityai/sd-turbo"
SAFETY_MODEL_ID = "CompVis/stable-diffusion-safety-checker"
FPS = 16
MAX_SEED = np.iinfo(np.int32).max

QUALITY_PROFILES = {
    "Schnell · 480p · 4 Schritte": {"area": 480 * 832, "steps": 4, "seconds": 70},
    "Ausgewogen · 480p · 8 Schritte": {"area": 480 * 832, "steps": 8, "seconds": 110},
    "Qualität · 576p · 12 Schritte": {"area": 576 * 1024, "steps": 12, "seconds": 145},
    "Sehr hoch · 720p · 16 Schritte": {"area": 720 * 1280, "steps": 16, "seconds": 300},
    "Maximum · 720p · 24 Schritte": {"area": 720 * 1280, "steps": 24, "seconds": 450},
}

ASPECT_RATIOS = {
    "Automatisch vom Startbild": None,
    "16:9 Querformat": 16 / 9,
    "9:16 Hochformat": 9 / 16,
    "1:1 Quadrat": 1.0,
    "4:3 Querformat": 4 / 3,
    "3:4 Hochformat": 3 / 4,
}

DEFAULT_NEGATIVE_PROMPT = (
    "static, frozen, blurry, low quality, compression artifacts, text, subtitles, watermark, "
    "deformed anatomy, extra limbs, duplicated subject, incoherent motion, camera jitter"
)


print("Loading Wan 2.2 pipeline …")
wan_pipe = WanImageToVideoPipeline.from_pretrained(
    WAN_MODEL_ID,
    transformer=WanTransformer3DModel.from_pretrained(
        WAN_WEIGHTS_ID,
        subfolder="transformer",
        torch_dtype=torch.bfloat16,
        device_map="cuda",
    ),
    transformer_2=WanTransformer3DModel.from_pretrained(
        WAN_WEIGHTS_ID,
        subfolder="transformer_2",
        torch_dtype=torch.bfloat16,
        device_map="cuda",
    ),
    torch_dtype=torch.bfloat16,
)
wan_pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(wan_pipe.scheduler.config, shift=8.0)
wan_pipe.to("cuda")
wan_pipe.vae.enable_tiling()

quantize_(wan_pipe.text_encoder, Int8WeightOnlyConfig())
quantize_(wan_pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
quantize_(wan_pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
spaces.aoti_blocks_load(wan_pipe.transformer, "zerogpu-aoti/Wan2", variant="fp8da")
spaces.aoti_blocks_load(wan_pipe.transformer_2, "zerogpu-aoti/Wan2", variant="fp8da")

print("Loading compact text-to-image starter …")
safety_checker = StableDiffusionSafetyChecker.from_pretrained(
    SAFETY_MODEL_ID,
    torch_dtype=torch.float16,
)
feature_extractor = CLIPImageProcessor.from_pretrained(SAFETY_MODEL_ID)
image_pipe = AutoPipelineForText2Image.from_pretrained(
    IMAGE_MODEL_ID,
    torch_dtype=torch.float16,
    variant="fp16",
    safety_checker=safety_checker,
    feature_extractor=feature_extractor,
    requires_safety_checker=True,
)
image_pipe.to("cuda")
image_pipe.set_progress_bar_config(disable=True)
print("Pipelines ready.")


def _target_size(image: Image.Image | None, aspect_choice: str, area: int) -> tuple[int, int]:
    selected_ratio = ASPECT_RATIOS.get(aspect_choice)
    if selected_ratio is None:
        selected_ratio = image.width / image.height if image is not None else 16 / 9
    width = max(16, round(np.sqrt(area * selected_ratio)) // 16 * 16)
    height = max(16, round(np.sqrt(area / selected_ratio)) // 16 * 16)
    return width, height


def _cover_resize(image: Image.Image, size: tuple[int, int]) -> Image.Image:
    image = image.convert("RGB")
    target_width, target_height = size
    scale = max(target_width / image.width, target_height / image.height)
    resized = image.resize(
        (round(image.width * scale), round(image.height * scale)),
        Image.Resampling.LANCZOS,
    )
    left = (resized.width - target_width) // 2
    top = (resized.height - target_height) // 2
    return resized.crop((left, top, left + target_width, top + target_height))


def _to_pil_frame(frame) -> Image.Image:
    """Normalize Diffusers frame outputs to an RGB PIL image."""
    if isinstance(frame, Image.Image):
        return frame.convert("RGB")
    if isinstance(frame, torch.Tensor):
        frame = frame.detach().cpu().numpy()

    array = np.asarray(frame)
    if array.ndim == 3 and array.shape[0] in (1, 3, 4) and array.shape[-1] not in (1, 3, 4):
        array = np.moveaxis(array, 0, -1)
    if np.issubdtype(array.dtype, np.floating):
        if array.size and float(array.min()) < 0:
            array = (array + 1.0) / 2.0
        array = np.clip(array, 0.0, 1.0) * 255.0
    array = np.clip(array, 0, 255).astype(np.uint8)
    return Image.fromarray(array).convert("RGB")


def _frame_count(duration_seconds: float) -> int:
    desired = int(round(float(duration_seconds) * FPS))
    return int(np.clip(round((desired - 1) / 4) * 4 + 1, 17, 97))


def estimate_gpu_duration(
    start_mode,
    uploaded_image,
    image_prompt,
    motion_prompt,
    extension_prompt,
    quality,
    aspect_ratio,
    clip_duration,
    extensions,
    negative_prompt,
    seed,
    randomize_seed,
    progress=None,
):
    del uploaded_image, image_prompt, motion_prompt, extension_prompt, aspect_ratio
    del negative_prompt, seed, randomize_seed, progress
    profile = QUALITY_PROFILES.get(quality, QUALITY_PROFILES["Ausgewogen · 480p · 8 Schritte"])
    frame_factor = _frame_count(float(clip_duration)) / _frame_count(2.1)
    total = profile["seconds"] * frame_factor * (int(extensions) + 1)
    if start_mode == "Text → Startbild → Video":
        total += 20
    # xlarge is billed at 2× quota, so 1,200s maps to the 40-minute PRO allowance.
    return min(1200, max(60, round(total)))


@spaces.GPU(duration=estimate_gpu_duration, size="xlarge")
def generate_extended_video(
    start_mode,
    uploaded_image,
    image_prompt,
    motion_prompt,
    extension_prompt,
    quality,
    aspect_ratio,
    clip_duration,
    extensions,
    negative_prompt,
    seed,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True),
):
    if not motion_prompt or not str(motion_prompt).strip():
        raise gr.Error("Bitte beschreibe die Bewegung im Video-Prompt.")

    extensions = int(extensions)
    segment_count = extensions + 1
    profile = QUALITY_PROFILES[quality]
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    if start_mode == "Bild → Video":
        if uploaded_image is None:
            raise gr.Error("Bitte lade ein Startbild hoch.")
        initial_source = uploaded_image.convert("RGB")
        target = _target_size(initial_source, aspect_ratio, profile["area"])
        first_frame = _cover_resize(initial_source, target)
    else:
        if not image_prompt or not str(image_prompt).strip():
            raise gr.Error("Bitte gib einen Prompt für das Startbild ein.")
        target = _target_size(None, aspect_ratio, profile["area"])
        progress(0.03, desc="Startbild aus Text erzeugen …")
        first_frame = image_pipe(
            prompt=str(image_prompt).strip(),
            negative_prompt="blurry, low quality, text, watermark",
            width=target[0],
            height=target[1],
            num_inference_steps=2,
            guidance_scale=0.0,
            generator=torch.Generator(device="cuda").manual_seed(current_seed),
        ).images[0].convert("RGB")

    current_frame = first_frame
    all_frames: list[Image.Image] = []
    clip_paths: list[str] = []
    extension_text = str(extension_prompt or "").strip()
    base_prompt = str(motion_prompt).strip()
    num_frames = _frame_count(float(clip_duration))
    print(
        f"Generation settings: quality={quality}, duration={clip_duration}s, "
        f"frames={num_frames}, extensions={extensions}, size={target[0]}x{target[1]}"
    )

    last_combined_path = None
    last_archive_path = None

    for index in range(segment_count):
        fraction = 0.08 + (index / segment_count) * 0.78
        progress(fraction, desc=f"Clip {index + 1}/{segment_count} mit Wan erzeugen …")
        prompt = base_prompt
        if index > 0:
            prompt = extension_text or (
                base_prompt
                + ", continue the action naturally from the previous clip, seamless temporal continuity"
            )

        try:
            frames = wan_pipe(
                image=current_frame,
                prompt=prompt,
                negative_prompt=negative_prompt,
                height=target[1],
                width=target[0],
                num_frames=num_frames,
                guidance_scale=1.0,
                guidance_scale_2=1.0,
                num_inference_steps=profile["steps"],
                generator=torch.Generator(device="cuda").manual_seed(
                    (current_seed + index) % MAX_SEED
                ),
            ).frames[0]
        except Exception as exc:
            traceback.print_exc()
            gc.collect()
            torch.cuda.empty_cache()
            if last_combined_path is None or last_archive_path is None:
                raise gr.Error(
                    f"Schon der erste Clip konnte nicht erzeugt werden: {type(exc).__name__}: {exc}"
                ) from exc

            status = (
                f"⚠️ Verlängerung {index} ist fehlgeschlagen. Die zuvor erzeugten {index} "
                f"Clip(s) sind unten als Video und ZIP gesichert. Lade den angezeigten letzten "
                f"Frame herunter und verwende ihn mit ‚Bild → Video‘, um dort fortzusetzen. "
                f"Fehler: {type(exc).__name__}: {exc}"
            )
            yield (
                last_combined_path,
                first_frame,
                current_frame,
                last_archive_path,
                current_seed,
                status,
            )
            return

        frames = [_to_pil_frame(frame) for frame in frames]

        clip = tempfile.NamedTemporaryFile(suffix=f"-clip-{index + 1}.mp4", delete=False)
        clip.close()
        export_to_video(frames, clip.name, fps=FPS)
        clip_paths.append(clip.name)

        if index == 0:
            all_frames.extend(frames)
        else:
            all_frames.extend(frames[1:])
        current_frame = frames[-1].convert("RGB")

        progress(
            0.08 + ((index + 0.9) / segment_count) * 0.78,
            desc=f"Checkpoint {index + 1}/{segment_count} sichern …",
        )
        combined = tempfile.NamedTemporaryFile(suffix="-wan-checkpoint.mp4", delete=False)
        combined.close()
        export_to_video(all_frames, combined.name, fps=FPS)

        archive = tempfile.NamedTemporaryFile(suffix="-wan-checkpoint.zip", delete=False)
        archive.close()
        with zipfile.ZipFile(archive.name, "w", compression=zipfile.ZIP_DEFLATED) as bundle:
            for clip_index, path in enumerate(clip_paths, start=1):
                bundle.write(path, arcname=f"clip-{clip_index:02d}.mp4")
            bundle.write(combined.name, arcname="complete-video.mp4")

        last_combined_path = combined.name
        last_archive_path = archive.name
        completed = index + 1
        total_seconds = len(all_frames) / FPS
        if completed == segment_count:
            status = (
                f"✅ Fertig: {completed} Clips ({extensions} Verlängerungen), "
                f"ca. {total_seconds:.1f} Sekunden, {target[0]}×{target[1]}, Seed {current_seed}."
            )
        else:
            status = (
                f"💾 Checkpoint gesichert: {completed}/{segment_count} Clips, "
                f"ca. {total_seconds:.1f} Sekunden. Video, ZIP und letzter Frame sind bereits verfügbar."
            )

        gc.collect()
        torch.cuda.empty_cache()
        yield combined.name, first_frame, current_frame, archive.name, current_seed, status

    # The loop yields the final checkpoint as its last result.
    return


CSS = """
.gradio-container { max-width: 1180px !important; }
.hero { text-align: center; margin: 0 auto 1.2rem; }
.hero h1 { font-size: 2.35rem; margin-bottom: .15rem; }
.flow { padding: 12px 16px; border-radius: 14px; border: 1px solid var(--border-color-primary); }
"""

with gr.Blocks() as demo:
    gr.HTML(
        "<div class='hero'><h1>♾️ Wan Extender</h1>"
        "<p>Startframe erzeugen oder hochladen · Wan-Clip generieren · automatisch verlängern</p></div>"
    )
    gr.Markdown(
        "**Logik:** Bei zwei Verlängerungen werden drei Clips erzeugt. Der letzte Frame jedes Clips "
        "wird automatisch zum ersten Frame des nächsten. Mehr Verlängerungen benötigen entsprechend "
        "mehr ZeroGPU-Kontingent."
    )

    with gr.Row():
        with gr.Column(scale=5, elem_classes="flow"):
            start_mode = gr.Radio(
                ["Text → Startbild → Video", "Bild → Video"],
                value="Text → Startbild → Video",
                label="Startmodus",
            )
            image_prompt = gr.Textbox(
                label="Prompt für das Startbild",
                placeholder="A cinematic wide shot of …",
                lines=3,
            )
            uploaded_image = gr.Image(
                type="pil",
                label="Optionales Startbild (für Modus Bild → Video)",
                sources=["upload", "clipboard"],
            )
            motion_prompt = gr.Textbox(
                label="Video- und Bewegungs-Prompt",
                placeholder="The camera slowly moves forward while …",
                lines=3,
            )
            extension_prompt = gr.Textbox(
                label="Prompt für Verlängerungen (optional)",
                placeholder="Leer lassen, um den Video-Prompt natürlich fortzuführen.",
                lines=2,
            )

        with gr.Column(scale=4, elem_classes="flow"):
            quality = gr.Dropdown(
                list(QUALITY_PROFILES),
                value="Ausgewogen · 480p · 8 Schritte",
                label="Qualitätsprofil",
            )
            aspect_ratio = gr.Dropdown(
                list(ASPECT_RATIOS),
                value="16:9 Querformat",
                label="Seitenverhältnis",
            )
            clip_duration = gr.Slider(
                1.1,
                6.0,
                value=2.1,
                step=0.1,
                label="Dauer pro Clip (Sekunden)",
                info="Bis 5 Sekunden regulär; 5,1–6 Sekunden experimenteller Langclip-Modus.",
            )
            extensions = gr.Slider(
                0,
                5,
                value=1,
                step=1,
                label="Anzahl Verlängerungen",
                info="0 = ein Clip, 5 = sechs Clips insgesamt",
            )
            with gr.Accordion("Erweiterte Einstellungen", open=False):
                negative_prompt = gr.Textbox(
                    value=DEFAULT_NEGATIVE_PROMPT,
                    label="Negative Prompt",
                    lines=3,
                )
                seed = gr.Number(value=42, precision=0, label="Seed")
                randomize_seed = gr.Checkbox(value=True, label="Seed zufällig wählen")

            generate_button = gr.Button("Video erzeugen und verlängern", variant="primary", size="lg")
            status = gr.Markdown()

    output_video = gr.Video(label="Zusammenhängendes Ergebnis", autoplay=True)
    gr.Markdown(
        "Nach jedem Clip wird sofort ein Checkpoint ausgegeben. Bei einem Fehler in einer späteren "
        "Verlängerung bleiben Video, ZIP und letzter Frame erhalten."
    )
    with gr.Row():
        generated_first_frame = gr.Image(label="Tatsächlicher erster Frame", interactive=False)
        generated_last_frame = gr.Image(label="Letzter Frame der letzten Verlängerung", interactive=False)
    clips_download = gr.File(label="Aktuellen Checkpoint als ZIP herunterladen")

    generate_button.click(
        generate_extended_video,
        inputs=[
            start_mode,
            uploaded_image,
            image_prompt,
            motion_prompt,
            extension_prompt,
            quality,
            aspect_ratio,
            clip_duration,
            extensions,
            negative_prompt,
            seed,
            randomize_seed,
        ],
        outputs=[
            output_video,
            generated_first_frame,
            generated_last_frame,
            clips_download,
            seed,
            status,
        ],
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        show_error=True,
        theme=gr.themes.Soft(),
        css=CSS,
    )