"""
FireRed local engine — mirrors the working parent Space
kulkas2pintu/FireRed-Image-Edit-1.0-Fast.

The pipeline is loaded at MODULE IMPORT (top level), exactly like the parent,
so ZeroGPU's `spaces` library virtualizes the CUDA placement. Loading it
lazily inside the @spaces.GPU worker instead triggers:
    RuntimeError: NVML_SUCCESS == r INTERNAL ASSERT FAILED (CUDACachingAllocator)
which is why the previous version always errored on the FireRed step.
"""
import os
import gc
import random

import numpy as np
import torch
from PIL import Image

import spaces  # type: ignore

MAX_SEED = np.iinfo(np.int32).max
LANCZOS = getattr(Image, "Resampling", Image).LANCZOS

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16

print("[firered_engine] torch:", torch.__version__, "· device:", device)

# NOTE: no prefetch snapshot_download here. from_pretrained (below) pulls only
# the files it actually needs; a full-repo prefetch of the "AIO" Qwen repo
# duplicates/expands storage and blew the Space's 150 GB disk limit.

from diffusers import FlowMatchEulerDiscreteScheduler  # noqa: F401
from firered import (
    QwenImageEditPlusPipeline,
    QwenImageTransformer2DModel,
    QwenDoubleStreamAttnProcessorFA3,
)

# Load once, at import — identical to the parent Space's proven pattern.
pipe = QwenImageEditPlusPipeline.from_pretrained(
    "FireRedTeam/FireRed-Image-Edit-1.1",
    transformer=QwenImageTransformer2DModel.from_pretrained(
        "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19",
        torch_dtype=dtype,
        device_map="cuda",
    ),
    torch_dtype=dtype,
).to(device)

# FP8-quantize the transformer. FireRed AND the resident Wan-Animate pipeline are
# both kept resident, and ZeroGPU packs ALL resident tensors to disk at startup;
# in bf16 that pack overflowed the ~150 GB disk ("No space left on device" in
# torch.pack). FP8 ~halves FireRed's footprint so both fit. The transformer is
# already on-device (device_map="cuda"), so quantize in place — no later .to()
# (torchao Float8Tensor can't be device-transferred). Non-fatal if unavailable.
try:
    from torchao.quantization import (
        quantize_ as _q8, Float8DynamicActivationFloat8WeightConfig as _F8,
    )
    _q8(pipe.transformer, _F8())
    print("[firered_engine] FP8 quantized transformer.")
except Exception as _q8e:
    print(f"[firered_engine] FP8 quantization skipped: {_q8e}")

# NOTE: the FA3 attention processor is intentionally NOT enabled here.
# In this combined Space flash-attn IS installed (for Wan), so the FA3 kernels
# would actually run — and on the ZeroGPU Blackwell GPU they hard-kill the
# worker ("GPU task aborted", no Python traceback). The parent FireRed Space
# has no flash-attn, so it silently falls back to default attention and works;
# we do the same explicitly here for reliability.
# (QwenDoubleStreamAttnProcessorFA3 imported above is left unused on purpose.)


def _resize_to_qwen(image_pil: Image.Image, max_dim: int = 1024):
    """Round to multiples of 8, cap at 1024 on the longest side."""
    w, h = image_pil.size
    if w > h:
        nw = max_dim
        nh = int(nw * h / w)
    else:
        nh = max_dim
        nw = int(nh * w / h)
    return (nw // 8) * 8, (nh // 8) * 8


# duration=70: the model is RESIDENT (loaded at boot), so a single-image edit at
# 1024px / 4 steps only takes ~30-50s of actual GPU time. The old 200 (→400s
# after the xlarge ×2) both wasted quota and exceeded ZeroGPU's per-call max,
# so every run was rejected with "requested GPU duration (400s) is larger than
# the maximum allowed". 70 → ~140s reserved, comfortably under the cap.
@spaces.GPU(duration=70, size='xlarge')
def firered_edit(image_pil, prompt, seed=0, randomize_seed=True,
                 guidance_scale=1.0, steps=4):
    """
    Edit an image with FireRed/Qwen-Image-Edit-Rapid.
    Returns (result_pil, seed_used).
    """
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    if image_pil is None:
        raise ValueError("image is required.")
    if not prompt or not str(prompt).strip():
        raise ValueError("prompt is required.")

    if image_pil.mode != "RGB":
        if image_pil.mode in ("RGBA", "LA"):
            bg = Image.new("RGB", image_pil.size, (255, 255, 255))
            mask = image_pil.split()[-1] if "A" in image_pil.mode else None
            bg.paste(image_pil, mask=mask)
            image_pil = bg
        else:
            image_pil = image_pil.convert("RGB")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(int(seed))

    width, height = _resize_to_qwen(image_pil, max_dim=1024)
    negative_prompt = (
        "worst quality, low quality, bad anatomy, bad hands, text, error, "
        "missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, "
        "signature, watermark, username, blurry"
    )

    try:
        result_image = pipe(
            image=[image_pil],
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            num_inference_steps=int(steps),
            generator=generator,
            true_cfg_scale=float(guidance_scale),
        ).images[0]
    finally:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    return result_image, int(seed)


# ─── Video (frame-by-frame) clothing edit ───────────────────────────────────
_NEG = (
    "worst quality, low quality, bad anatomy, bad hands, text, error, "
    "missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, "
    "signature, watermark, username, blurry"
)
_ANCHOR = (
    " Keep the new clothing identical in color, texture, pattern and style to "
    "the reference image. Keep the person's face, hair, body, pose and the "
    "background completely unchanged."
)


def _edit_video_duration(frames_dir, out_dir, prompt, n_frames, seed=0,
                         steps=4, guidance_scale=1.0, multiplier=1, progress=None):
    # Model is RESIDENT + FP8, so no packing/reload cost here — just ~7 s/keyframe
    # edit + fast RIFE. Kept well under ZeroGPU's per-call ceiling: with xlarge the
    # effective limit is ~290 (duration ×2), so duration must stay ≲145. At ~10
    # keyframes this lands ~130 (→ ~260 reserved). The app caps keyframes to match.
    return int(min(140, 40 + int(n_frames) * 9))


# size='xlarge' is REQUIRED — but NOT because of FireRed alone. `spaces` packs
# ALL resident CUDA tensors process-wide into ONE GLOBAL pack (zero/torch/
# patching.py: cuda_aliases + tensor_packs are module globals) and worker_init
# moves the WHOLE pack into EVERY worker, for every decorated function.
# Post-FP8 FireRed is ~37.3 GB and WOULD fit `large`; FireRed + the resident Wan
# pipeline is ~62 GB (57.7 GiB) vs a ~47.5 GiB `large` partition and HF's own
# xlarge_threshold of 42.0 GiB. Forcing size='large' does not fail at schedule
# time — it OOMs inside worker_init and surfaces as a bare "GPU task aborted".
# (The old "~58 GB" figure was FireRed's PRE-FP8 size, so the conclusion was
# right for the wrong reason. De-residenting FireRed does NOT help: its worker
# still force-moves Wan's pack and cannot free it.)
@spaces.GPU(duration=_edit_video_duration, size='xlarge')
def firered_edit_video(frames_dir, out_dir, prompt, n_frames, seed=0,
                       steps=4, guidance_scale=1.0, multiplier=1, progress=None):
    """
    Edit the clothing on every extracted keyframe in `frames_dir` with the same
    prompt, RIFE-interpolate the result `multiplier`x for smooth motion, and
    write the final frames (frame_00000.png …) to `out_dir`.

    Frame 0 is edited from the prompt alone and becomes the appearance
    "anchor"; every later frame is edited with [frame, anchor] so the new outfit
    stays temporally consistent. All edits are forced to the first frame's size
    so they reassemble/interpolate cleanly. Returns `out_dir`.
    """
    import glob as _glob

    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    os.makedirs(out_dir, exist_ok=True)
    frame_paths = sorted(_glob.glob(os.path.join(frames_dir, "*.png")))
    out_size = None
    edited = []
    # Appended to every prompt so Qwen keeps the original composition instead of
    # zooming/recomposing (which was destroying the framing on close-up video).
    # Prompt for the FIRST keyframe (edited from the prompt alone → the anchor).
    frame_prompt = (prompt.rstrip(". ") +
                    ". Keep the exact same camera framing, zoom level, crop and "
                    "composition as the input image; do not zoom in, and keep "
                    "the background unchanged.")
    # Prompt for EVERY LATER keyframe: dress the person in this frame (image 1) in
    # the SAME outfit as the anchor (image 2), so the clothing is identical across
    # frames. Editing frames independently (the old behaviour) made every frame
    # come out differently — only the first looked right; this anchors them.
    anchored_prompt = (
        prompt.rstrip(". ") +
        ". Put the person in the first image into the exact same outfit as the "
        "person in the second image — identical color, texture, pattern and "
        "style. Keep the first image's face, hair, body, pose, camera framing and "
        "background completely unchanged; do not zoom or recompose.")

    anchor_pil = None
    for fp in frame_paths:
        img = Image.open(fp).convert("RGB")
        # Full resolution (1024 px). Qwen needs ~1 MP to preserve composition;
        # at low res it recomposes/zooms, which broke the framing.
        width, height = _resize_to_qwen(img, max_dim=1024)
        generator = torch.Generator(device=device).manual_seed(int(seed))
        if anchor_pil is None:
            # First keyframe → prompt only; its result becomes the outfit anchor.
            res = pipe(
                image=[img],
                prompt=frame_prompt,
                negative_prompt=_NEG,
                height=height, width=width,
                num_inference_steps=int(steps),
                generator=generator,
                true_cfg_scale=float(guidance_scale),
            ).images[0]
            anchor_pil = res
        else:
            # Later keyframes → [this frame, anchor] so the outfit matches frame 1.
            res = pipe(
                image=[img, anchor_pil],
                prompt=anchored_prompt,
                negative_prompt=_NEG,
                height=height, width=width,
                num_inference_steps=int(steps),
                generator=generator,
                true_cfg_scale=float(guidance_scale),
            ).images[0]
        if out_size is None:
            out_size = res.size
        elif res.size != out_size:
            res = res.resize(out_size, LANCZOS)
        edited.append(np.asarray(res).astype(np.float32) / 255.0)

    # RIFE-interpolate the edited keyframes up to a smoother fps (fast on GPU).
    if int(multiplier) > 1 and len(edited) >= 2:
        try:
            import rife_interp
            frames_out = rife_interp.interpolate_bits(
                np.stack(edited, axis=0), multiplier=int(multiplier))
        except Exception as e:
            print(f"[firered] RIFE interpolation failed, keeping keyframes: {e}")
            frames_out = edited
    else:
        frames_out = edited

    for i, f in enumerate(frames_out):
        arr = (np.clip(f, 0.0, 1.0) * 255.0).astype(np.uint8)
        Image.fromarray(arr).save(os.path.join(out_dir, f"frame_{i:05d}.png"))

    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return out_dir