import os, gc, sys, json, warnings, traceback, tempfile, time
from pathlib import Path

import torch
import gradio as gr
from PIL import Image

warnings.filterwarnings("ignore")
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

try:
    import psutil
    _HAS_PSUTIL = True
except ImportError:
    _HAS_PSUTIL = False


MODEL_REPO  = "rockapaper/wan2.2_14b_i2v_480p_lightning_nsfw_diffusers_Multiscene_SVI_sdnq_fp8_svd_r64"
EULER_REPO  = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
HF_TOKEN    = os.environ.get("HF_TOKEN", None)
CACHE_ROOT  = "/data" if os.path.isdir("/data") else os.path.join(os.path.dirname(os.path.abspath(__file__)), "cache")
MODEL_DIR   = os.path.join(CACHE_ROOT, "model")
OFFLOAD_DIR = os.path.join(CACHE_ROOT, "offload")
OUTPUT_DIR  = os.path.join(CACHE_ROOT, "outputs")
GRADIO_TMP  = os.path.join(CACHE_ROOT, "gradio_tmp")

os.environ["GRADIO_TEMP_DIR"] = GRADIO_TMP

for _d in [MODEL_DIR, OFFLOAD_DIR, OUTPUT_DIR, GRADIO_TMP]:
    os.makedirs(_d, exist_ok=True)


# ── Reduced defaults for 16 GB RAM ──────────────────────────────────────────
DEFAULTS = dict(
    num_frames          = 81,
    fps                 = 16,
    guidance_scale      = 1.0,
    num_inference_steps = 4,
    height              = 480,
    width               = 832,
    seed                = 42,
    negative_prompt     = (
        "blurry, low quality, deformed, bad anatomy, extra limbs, watermark, "
        "text, censored, mosaic, static, frozen, flickering, artifacts, ugly"
    ),
)

TARGET_DTYPE = torch.bfloat16
_pipe        = None

# ── Global transformer swap state: only ONE transformer in RAM at a time ─────
_SWAP = {"sf": None, "model": None}


class ExtremeTeaCache:
    def __init__(self, rel_l1_thresh: float = 0.15, start_percent: float = 0.1):
        self.rel_l1_thresh      = rel_l1_thresh
        self.start_percent      = start_percent
        self.cache              = None
        self.prev_hidden_states = None
        self.step               = 0
        self.total_steps        = 0
        self.skipped            = 0

    def inject(self, pipe, total_steps: int):
        self.total_steps = total_steps
        original_fwd = pipe.transformer.forward

        def _patched_forward(*args, **kwargs):
            hs  = kwargs.get("hidden_states") or (args[0] if args else None)
            pct = self.step / max(1, self.total_steps)
            if (
                pct >= self.start_percent
                and self.cache is not None
                and self.prev_hidden_states is not None
                and hs is not None
            ):
                diff   = hs - self.prev_hidden_states
                rel_l1 = diff.abs().mean() / (self.prev_hidden_states.abs().mean() + 1e-8)
                if rel_l1 < self.rel_l1_thresh:
                    self.skipped += 1
                    self.step    += 1
                    return self.cache
            out = original_fwd(*args, **kwargs)
            self.cache              = out
            self.prev_hidden_states = hs.clone() if hs is not None else None
            self.step              += 1
            if self.step % 2 == 0:
                gc.collect()
            return out

        pipe.transformer.forward = _patched_forward
        print("[TeaCache] ✅ Injected")
        return pipe


def ram_str() -> str:
    if not _HAS_PSUTIL:
        return ""
    vm   = psutil.virtual_memory()
    used = vm.used      / 1024**3
    free = vm.available / 1024**3
    return f"  [RAM: {used:.1f}GB used | {free:.1f}GB free]"


def ensure_model_downloaded(progress_cb=None) -> str:
    from huggingface_hub import snapshot_download
    marker = os.path.join(MODEL_DIR, "model_index.json")
    if os.path.exists(marker):
        print(f"[Model] ✅ Already cached: {MODEL_DIR}")
        return MODEL_DIR
    print(f"[Model] 🔽 Downloading (~40GB) → {MODEL_DIR}...")
    if progress_cb:
        progress_cb(0, "Downloading model (~40GB)...")
    snapshot_download(
        repo_id         = MODEL_REPO,
        local_dir       = MODEL_DIR,
        token           = HF_TOKEN,
        ignore_patterns = ["*.msgpack", "*.h5", "flax_model*"],
    )
    print(f"[Model] ✅ Done")
    return MODEL_DIR


# ─────────────────────────────────────────────────────────────────────────────
# LAZY SWAP TRANSFORMER
# Only one WanTransformer3DModel lives in RAM at a time.
# When the pipeline calls transformer_2 after having used transformer,
# _LazyTransformer automatically unloads the previous one first.
# ─────────────────────────────────────────────────────────────────────────────
class _LazyTransformer:
    """
    Transparent proxy for WanTransformer3DModel.
    - Loads its weights on first forward call (or explicit _activate()).
    - Unloads the *other* transformer from _SWAP before loading itself.
    - Uses attn_implementation="sdpa" for memory-efficient attention.
    """

    def __init__(self, model_dir: str, subfolder: str, dtype):
        # Use object.__setattr__ so our __setattr__ doesn't recurse
        object.__setattr__(self, "_dir", model_dir)
        object.__setattr__(self, "_sf",  subfolder)
        object.__setattr__(self, "_dt",  dtype)
        object.__setattr__(self, "_cfg", None)

    # ── load on demand, evict sibling ────────────────────────────────────────
    def _activate(self):
        sf = object.__getattribute__(self, "_sf")
        if _SWAP["sf"] == sf:
            return _SWAP["model"]

        # Evict current resident
        if _SWAP["model"] is not None:
            print(f"[Swap] ⬇ Evicting {_SWAP['sf']}…{ram_str()}")
            del _SWAP["model"]
            _SWAP["model"] = None
            gc.collect()

        # Load self with SDPA memory-efficient attention
        md = object.__getattribute__(self, "_dir")
        dt = object.__getattribute__(self, "_dt")
        from diffusers import WanTransformer3DModel
        print(f"[Swap] ⬆ Loading {sf}…{ram_str()}")
        m = WanTransformer3DModel.from_pretrained(
            md, subfolder=sf,
            torch_dtype             = dt,
            low_cpu_mem_usage       = True,
            ignore_mismatched_sizes = True,
            attn_implementation     = "sdpa",   # ← memory-efficient SDPA
        )
        m.eval()
        gc.collect()
        _SWAP["sf"]    = sf
        _SWAP["model"] = m
        print(f"[Swap] ✅ {sf} active{ram_str()}")
        return m

    # Pre-warm: load without waiting for the pipeline to call forward
    def prefetch(self):
        self._activate()

    # ── nn.Module interface expected by the pipeline ─────────────────────────
    def __call__(self, *a, **kw):
        return self._activate()(*a, **kw)

    def forward(self, *a, **kw):
        return self._activate().forward(*a, **kw)

    def eval(self):
        return self

    def train(self, mode=True):
        return self

    def to(self, *a, **kw):
        return self   # weights stay on CPU; no-op

    def parameters(self):
        # Avoid loading model just to iterate params (e.g. for dtype checks)
        return iter([])

    def named_parameters(self, *a, **kw):
        return iter([])

    # ── attribute passthrough ────────────────────────────────────────────────
    @property
    def config(self):
        cfg = object.__getattribute__(self, "_cfg")
        if cfg is None:
            from diffusers import WanTransformer3DModel
            from diffusers.configuration_utils import FrozenDict
            raw = WanTransformer3DModel.load_config(
                object.__getattribute__(self, "_dir"),
                subfolder=object.__getattribute__(self, "_sf"),
            )
            # Pipeline accesses config via attribute (config.patch_size),
            # but load_config() returns a plain dict → wrap in FrozenDict.
            cfg = FrozenDict(raw)
            object.__setattr__(self, "_cfg", cfg)
        return cfg

    @property
    def dtype(self):
        return object.__getattribute__(self, "_dt")

    @property
    def device(self):
        return torch.device("cpu")

    @property
    def training(self):
        return False

    def __getattr__(self, name):
        # Private / already-defined attrs raise normally
        if name.startswith("_"):
            raise AttributeError(name)
        # Passthrough to loaded model
        return getattr(self._activate(), name)

    def __setattr__(self, name, value):
        if name.startswith("_"):
            object.__setattr__(self, name, value)
        else:
            setattr(self._activate(), name, value)


def encode_text(model_dir, prompt, negative_prompt, max_length=226):
    from transformers import UMT5EncoderModel, AutoTokenizer

    # Register SDNQ BEFORE loading the text encoder
    try:
        from sdnq import SDNQConfig  # noqa: F401
        print("[SDNQ] ✅ Registered in encode_text")
    except ImportError:
        import subprocess
        subprocess.run([sys.executable, "-m", "pip", "install", "sdnq", "-q"], check=True)
        from sdnq import SDNQConfig  # noqa: F401

    print(f"\n[UMT5] Loading text encoder...{ram_str()}")
    tokenizer    = AutoTokenizer.from_pretrained(model_dir, subfolder="tokenizer", use_fast=True)
    text_encoder = UMT5EncoderModel.from_pretrained(
        model_dir, subfolder="text_encoder",
        torch_dtype=TARGET_DTYPE, low_cpu_mem_usage=True,
        ignore_mismatched_sizes=True,
    )
    text_encoder.eval()
    print(f"[UMT5] ✅ Loaded{ram_str()}")

    with torch.inference_mode():
        def enc(text):
            ids = tokenizer(text, padding="max_length", max_length=max_length,
                            truncation=True, return_tensors="pt")
            return text_encoder(ids.input_ids, attention_mask=ids.attention_mask)[0].cpu()
        p_emb = enc(prompt)
        n_emb = enc(negative_prompt)

    del text_encoder, tokenizer
    gc.collect()
    print(f"[UMT5] ✅ Encoded & unloaded{ram_str()}")
    return p_emb, n_emb


def build_pipeline(model_dir, progress_cb=None):
    try:
        from sdnq import SDNQConfig  # noqa: F401
        print("[SDNQ] ✅ Registered")
    except ImportError:
        import subprocess
        subprocess.run([sys.executable, "-m", "pip", "install", "sdnq", "-q"], check=True)
        from sdnq import SDNQConfig  # noqa: F401

    from diffusers import WanImageToVideoPipeline, AutoencoderKLWan, FlowMatchEulerDiscreteScheduler

    # ── ① Lazy transformer proxies (NO weights loaded yet) ───────────────────
    if progress_cb: progress_cb(0.10, "① Creating lazy transformer proxies…")
    transformer   = _LazyTransformer(model_dir, "transformer",   TARGET_DTYPE)
    transformer_2 = _LazyTransformer(model_dir, "transformer_2", TARGET_DTYPE)
    # Pre-load T1 now (handles ~90% of denoising steps)
    if progress_cb: progress_cb(0.15, "① Loading transformer (High Noise)…")
    transformer.prefetch()

    # ── ② VAE (CPU, fp16, tiling+slicing) ───────────────────────────────────
    if progress_cb: progress_cb(0.55, "② Loading VAE…")
    print(f"\n[VAE] Loading...{ram_str()}")
    vae = AutoencoderKLWan.from_pretrained(
        model_dir, subfolder="vae",
        torch_dtype=TARGET_DTYPE, low_cpu_mem_usage=True,
    )
    vae.eval()
    vae.enable_tiling()
    vae.enable_slicing()
    gc.collect()
    print(f"[VAE] ✅ OK{ram_str()}")

    # ── ③ Scheduler ──────────────────────────────────────────────────────────
    if progress_cb: progress_cb(0.70, "③ Loading Scheduler…")
    euler_cache = os.path.join(CACHE_ROOT, "euler_scheduler")
    try:
        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
            EULER_REPO, subfolder="scheduler", cache_dir=euler_cache, token=HF_TOKEN,
        )
        print("[Scheduler] ✅ OK")
    except Exception as e:
        print(f"[Scheduler] fallback default ({e})")
        scheduler = FlowMatchEulerDiscreteScheduler()

    # ── ④ Pipeline ───────────────────────────────────────────────────────────
    if progress_cb: progress_cb(0.80, "④ Assembling pipeline…")
    pipe = WanImageToVideoPipeline(
        tokenizer=None, text_encoder=None,
        image_encoder=None, image_processor=None,
        transformer=transformer, transformer_2=transformer_2,
        vae=vae, scheduler=scheduler,
    )
    pipe.boundary_ratio   = 0.9
    pipe.expand_timesteps = False

    # ── ⑤ Attention slicing: process one head at a time (min VRAM/RAM peak) ──
    try:
        pipe.enable_attention_slicing(1)
        print("[AttnSlice] ✅ slice_size=1")
    except Exception as e:
        print(f"[AttnSlice] skip ({e})")

    gc.collect()

    # ── ⑥ TeaCache ───────────────────────────────────────────────────────────
    if progress_cb: progress_cb(0.92, "⑤ Injecting TeaCache…")
    num_steps = DEFAULTS["num_inference_steps"]
    if num_steps > 2:
        thresh = 0.15 if num_steps <= 6 else 0.12
        tc = ExtremeTeaCache(rel_l1_thresh=thresh, start_percent=0.1)
        pipe = tc.inject(pipe, total_steps=num_steps)
        pipe._teacache = tc
        print(f"[TeaCache] ✅ threshold={thresh}")
    else:
        pipe._teacache = None

    if progress_cb: progress_cb(1.0, "✅ Pipeline ready!")
    print("\n[Pipeline] 🎉 Ready!\n")
    return pipe


def get_pipeline(progress_cb=None):
    global _pipe
    if _pipe is not None:
        return _pipe
    print("=" * 60)
    print("  First load – might take 30-120 min on CPU, grab some tea")
    print("=" * 60)
    model_dir = ensure_model_downloaded(progress_cb)
    _pipe     = build_pipeline(model_dir, progress_cb)
    return _pipe


def generate(
    image_input, prompt, negative_prompt,
    num_frames, num_steps, guidance_scale,
    height, width, seed,
    progress=gr.Progress(track_tqdm=True),
):
    from diffusers.utils import export_to_video

    if image_input is None:
        return None, "❌ Upload an image first!"
    if not prompt.strip():
        return None, "❌ Write a prompt!"

    progress(0.0, "Standby...")
    t_start = time.time()
    logs    = []
    def log(m): logs.append(m); print(m)

    log(f"🎬 {num_frames} frames × {num_steps} steps")
    log(f"   Prompt: {prompt[:80]}{'...' if len(prompt)>80 else ''}")

    # Step 1 – encode text
    try:
        progress(0.05, "① Encoding text (UMT5)...")
        log(f"\n[Step 1] Encode text{ram_str()}")
        model_dir    = ensure_model_downloaded()
        p_emb, n_emb = encode_text(model_dir, prompt, negative_prompt)
        log(f"   p_emb: {p_emb.shape}  n_emb: {n_emb.shape}")
    except Exception as e:
        traceback.print_exc(); return None, f"❌ Text encode failed: {e}"

    # Step 2 – pipeline
    try:
        progress(0.10, "② Loading pipeline (first time = long wait)...")
        log(f"\n[Step 2] Get pipeline{ram_str()}")
        pipe = get_pipeline(progress_cb=lambda v, m: progress(0.10 + v*0.35, m))
    except Exception as e:
        traceback.print_exc(); return None, f"❌ Pipeline failed: {e}"

    if hasattr(pipe, "_teacache") and pipe._teacache:
        tc = pipe._teacache
        tc.cache = tc.prev_hidden_states = None
        tc.step = tc.skipped = 0
        tc.total_steps = num_steps
    pipe.scheduler.config.num_inference_steps = num_steps

    # Step 3 – prepare image
    progress(0.47, "③ Preparing image...")
    # Gradio 6.x can pass image as a dict {"path": "...", "url": "..."} or str path
    if isinstance(image_input, dict):
        image_input = image_input.get("path") or image_input.get("url")
    if isinstance(image_input, str):
        image_input = Image.open(image_input)
    if not isinstance(image_input, Image.Image):
        image_input = Image.fromarray(image_input)
    img = image_input.convert("RGB").resize((width, height), Image.Resampling.LANCZOS)
    log(f"\n[Step 3] Resized → {width}×{height}")

    # Step 4 – render
    generator = torch.Generator(device="cpu").manual_seed(seed)
    progress(0.50, f"④ Rendering {num_frames}f × {num_steps}s – go get some coffee...")
    log(f"\n[Step 4] Rendering...{ram_str()}")
    log(f"   ⚠️  CPU + swap: transformer swaps once at boundary_ratio=0.9")

    try:
        with torch.inference_mode():
            result = pipe(
                image=img,
                prompt_embeds=p_emb, negative_prompt_embeds=n_emb,
                num_frames=num_frames, height=height, width=width,
                guidance_scale=guidance_scale, num_inference_steps=num_steps,
                generator=generator,
            )
        frames = result.frames[0]
    except MemoryError as e:
        gc.collect(); return None, f"❌ OOM: {e}"
    except Exception as e:
        traceback.print_exc(); return None, f"❌ Render failed: {e}"

    # Step 5 – export
    progress(0.96, "⑤ Exporting MP4...")
    output_path = os.path.join(OUTPUT_DIR, f"wan22_{int(time.time())}.mp4")
    try:
        export_to_video(frames, output_path, fps=DEFAULTS["fps"])
    except Exception as e:
        traceback.print_exc(); return None, f"❌ Export failed: {e}"

    elapsed = time.time() - t_start
    tc_info = ""
    if hasattr(pipe, "_teacache") and pipe._teacache:
        tc = pipe._teacache
        tc_info = f"\n   ⚡ TeaCache skipped: {tc.skipped}/{num_steps} ({100*tc.skipped/max(1,num_steps):.0f}%)"

    summary = f"✅ Done!\n   Time: {elapsed/3600:.1f}h\n   File: {output_path}{tc_info}"
    log(summary)
    progress(1.0, "✅ Done!")
    gc.collect()
    return output_path, "\n".join(logs)


CSS = """
#title-row { text-align: center; margin-bottom: 10px; }
#title-row h1 { font-size: 1.6em; font-weight: bold; }
.log-box textarea { font-family: monospace; font-size: 12px; }
"""

DESCRIPTION = """
**WAN 2.2 14B I2V** · Lightning + SVI · SDNQ FP8+SVD · CPU  
Upload image + write prompt → Generate video. Lightning LoRA baked in.
"""

def build_ui():
    with gr.Blocks(css=CSS, title="WAN 2.2 I2V – CPU") as demo:
        with gr.Row(elem_id="title-row"):
            with gr.Column():
                gr.HTML("<h1>🎬 WAN 2.2 14B I2V · SDNQ FP8+SVD</h1>")
                gr.Markdown(DESCRIPTION)

        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="🖼️ Reference Image", type="pil",
                    sources=["upload", "clipboard"],
                )
                prompt = gr.Textbox(
                    label="✏️ Prompt", lines=3,
                    value="The woman walks forward slowly, camera follows smoothly.",
                )
                negative_prompt = gr.Textbox(
                    label="🚫 Negative Prompt", lines=2,
                    value=DEFAULTS["negative_prompt"],
                )
                generate_btn = gr.Button("🎬 Generate Video", variant="primary", size="lg")

                with gr.Accordion("⚙️ Advanced", open=False):
                    with gr.Row():
                        num_frames = gr.Slider(label="Frames", minimum=17, maximum=121, step=8, value=DEFAULTS["num_frames"])
                        num_steps  = gr.Slider(label="Steps",  minimum=1,  maximum=20,  step=1, value=DEFAULTS["num_inference_steps"], info="Lightning: 4 steps optimal")
                    with gr.Row():
                        guidance = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=8.0, step=0.5, value=DEFAULTS["guidance_scale"], info="1.0 for distilled model")
                        seed     = gr.Number(label="Seed", value=DEFAULTS["seed"], precision=0)
                    with gr.Row():
                        height = gr.Dropdown(label="Height (px)", choices=[480, 720], value=DEFAULTS["height"])
                        width  = gr.Dropdown(label="Width (px)",  choices=[832, 1280], value=DEFAULTS["width"])
                    gr.Markdown("""
> **CPU Free Tier:** 4 steps × 81 frames ≈ 2-8h · TeaCache speeds up ~30-50%
                    """)

            with gr.Column(scale=1):
                video_output = gr.Video(label="🎥 Output", autoplay=True)
                log_output   = gr.Textbox(label="📋 Log", lines=15, interactive=False, elem_classes=["log-box"])

        gr.Examples(
            examples=[
                [None, "The woman walks forward slowly, sunlight on her face, cinematic.", DEFAULTS["negative_prompt"]],
                [None, "Close-up of hands assembling a mechanical watch, shallow depth of field.", DEFAULTS["negative_prompt"]],
                [None, "Ocean waves crash on a rocky shore, slow motion, dramatic lighting.", DEFAULTS["negative_prompt"]],
            ],
            inputs=[image_input, prompt, negative_prompt],
            label="💡 Example prompts",
        )

        generate_btn.click(
            fn=generate,
            inputs=[image_input, prompt, negative_prompt, num_frames, num_steps, guidance, height, width, seed],
            outputs=[video_output, log_output],
        )

    return demo

if __name__ == "__main__":
    demo = build_ui()
    demo.queue(max_size=2).launch(
        server_name = "0.0.0.0",
        server_port = 7860,
        show_error  = True,
        share       = True,
    )