"""Weight plumbing for the ActionRoPE Space.

The upstream ``actionrope.models.wan.load_pipeline`` loads the *base* Wan2.2-TI2V-5B
DiT shards (3 files, ~20 GB) and then overwrites them with the ActionRoPE
checkpoint. The ActionRoPE checkpoint is itself a full DiffSynth-native DiT state
dict (827 tensors: blocks.* / head.* / patch_embedding / text_embedding /
time_embedding / time_projection / arope_mask_embedding), so the base shards are
redundant. This module builds ``diffsynth.models.wan_video_dit.WanModel`` directly
from the registered TI2V-5B config and strict-loads the ActionRoPE checkpoint
(exactly what upstream ``load_checkpoint`` + ``TrainTask.build_module`` do), loads
the Wan2.2 VAE the same way, and returns a minimal pipeline object plus a
``Generator`` whose ``__call__`` is the upstream one — the paper's sampling path,
unchanged.
"""
from __future__ import annotations

import os


def download_space_weights() -> tuple[str, str, str]:
    """Fetch the three runtime weight files (ActionRoPE DiT, its NULL-prompt table, the Wan2.2 VAE).

    Returns absolute (ckpt_path, null_prompt_path, vae_path).
    """
    from huggingface_hub import hf_hub_download

    ckpt_path = hf_hub_download(
        "teawhite/ActionRoPE", "wan2.2-ti2v-5b/actionrope_cns.safetensors", repo_type="model"
    )
    null_prompt_path = hf_hub_download(
        "teawhite/ActionRoPE", "wan2.2-ti2v-5b/umt5_null_prompt.pt", repo_type="model"
    )
    vae_path = hf_hub_download("Wan-AI/Wan2.2-TI2V-5B", "Wan2.2_VAE.pth", repo_type="model")
    return ckpt_path, null_prompt_path, vae_path


class _VAEPipeline:
    """BasePipeline holding only the VAE.

    Subclassing (rather than cherry-picking bound methods) keeps every helper
    ``wan.py`` reaches for — ``preprocess_image``, ``vae_output_to_video`` (which
    internally calls ``vae_output_to_image``), ``torch_dtype``, ``device``.
    """

    def __init__(self, vae, device, torch_dtype):
        from diffsynth.diffusion.base_pipeline import BasePipeline

        class _Pipe(BasePipeline):
            def __init__(self):
                super().__init__(device=device, torch_dtype=torch_dtype)
                self.vae = vae

        self._pipe = _Pipe()
        # wan.py accesses pipe.vae / pipe.device / pipe.torch_dtype / helpers on the object itself
        self.vae = vae
        self.device = device
        self.torch_dtype = torch_dtype

    def __getattr__(self, name):
        # delegate any BasePipeline helper (preprocess_image, vae_output_to_*) to the real pipeline object
        return getattr(self.__dict__["_pipe"], name)

    def encode(self, videos, device, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)):
        return self.vae.encode(videos, device=device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)

    def decode(self, hidden_states, device, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)):
        return self.vae.decode(hidden_states, device=device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)


def build_pipe(ckpt_path: str, vae_path: str, device, torch_dtype) -> tuple:
    """Build (pipeline, dit, meta): DiT from the ActionRoPE ckpt, VAE from Wan2.2_VAE.pth.

    No base-DiT download, no text encoder — the DiT state dict is complete in the
    ActionRoPE checkpoint, and prompts come from the shipped ``umt5_null_prompt.pt``
    embedding table.
    """
    from safetensors import safe_open

    from diffsynth.core.loader.file import load_state_dict
    from diffsynth.models.wan_video_dit import WanModel
    from diffsynth.models.wan_video_vae import WanVideoVAE38
    from diffsynth.utils.state_dict_converters.wan_video_vae import WanVideoVAEStateDictConverter

    from actionrope.checkpoint import read_metadata
    from actionrope.models.wan_dit import install_mask_channel

    # Registered DiffSynth config for Wan-AI/Wan2.2-TI2V-5B
    # (diffsynth/configs/model_configs.py, model_hash 1f5ab7703c6fc803fdded85ff040c316)
    ti2v_5b_kwargs = {
        "has_image_input": False, "patch_size": [1, 2, 2], "in_dim": 48, "dim": 3072, "ffn_dim": 14336,
        "freq_dim": 256, "text_dim": 4096, "out_dim": 48, "num_heads": 24, "num_layers": 30, "eps": 1e-06,
        "seperated_timestep": True, "require_clip_embedding": False, "require_vae_embedding": False,
        "fuse_vae_embedding_in_latents": True,
    }

    meta = read_metadata(ckpt_path)
    sd = {}
    with safe_open(ckpt_path, framework="pt", device="cpu") as f:
        for k in f.keys():
            sd[k] = f.get_tensor(k)
    has_mask = any(k.startswith("arope_mask_embedding.") for k in sd)

    dit = WanModel(**ti2v_5b_kwargs)
    if has_mask:
        install_mask_channel(dit)
    dit.load_state_dict(sd, strict=True)  # raises if the checkpoint is not the full TI2V-5B DiT
    del sd
    dit = dit.to(device=device, dtype=torch_dtype).eval().requires_grad_(False)
    # RoPE tables are plain tensors (not buffers): move them once (mirrors TrainTask.on_prepared)
    dit.freqs = tuple(t.to(device) for t in dit.freqs)

    vae_sd = WanVideoVAEStateDictConverter(load_state_dict(vae_path, torch_dtype=torch_dtype, device="cpu"))
    vae = WanVideoVAE38()
    vae.load_state_dict(vae_sd, strict=True)
    del vae_sd
    vae = vae.to(device=device, dtype=torch_dtype).eval().requires_grad_(False)
    vae.mean = vae.mean.to(device)
    vae.std = vae.std.to(device)

    print(f"[arope_loader] DiT {sum(p.numel() for p in dit.parameters()) / 1e9:.2f} B params, "
          f"mask_channel={has_mask}, metadata={dict(sorted(meta.items()))}")
    return _VAEPipeline(vae, device, torch_dtype), dit, meta


class Generator:
    """Upstream ``actionrope.models.wan.Generator.__call__`` on a directly-built pipe.

    The upstream ``__init__`` calls ``load_pipeline(model_dir)`` which needs the base
    DiT shards; everything else in it is replicated here verbatim.
    """

    def __init__(self, ckpt_path: str, vae_path: str, null_prompt_path: str,
                 device="cuda", torch_dtype=None):
        import torch

        from actionrope.checkpoint import meta_bool, meta_float
        from actionrope.models.wan import SHIFT, PromptEmbeddings
        from actionrope.models.wan_dit import MASK_EMBEDDING

        torch_dtype = torch_dtype or torch.bfloat16
        self.device = device
        self.pipe, self.dit, self.meta = build_pipe(ckpt_path, vae_path, device, torch_dtype)
        self.mask_channel = hasattr(self.dit, MASK_EMBEDDING)
        self.shift = meta_float(self.meta, "cns_shift", SHIFT)
        self.cns = (
            {
                "shift_known": meta_float(self.meta, "cns_shift_known"),
                "alpha": meta_float(self.meta, "cns_alpha", 1.0),
            }
            if meta_bool(self.meta, "cns")
            else None
        )
        self.prompts = PromptEmbeddings("", device, null_prompt_path)

    # The full paper sampling path, verbatim from actionrope.models.wan.Generator.
    from actionrope.models.wan import Generator as _Upstream  # noqa: F401  (readability marker)

    def __call__(self, first_frame, frame_offset_px, prompt=None, seed=0, steps=50, shift=None, cns=None):
        import numpy as np

        from actionrope.models import wan as wan_module
        from actionrope.prompts import NULL_PROMPT

        z0 = wan_module.encode_first_frame(self.pipe, first_frame)
        context = self.prompts(prompt if prompt is not None else NULL_PROMPT)
        offset = wan_module.cell_offsets(frame_offset_px).unsqueeze(0).to(self.device)
        mask = wan_module.known_input(offset) if self.mask_channel else None
        use_cns = self.cns is not None if cns is None else bool(cns)
        if use_cns and self.cns is None:
            raise ValueError("cns=True needs a checkpoint trained with CNS")
        latents = wan_module.sample(
            self.dit, z0, context, offset, mask,
            steps=steps, shift=self.shift if shift is None else shift, seed=seed,
            cns=self.cns if use_cns else None,
        )
        return wan_module.decode(self.pipe, latents), float(wan_module.FPS)