import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  (must precede torch / CUDA imports)
import torch  # noqa: E402
import numpy as np  # noqa: E402
import tempfile  # noqa: E402
import subprocess  # noqa: E402
import gradio as gr  # noqa: E402
from PIL import Image  # noqa: E402
from huggingface_hub import snapshot_download  # noqa: E402
from safetensors.torch import load_file  # noqa: E402
from diffusers import WanAnimate2Pipeline, WanAnimate2Transformer3DModel  # noqa: E402
from diffusers.utils import export_to_video, load_image  # noqa: E402

# ---------------------------------------------------------------------------
# ZeroGPU eager BlockMask workaround for flex_attention
# ---------------------------------------------------------------------------
import torch.nn.attention.flex_attention as _flex_mod  # noqa: E402

_orig_create_block_mask = _flex_mod.create_block_mask


def _create_block_mask_no_compile(*args, **kwargs):
    kwargs["_compile"] = False
    return _orig_create_block_mask(*args, **kwargs)


_flex_mod.create_block_mask = _create_block_mask_no_compile

try:
    import diffusers.models.transformers.transformer_wan_animate_2 as _wa2  # noqa: E402

    _wa2.create_block_mask = _create_block_mask_no_compile
except Exception as _e:  # pragma: no cover
    print(f"[patch] could not patch transformer_wan_animate_2.create_block_mask: {_e!r}")

from torch.nn.attention.flex_attention import BlockMask as _BlockMask  # noqa: E402
import torch.nn.functional as _F  # noqa: E402
import diffusers.models.attention_dispatch as _attn_dispatch  # noqa: E402

_dense_mask_cache = {}


def _blockmask_to_dense_bool(block_mask, seq_len_q, seq_len_kv, device):
    key = (id(block_mask), seq_len_q, seq_len_kv)
    cached = _dense_mask_cache.get(key)
    if cached is not None:
        return cached
    mask_mod = block_mask.mask_mod
    zero = torch.zeros((), dtype=torch.long, device=device)
    kv_idx = torch.arange(seq_len_kv, device=device)
    dense = torch.empty((seq_len_q, seq_len_kv), dtype=torch.bool, device=device)
    q_chunk = max(1, min(seq_len_q, 512))
    for q0 in range(0, seq_len_q, q_chunk):
        q1 = min(seq_len_q, q0 + q_chunk)
        rows = q1 - q0
        qg = torch.arange(q0, q1, device=device).view(rows, 1).expand(rows, seq_len_kv)
        kg = kv_idx.view(1, seq_len_kv).expand(rows, seq_len_kv)
        dense[q0:q1] = mask_mod(zero, zero, qg, kg).to(torch.bool)
    dense = dense.view(1, 1, seq_len_q, seq_len_kv)
    if len(_dense_mask_cache) > 4:
        _dense_mask_cache.clear()
    _dense_mask_cache[key] = dense
    return dense


def _flex_backend_sdpa(
    query, key, value, attn_mask=None, is_causal=False, scale=None,
    enable_gqa=False, return_lse=False, _parallel_config=None,
):
    batch_size, seq_len_q, _, _ = query.shape
    seq_len_kv = key.shape[1]
    q = query.permute(0, 2, 1, 3)
    k = key.permute(0, 2, 1, 3)
    v = value.permute(0, 2, 1, 3)
    sdpa_mask = None
    if isinstance(attn_mask, _BlockMask):
        sdpa_mask = _blockmask_to_dense_bool(attn_mask, seq_len_q, seq_len_kv, query.device)
    elif attn_mask is not None and torch.is_tensor(attn_mask):
        sdpa_mask = attn_mask
    out = _F.scaled_dot_product_attention(
        q, k, v, attn_mask=sdpa_mask, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa,
    )
    out = out.permute(0, 2, 1, 3)
    if return_lse:
        return out, None
    return out


_attn_dispatch._native_flex_attention = _flex_backend_sdpa
try:
    _reg = _attn_dispatch._AttentionBackendRegistry
    _reg._backends[_attn_dispatch.AttentionBackendName.FLEX] = _flex_backend_sdpa
    print("[patch] FLEX backend replaced with memory-efficient SDPA path")
except Exception as _e:  # pragma: no cover
    print(f"[patch] could not re-register FLEX backend: {_e!r}")

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

# Enhanced negative prompt to eliminate artifacts, pixelation, and motion blur
DEFAULT_NEGATIVE_PROMPT = (
    "blur, low resolution, artifacts, pixelated, jitter, distortion, noisy, unnatural motion, "
    "morphing, overexposed, static, bad anatomy, bad hands, deformed limbs, floating limbs, "
    "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，"
    "低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的"
)


def _remap_transformer_key(key: str) -> str:
    if not key.startswith("blocks."):
        return key
    parts = key.split(".")
    if len(parts) < 4 or parts[2] != "block":
        return key
    n = parts[1]
    rest = parts[3:]
    sub = rest[0]
    if sub in ("self_attn", "cross_attn"):
        proj = rest[1]
        tail = ".".join(rest[2:])
        proj_map = {
            "q": "to_q",
            "k": "to_k",
            "v": "to_v",
            "o": "to_out.0",
            "k_img": "add_k_proj",
            "v_img": "add_v_proj",
            "norm_q": "norm_q",
            "norm_k": "norm_k",
            "norm_k_img": "norm_added_k",
        }
        mapped = proj_map.get(proj, proj)
        new = f"blocks.{n}.{sub}.{mapped}"
        return f"{new}.{tail}" if tail else new
    return f"blocks.{n}." + ".".join(rest)


def _load_transformer():
    ckpt_dir = snapshot_download(MODEL_ID, allow_patterns=["transformer/*"])
    tdir = os.path.join(ckpt_dir, "transformer")
    import glob

    remapped = {}
    for shard in sorted(glob.glob(os.path.join(tdir, "*.safetensors"))):
        sd = load_file(shard)
        for k, v in sd.items():
            remapped[_remap_transformer_key(k)] = v.to(torch.bfloat16)
        del sd

    model = WanAnimate2Transformer3DModel.from_config(
        WanAnimate2Transformer3DModel.load_config(tdir)
    ).to(torch.bfloat16)
    missing, unexpected = model.load_state_dict(remapped, strict=False)
    missing = [m for m in missing if "kv_cache" not in m]
    if missing:
        raise RuntimeError(f"Missing transformer keys after remap: {missing[:20]} (total {len(missing)})")
    if unexpected:
        raise RuntimeError(f"Unexpected transformer keys after remap: {unexpected[:20]} (total {len(unexpected)})")
    return model


transformer = _load_transformer()
pipe = WanAnimate2Pipeline.from_pretrained(
    MODEL_ID, transformer=transformer, torch_dtype=torch.bfloat16
)
pipe.to("cuda")

AOTI_REPO = "multimodalart/Wan2.2-Animate-2-14B-Distilled-aoti"
AOTI_FFN_DIR = "WanAnimate2TransformerBlockFFN"
try:
    from pathlib import Path as _Path
    from spaces.zero.torch.aoti import aoti_load_from_module_dir

    _aoti_root = _Path(snapshot_download(AOTI_REPO, allow_patterns=[f"{AOTI_FFN_DIR}/*"]))
    _ffn_pkg_dir = _aoti_root / AOTI_FFN_DIR
    if (_ffn_pkg_dir / "package.pt2").exists():
        _ffns = [blk.ffn for blk in pipe.transformer.blocks]
        aoti_load_from_module_dir(_ffns, _ffn_pkg_dir)
        print(f"[aoti] patched {len(_ffns)} block.ffn modules from {AOTI_REPO}/{AOTI_FFN_DIR}")
    else:
        print(f"[aoti] {AOTI_FFN_DIR}/package.pt2 not found in {AOTI_REPO}; running eager")
except Exception as _e:  # pragma: no cover
    print(f"[aoti] AoTI load failed ({_e!r}); running eager")


def _trim_video(src_path: str, max_seconds: float, target_fps: int = 24) -> str:
    out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    try:
        subprocess.run(
            [
                "ffmpeg", "-y", "-i", src_path,
                "-t", str(max_seconds),
                "-r", str(target_fps),
                "-an",
                "-c:v", "libx264", "-pix_fmt", "yuv420p",
                out,
            ],
            check=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        return out
    except Exception as e:
        print(f"[trim] ffmpeg failed ({e!r}); using original video")
        return src_path


def _export_hq_video(frames, target_fps: int = 24) -> str:
    """Exports and recompresses video using H.264 CRF 17 for crystal clear outputs."""
    raw_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    export_to_video(frames, raw_path, fps=target_fps)

    hq_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    try:
        subprocess.run(
            [
                "ffmpeg", "-y", "-i", raw_path,
                "-c:v", "libx264",
                "-crf", "17",
                "-preset", "slow",
                "-pix_fmt", "yuv420p",
                "-movflags", "+faststart",
                hq_path,
            ],
            check=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        os.remove(raw_path)
        return hq_path
    except Exception:
        return raw_path


def _estimate_duration(image, driving_video, prompt, max_seconds=60.0, height=720,
                       width=720, num_inference_steps=14, *args, **kwargs):
    try:
        segments = max(1, int(np.ceil(float(max_seconds) / 3.4)))
    except Exception:
        segments = 1
    try:
        pixels = float(height) * float(width)
        steps = max(1, int(num_inference_steps))
    except Exception:
        pixels, steps = 720 * 720, 14
    P0 = 320.0 * 480.0
    per_step = 3.25 * (pixels / P0) ** 3.1
    per_segment = steps * per_step
    total = 10.0 + segments * per_segment
    return int(min(7200, total * 1.2 + 15))


@spaces.GPU(duration=_estimate_duration, size="xlarge")
def animate(
    image,
    driving_video,
    prompt,
    max_seconds: float = 60.0,
    height: int = 720,
    width: int = 720,
    num_inference_steps: int = 14,
    guidance_scale: float = 2.5,
    sample_shift: float = 5.0,
    negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
    seed: int = 0,
    progress=gr.Progress(track_tqdm=True),
):
    if image is None:
        raise gr.Error("Please provide a reference character image.")
    if driving_video is None:
        raise gr.Error("Please provide a driving video.")

    if isinstance(image, str):
        image = load_image(image)
    if not isinstance(image, Image.Image):
        image = Image.fromarray(np.asarray(image))
    image = image.convert("RGB")

    height = int(height) - (int(height) % 16)
    width = int(width) - (int(width) % 16)

    trimmed = _trim_video(driving_video, float(max_seconds), target_fps=24)

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

    # Append quality booster tokens to the user prompt
    quality_prompt = f"{prompt.strip() if prompt else 'a person'}, 8k resolution, photorealistic, cinematic lighting, sharp focus, highly detailed, master quality"

    try:
        output = pipe(
            image=image,
            driving_video=trimmed,
            prompt=quality_prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            fps=24,
            num_inference_steps=int(num_inference_steps),
            guidance_scale=float(guidance_scale),
            sample_shift=float(sample_shift),
            seed=int(seed),
            generator=generator,
            output_type="np",
        )
    except Exception as exc:
        import traceback
        tb = traceback.format_exc()
        print("[animate] inference failed:\n" + tb, flush=True)
        raise gr.Error(f"{type(exc).__name__}: {exc}\n{tb[-1500:]}")
        
    frames = output.frames[0]
    return _export_hq_video(frames, target_fps=24)


CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;600;700;900&family=Space+Grotesk:wght@500;700&display=swap');

:root {
    --primary-gradient: linear-gradient(135deg, #FF3366 0%, #BA13F9 50%, #4361EE 100%);
    --surface-glass: rgba(18, 22, 36, 0.75);
    --surface-card: rgba(28, 33, 53, 0.65);
    --neon-border: rgba(186, 19, 249, 0.35);
    --neon-glow: 0 8px 32px 0 rgba(186, 19, 249, 0.25);
    --text-primary: #F8FAFC;
    --text-secondary: #94A3B8;
}

body, .gradio-container {
    font-family: 'Outfit', sans-serif !important;
    background: radial-gradient(circle at 10% 20%, rgba(90, 24, 154, 0.2) 0%, transparent 40%),
                radial-gradient(circle at 90% 80%, rgba(255, 51, 102, 0.15) 0%, transparent 40%),
                #0B0E17 !important;
    color: var(--text-primary) !important;
    min-height: 100vh;
}

#main-container {
    max-width: 1280px;
    margin: 0 auto;
    padding: 24px 16px;
}

.hero-header {
    text-align: center;
    padding: 32px 20px;
    margin-bottom: 24px;
    background: var(--surface-glass);
    border-radius: 24px;
    border: 1px solid var(--neon-border);
    backdrop-filter: blur(20px);
    -webkit-backdrop-filter: blur(20px);
    box-shadow: var(--neon-glow);
}

.hero-title {
    font-family: 'Space Grotesk', sans-serif !important;
    font-size: 3rem !important;
    font-weight: 900 !important;
    letter-spacing: -1px;
    background: var(--primary-gradient);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 8px;
    text-transform: uppercase;
}

.hero-subtitle {
    font-size: 1.15rem;
    color: var(--text-secondary);
    max-width: 600px;
    margin: 0 auto 16px auto;
}

.credit-badge {
    display: inline-flex;
    align-items: center;
    gap: 8px;
    background: rgba(255, 255, 255, 0.06);
    border: 1px solid rgba(255, 255, 255, 0.12);
    padding: 6px 18px;
    border-radius: 9999px;
    font-size: 0.85rem;
    font-weight: 600;
    color: #E2E8F0;
    letter-spacing: 0.5px;
    box-shadow: 0 4px 15px rgba(0,0,0,0.3);
}

.credit-badge span {
    background: var(--primary-gradient);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    font-weight: 700;
}

.glass-panel {
    background: var(--surface-card) !important;
    border-radius: 20px !important;
    border: 1px solid rgba(255, 255, 255, 0.08) !important;
    backdrop-filter: blur(16px);
    box-shadow: 0 10px 30px rgba(0,0,0,0.35);
    padding: 16px !important;
}

.glow-button {
    background: var(--primary-gradient) !important;
    border: none !important;
    border-radius: 14px !important;
    color: #FFFFFF !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
    letter-spacing: 0.5px;
    box-shadow: 0 0 25px rgba(186, 19, 249, 0.45) !important;
    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
    cursor: pointer !important;
    padding: 12px 24px !important;
}

.glow-button:hover {
    transform: translateY(-2px) scale(1.01);
    box-shadow: 0 0 35px rgba(255, 51, 102, 0.65) !important;
}

.accordion-glass {
    background: rgba(18, 22, 36, 0.5) !important;
    border: 1px solid rgba(255, 255, 255, 0.08) !important;
    border-radius: 16px !important;
    margin-top: 20px;
}

footer {
    display: none !important;
}
"""

with gr.Blocks(theme=gr.themes.Base(), css=CUSTOM_CSS) as demo:
    with gr.Column(elem_id="main-container"):
        gr.HTML(
            """
            <div class="hero-header">
                <h1 class="hero-title">PR MOTIONS AI</h1>
                <p class="hero-subtitle">Ultra HD Character Motion Synthesis powered by Wan2.2 Animate</p>
                <div class="credit-badge">
                    ⚡ Created & Optimized by <span>Prince</span>
                </div>
            </div>
            """
        )

        with gr.Row(equal_height=True):
            with gr.Column(elem_classes=["glass-panel"], scale=1):
                image = gr.Image(label="Reference Character", type="filepath", height=320)
                driving_video = gr.Video(label="Motion Source Video", height=320)
                prompt = gr.Textbox(
                    label="Animation Prompt",
                    placeholder="Describe character appearance, action & lighting...",
                    value="a person",
                    lines=2,
                )
                run = gr.Button("✨ Render HD Motion Video", variant="primary", elem_classes=["glow-button"])
                
            with gr.Column(elem_classes=["glass-panel"], scale=1):
                output_video = gr.Video(label="Rendered Motion Result (Ultra HD)", height=540)

        with gr.Accordion("⚙️ Advanced Quality & Production Settings", open=False, elem_classes=["accordion-glass"]):
            max_seconds = gr.Slider(
                label="Timeline Duration (Seconds: 60s = 1 min, 600s = 10 min)",
                minimum=60.0,
                maximum=600.0,
                step=10.0,
                value=60.0,
            )
            with gr.Row():
                height = gr.Slider(label="Render Height", minimum=320, maximum=1024, step=16, value=720)
                width = gr.Slider(label="Render Width", minimum=320, maximum=1024, step=16, value=720)
            with gr.Row():
                num_inference_steps = gr.Slider(
                    label="Denoising Steps (Higher = Sharper & Cleaner)", minimum=6, maximum=30, step=1, value=14,
                )
                guidance_scale = gr.Slider(
                    label="Guidance Scale (CFG Prompt Adherence)", minimum=1.0, maximum=7.0, step=0.5, value=2.5,
                )
            sample_shift = gr.Slider(
                label="Sample Shift", minimum=1.0, maximum=12.0, step=0.5, value=5.0,
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT, lines=2,
            )
            seed = gr.Number(label="Seed", value=0, precision=0)

        gr.Examples(
            examples=[
                ["examples/animate_character.jpeg", "examples/animate_driving.mp4", "a person dancing"],
                ["examples/replace_character.jpeg", "examples/replace_driving.mp4", "a person"],
            ],
            inputs=[image, driving_video, prompt],
            outputs=output_video,
            fn=animate,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        fn=animate,
        inputs=[
            image, driving_video, prompt, max_seconds, height, width,
            num_inference_steps, guidance_scale, sample_shift, negative_prompt, seed,
        ],
        outputs=output_video,
        api_name="animate",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True)