# SPDX-License-Identifier: Apache-2.0
#
# SANA-WM STREAMING overlay materializer.
#
# Translates the NVlabs SANA-WM_streaming HF layout (self-forcing DMD
# checkpoint) into a Diffusers-style directory tree consumable by SGLang's
# ComposedPipelineBase (SanaWMTwoStagePipeline / SanaWMRealtimePipeline).
#
# Source layout (Efficient-Large-Model/SANA-WM_streaming):
#   sana_dit/model.pt                      (DMD self-forcing TRAINING ckpt:
#                                           DiT weights under `generator` with
#                                           a wrapping prefix, bf16, 872 tensors)
#   ltx2_causal_vae/{config.json, diffusion_pytorch_model.safetensors}
#   refiner_diffusers/{transformer,connectors}/...
#   gemma3_12b/...                          (refiner text encoder + tokenizer)
#
# Output layout (Diffusers-style, written to ``output_dir``):
#   model_index.json                        (copied from overlay root by the runner)
#   transformer/{config.json, diffusion_pytorch_model.safetensors}
#                                           (CONVERTED from sana_dit/model.pt)
#   vae/...                                 (linked: ltx2 CAUSAL vae)
#   text_encoder/..., tokenizer/...         (downloaded from google/gemma-2-2b-it)
#   scheduler/scheduler_config.json
#   refiner/{transformer,connectors,text_encoder}/...

from __future__ import annotations

import json
import os
from typing import Any

from huggingface_hub import snapshot_download
from sglang.multimodal_gen.runtime.utils.model_overlay import (
    _copytree_link_or_copy,
    _ensure_dir,
    _link_or_copy_file,
)

# Stage-1 text encoder (caption_channels=2304 == Gemma-2-2b hidden size; NOT
# interchangeable with the refiner's Gemma3). Users must accept the Gemma
# license once on HuggingFace before materialization runs.
GEMMA2_REPO_ID = "google/gemma-2-2b-it"
GEMMA2_REVISION = "299a8560bedf22ed1c72a8a11e7dce4a7f9f51f8"

# Prefix candidates for the DMD checkpoint's DiT weights (under `generator`).
_DMD_PREFIXES = ("model.", "", "module.", "model.module.", "_orig_mod.", "model._orig_mod.")

# Opt-in disk reclaim (set to "1"): delete sana_dit/model.pt from the HF cache
# once materialization succeeds. The ~32 GB DMD training ckpt (generator +
# critic + both optimizer states) is dead weight after conversion — everything
# else in the source snapshot stays hardlinked into the materialized dir, but
# the ckpt is only ever read here. OFF by default; know the tradeoffs:
#   - A future RE-materialization (changed overlay) will fail the source
#     completeness check and force-redownload the FULL source repo (~102 GB,
#     not just model.pt); with download disabled it fails hard instead.
#   - Eviction only acts on a classic HF-cache layout (snapshot symlink whose
#     blob lives under the same repo's blobs/ and is referenced by no other
#     snapshot). User-managed local source dirs are never touched.
_EVICT_DMD_CKPT_ENV = "SANAWM_OVERLAY_EVICT_DMD_CKPT"


def _write_json(path: str, payload: dict[str, Any]) -> None:
    _ensure_dir(os.path.dirname(path))
    with open(path, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, sort_keys=True)
        f.write("\n")


def _transformer_config() -> dict[str, Any]:
    # NVlabs SanaMSVideoCamCtrl_1600M_P1_D20 (depth=20, hidden=2240, heads=20,
    # linear_head_dim=112) — the streaming DMD ckpt shares the bidirectional
    # architecture; the repo ships no config.yaml, so these are the (verified)
    # 1600M values.
    return {
        "_class_name": "SanaWMTransformer3DModel",
        "_diffusers_version": "0.37.0.dev0",
        "patch_size": 1,
        "patch_size_t": 1,
        "in_channels": 128,
        "out_channels": 128,
        "num_layers": 20,
        "num_attention_heads": 20,
        "attention_head_dim": 112,
        "linear_head_dim": 112,
        "num_cross_attention_heads": 20,
        "cross_attention_head_dim": 112,
        "cross_attention_dim": 2240,
        "cross_norm": True,
        "caption_channels": 2304,
        "model_max_length": 300,
        "mlp_ratio": 3.0,
        "qk_norm": True,
        "softmax_every_n": 4,
        "conv_kernel_size": 4,
        "k_conv_only": True,
        "cam_attn_compress": 1,
        "init_cam_from_base": True,
        "use_chunk_plucker_post_attn": True,
        "chunk_plucker_channels": 48,
        "chunk_plucker_post_attn_blocks": 20,
        "chunk_split_strategy": "first_chunk_plus_one",
        "ffn_type": "GLUMBConvTemp",
        "t_kernel_size": 3,
        "mlp_acts": ["silu", "silu", None],
        "pos_embed_type": "wan_rope",
        "vae_temporal_stride": 8,
        "vae_spatial_stride": 32,
    }


def _convert_dmd_dit(
    *, source_dir: str, overlay_dir: str, output_dir: str
) -> None:
    """sana_dit/model.pt (DMD training ckpt) -> transformer safetensors.

    The DiT weights live under ``generator`` with a wrapping prefix (e.g.
    ``model.``). The wrapping is auto-detected by matching the stripped keys
    against the reference key list shipped with the overlay; an exact 872/872
    match is required (silent partial loads are how checkpoints rot)."""
    import torch
    from safetensors.torch import save_file

    ref_path = os.path.join(overlay_dir, "_overlay", "dit_key_list.txt")
    with open(ref_path, encoding="utf-8") as f:
        ref_keys = {ln.strip() for ln in f if ln.strip()}

    pt_path = os.path.join(source_dir, "sana_dit", "model.pt")
    if not os.path.isfile(pt_path):
        raise FileNotFoundError(f"SANA-WM streaming DiT not found at {pt_path}.")
    # mmap keeps the ~32 GB DMD training ckpt (generator + critic + both
    # optimizer states) on disk; only the ~3.2 GB of generator tensors that
    # save_file() reads are ever paged in (and those pages are reclaimable
    # cache, not anonymous memory). Fall back for legacy non-zipfile
    # serialization (mmap unsupported there) and for old torch (no kwarg).
    try:
        ckpt = torch.load(pt_path, map_location="cpu", weights_only=False, mmap=True)
    except (RuntimeError, TypeError, ValueError):
        ckpt = torch.load(pt_path, map_location="cpu", weights_only=False)
    state = ckpt.get("generator", ckpt)
    if not isinstance(state, dict):
        state = state.state_dict()

    chosen = None
    for prefix in _DMD_PREFIXES:
        mapped = {
            k[len(prefix):]: v
            for k, v in state.items()
            if k.startswith(prefix) and k[len(prefix):] in ref_keys
        }
        if len(mapped) == len(ref_keys):
            chosen = (prefix, mapped)
            break
    if chosen is None:
        best = max(
            _DMD_PREFIXES,
            key=lambda p: sum(1 for k in state if k.startswith(p) and k[len(p):] in ref_keys),
        )
        raise RuntimeError(
            f"DMD checkpoint key mapping failed: no prefix maps all {len(ref_keys)} "
            f"reference tensors (best candidate {best!r}). The checkpoint layout "
            "changed; refusing a partial conversion."
        )
    prefix, mapped = chosen
    out = os.path.join(output_dir, "transformer")
    _ensure_dir(out)
    _write_json(os.path.join(out, "config.json"), _transformer_config())
    save_file(
        {k: v.contiguous() for k, v in mapped.items()},
        os.path.join(out, "diffusion_pytorch_model.safetensors"),
    )


def _maybe_evict_dmd_ckpt(source_dir: str) -> None:
    """Opt-in: remove the converted-and-now-unneeded DMD ckpt (link + blob).

    Called as the LAST materialization step — evicting any earlier would
    strand a half-materialized run with no source to retry from. (A residual
    window remains: the framework commits the output dir — marker write +
    rename — only after this returns. Tiny, and recoverable via the
    force-redownload path; documented at _EVICT_DMD_CKPT_ENV.)

    Deliberately conservative: acts only on a classic HF-cache snapshot
    symlink, and removes the underlying blob only when it lives under this
    repo's own blobs/ dir and no other snapshot still references it (HF blobs
    are content-addressed and can be shared across revisions). A regular-file
    model.pt means a user-managed source dir — never deleted.
    """
    if os.environ.get(_EVICT_DMD_CKPT_ENV, "") != "1":
        return
    pt_path = os.path.join(source_dir, "sana_dit", "model.pt")
    if not os.path.islink(pt_path):
        return
    blob_path = os.path.realpath(pt_path)
    # source_dir = <repo_cache>/snapshots/<revision>
    repo_root = os.path.dirname(os.path.dirname(os.path.normpath(source_dir)))
    blobs_dir = os.path.join(repo_root, "blobs")
    snapshots_dir = os.path.join(repo_root, "snapshots")
    try:
        if os.path.commonpath([blob_path, blobs_dir]) != blobs_dir:
            return
        os.unlink(pt_path)
        for root, _, files in os.walk(snapshots_dir):
            for name in files:
                link = os.path.join(root, name)
                if os.path.islink(link) and os.path.realpath(link) == blob_path:
                    return  # another snapshot still needs this blob
        os.unlink(blob_path)
    except OSError:
        pass


def _patch_class_name(config_path: str, new_class_name: str) -> None:
    if not os.path.isfile(config_path):
        return
    with open(config_path, encoding="utf-8") as f:
        cfg = json.load(f)
    cfg["_class_name"] = new_class_name
    with open(config_path, "w", encoding="utf-8") as f:
        json.dump(cfg, f, indent=2, sort_keys=True)
        f.write("\n")


def _materialize_refiner(*, source_dir: str, output_dir: str) -> None:
    dst = os.path.join(output_dir, "refiner")
    for sub in ("transformer", "connectors"):
        src = os.path.join(source_dir, "refiner_diffusers", sub)
        if not os.path.isdir(src):
            raise FileNotFoundError(f"refiner component missing: {src}")
        _copytree_link_or_copy(src, os.path.join(dst, sub))
    gemma3 = os.path.join(source_dir, "gemma3_12b")
    if not os.path.isdir(gemma3):
        raise FileNotFoundError(f"refiner text encoder missing: {gemma3}")
    _copytree_link_or_copy(gemma3, os.path.join(dst, "text_encoder"))
    # SGLang routes loadable modules by _class_name; rewrite the copies so the
    # framework loader picks sglang's video-only refiner implementation.
    _patch_class_name(
        os.path.join(dst, "transformer", "config.json"), "SanaWMLTX2VideoRefiner"
    )
    _patch_class_name(
        os.path.join(dst, "connectors", "config.json"), "LTX2TextConnectors"
    )


def _materialize_text_encoder_and_tokenizer(*, output_dir: str) -> None:
    cached_repo = snapshot_download(
        repo_id=GEMMA2_REPO_ID,
        revision=GEMMA2_REVISION,
        allow_patterns=[
            "config.json", "generation_config.json",
            "model.safetensors", "model-*.safetensors", "model.safetensors.index.json",
            "tokenizer.json", "tokenizer.model", "tokenizer_config.json",
            "special_tokens_map.json", "added_tokens.json",
        ],
        max_workers=4,
    )
    text_encoder_out = os.path.join(output_dir, "text_encoder")
    tokenizer_out = os.path.join(output_dir, "tokenizer")
    _ensure_dir(text_encoder_out)
    _ensure_dir(tokenizer_out)
    tokenizer_filenames = {
        "tokenizer.json", "tokenizer.model", "tokenizer_config.json",
        "special_tokens_map.json", "added_tokens.json",
    }
    encoder_filenames = {
        "config.json", "generation_config.json",
        "model.safetensors", "model.safetensors.index.json",
    }
    for entry in sorted(os.listdir(cached_repo)):
        src = os.path.join(cached_repo, entry)
        if not os.path.isfile(src):
            continue
        if entry in tokenizer_filenames:
            _link_or_copy_file(src, os.path.join(tokenizer_out, entry))
        elif entry in encoder_filenames or (
            entry.startswith("model-") and entry.endswith(".safetensors")
        ):
            _link_or_copy_file(src, os.path.join(text_encoder_out, entry))


def materialize(
    *,
    overlay_dir: str,
    source_dir: str,
    output_dir: str,
    manifest: dict[str, Any],
) -> None:
    _convert_dmd_dit(
        source_dir=source_dir, overlay_dir=overlay_dir, output_dir=output_dir
    )
    _copytree_link_or_copy(
        os.path.join(source_dir, "ltx2_causal_vae"), os.path.join(output_dir, "vae")
    )
    _materialize_refiner(source_dir=source_dir, output_dir=output_dir)
    _write_json(
        os.path.join(output_dir, "scheduler", "scheduler_config.json"),
        {
            "_class_name": "FlowMatchEulerDiscreteScheduler",
            "_diffusers_version": "0.37.0.dev0",
            "num_train_timesteps": 1000,
            # Dense-path shift; the streaming self-forcing path constructs its
            # own shift=1.0 scheduler with the explicit distilled sigma list.
            "shift": 9.95,
            "use_dynamic_shifting": False,
        },
    )
    _materialize_text_encoder_and_tokenizer(output_dir=output_dir)
    _maybe_evict_dmd_ckpt(source_dir)