"""WAMU_v3 image-to-video pipeline loading and sizing helpers.

Baseline (non-AOT) loading per the SRS's "baseline first, AOT after" sequencing: FR-6's
AOT-compiled-transformer requirement is deferred until AOT-package compatibility with WAMU_v3
is confirmed (open item C-2, tracked in issue #3). Weight quantization (below) is independent
of AOT compilation and is applied now: the full bf16 pipe packs to ~69GB, which made ZeroGPU's
per-call GPU init (loading the packed weights onto the GPU) exceed the requested call duration.
"""

from __future__ import annotations

import torch
import torch._dynamo
from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
from PIL import Image
from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, Int8WeightOnlyConfig, quantize_

MODEL_ID = "thornmaze/WAMU_v3_WAN2.2_I2V_LIGHTNING"

MAX_DIM = 832
MIN_DIM = 480
SQUARE_DIM = 640
MULTIPLE_OF = 16

FIXED_FPS = 16
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 321
MIN_DURATION = round(MIN_FRAMES_MODEL / FIXED_FPS, 1)
MAX_DURATION = round(MAX_FRAMES_MODEL / FIXED_FPS, 1)

DEFAULT_PROMPT = "make this image come alive, cinematic motion, smooth animation"
DEFAULT_NEGATIVE_PROMPT = (
    "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, "
    "最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, "
    "画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, "
    "杂乱的背景, 三条腿, 背景人很多, 倒着走"
)


def load_pipeline() -> WanImageToVideoPipeline:
    pipe = WanImageToVideoPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16).to("cuda")

    # Quantize to shrink the ~69GB bf16 pipe (drops per-call ZeroGPU init time and peak memory).
    # Independent of AOT compilation (FR-6/C-2, deferred) — same recipe as the reference project.
    quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
    torch._dynamo.reset()
    quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
    torch._dynamo.reset()
    quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
    torch._dynamo.reset()

    return pipe


def resize_image(image: Image.Image) -> Image.Image:
    """Resize/crop to the model's supported dimension range, keeping a multiple-of-16 size."""
    width, height = image.size
    if width == height:
        return image.resize((SQUARE_DIM, SQUARE_DIM), Image.LANCZOS)

    aspect_ratio = width / height
    max_ar = MAX_DIM / MIN_DIM
    min_ar = MIN_DIM / MAX_DIM

    image_to_resize = image
    if aspect_ratio > max_ar:
        crop_width = int(round(height * max_ar))
        left = (width - crop_width) // 2
        image_to_resize = image.crop((left, 0, left + crop_width, height))
        target_w, target_h = MAX_DIM, MIN_DIM
    elif aspect_ratio < min_ar:
        crop_height = int(round(width / min_ar))
        top = (height - crop_height) // 2
        image_to_resize = image.crop((0, top, width, top + crop_height))
        target_w, target_h = MIN_DIM, MAX_DIM
    else:
        if width > height:
            target_w = MAX_DIM
            target_h = int(round(target_w / aspect_ratio))
        else:
            target_h = MAX_DIM
            target_w = int(round(target_h * aspect_ratio))

    final_w = max(MIN_DIM, min(MAX_DIM, round(target_w / MULTIPLE_OF) * MULTIPLE_OF))
    final_h = max(MIN_DIM, min(MAX_DIM, round(target_h / MULTIPLE_OF) * MULTIPLE_OF))
    return image_to_resize.resize((final_w, final_h), Image.LANCZOS)


def get_num_frames(duration_seconds: float) -> int:
    raw = int(round(duration_seconds * FIXED_FPS))
    raw = max(MIN_FRAMES_MODEL, min(MAX_FRAMES_MODEL, raw))
    return ((raw - 1) // 4) * 4 + 1


def estimate_generation_seconds(
    width: int, height: int, num_frames: int, steps: int, guidance_scale: float
) -> float:
    """Rough @spaces.GPU duration budget, calibrated from the reference wan2-2-i2v-v3 Space.

    Exact limits are still an open item (SRS Section 4, #3) pending empirical testing on the
    dev Space; this heuristic is a starting point, not a final calibration.
    """
    base_frames_height_width = 81 * 832 * 624
    base_step_duration = 5.0
    factor = num_frames * width * height / base_frames_height_width
    step_duration = base_step_duration * factor**1.5
    gen_time = steps * step_duration
    if guidance_scale > 1:
        gen_time *= 2.4

    # The reference project's "15" was calibrated with AOT-compiled transformers (FR-6, deferred
    # here per C-2/issue #3). Without AOT, per-call ZeroGPU init (loading the quantized-but-not-
    # compiled pipe onto the GPU) plus eager-mode execution both take longer, so pad generously
    # until this can be measured empirically on the dev Space.
    init_overhead = 60
    return init_overhead + gen_time