#!/usr/bin/env python3
"""建 ActionRoPE 的 UMT5 文本表（动作从句已剥离）。

数据集自带的 text_table_eybx.pt 是 1,253 键，每一条都含 "the player is {动作}"
—— 那是动作走文本的老方案。ActionRoPE 的动作走 RoPE，文本只留场景，剥完
只剩 360 条不同的串（18 纯场景 + 342 转场），所以要重建一张表。

编码走 DiffSynth 自己的 tokenizer + WanTextEncoder，与训练/推理时
WanVideoUnit_PromptEmbedder 的路径逐字一致：
  · HuggingfaceTokenizer(seq_len=512, clean='whitespace')
  · prompt_emb[:, seq_len:] = 0   （尾部 pad 清零）
产物与老表同构：dict[str, bf16[512, 4096]]，训练时直接查表，不必挂 11 GB 的 UMT5。

    python -m actionrope.build_text_table --out data/latent/text_table_actionrope.pt
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import sys
import time

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from actionrope.prompts import NULL_PROMPT, strip_action  # noqa: E402

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_CKPT = os.path.join(ROOT, "models", "Wan2.2-TI2V-5B")


def collect_keys(meta_glob: str, null_prompt: str) -> list[str]:
    """从 meta/*.jsonl 收全部逐 cell prompt，剥掉动作从句后去重。

    走 meta 而不是 22,782 个 .pt：meta 只有 65 MB，.pt 要读 72 GB，
    两者的 prompts 字段是同一份（encode_latents.py 直接从 meta 抄进 sidecar）。
    """
    keys, n_rows, n_cells = set(), 0, 0
    files = sorted(glob.glob(meta_glob))
    if not files:
        raise FileNotFoundError(f"没有 meta 文件: {meta_glob}")
    for f in files:
        with open(f, encoding="utf-8") as fh:
            for line in fh:
                r = json.loads(line)
                n_rows += 1
                for p in r["prompts"]:
                    n_cells += 1
                    keys.add(strip_action(p))
    print(f"扫过 {n_rows:,} 个 clip / {n_cells:,} 个 cell -> {len(keys):,} 个不同串")
    # 场景 dropout / CFG 用的空条件串。剥离后 prompts_bossdrop 的 21 条会全部
    # 塌成同一个「没有场景」的串，这里显式加进表里，训练侧才有得查。
    if null_prompt not in keys:
        keys.add(null_prompt)
        print(f"追加空条件串: {null_prompt!r}")
    return sorted(keys)


def main():
    ap = argparse.ArgumentParser(description="ActionRoPE UMT5 文本表")
    ap.add_argument("--meta", default=os.path.join(ROOT, "data/meta/*.jsonl"))
    ap.add_argument("--out", default=os.path.join(ROOT, "data/latent/text_table_actionrope.pt"))
    ap.add_argument("--ckpt", default=DEFAULT_CKPT, help="Wan2.2-TI2V-5B 本地目录")
    ap.add_argument("--null_prompt", default=NULL_PROMPT,
                    help="场景 dropout / CFG 的空条件串")
    ap.add_argument("--batch", type=int, default=32)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--dry_run", action="store_true", help="只收串、不加载 UMT5")
    args = ap.parse_args()

    keys = collect_keys(args.meta, args.null_prompt)
    plain = [k for k in keys if "burning away into" not in k and "cutting to" not in k]
    print(f"  纯场景 {len(plain)} 条 / 转场 {len(keys) - len(plain)} 条")
    if args.dry_run:
        for k in plain:
            print("   ", k)
        return

    from diffsynth.pipelines.wan_video import WanVideoPipeline, ModelConfig

    # 只加载 text_encoder：DiT 和 VAE 这一步用不上，省 15 GB 显存。
    # redirect_common_files=False —— 本地已有 .pth，不要去 ModelScope 拉 safetensors。
    pipe = WanVideoPipeline.from_pretrained(
        torch_dtype=torch.bfloat16,
        device=args.device,
        model_configs=[ModelConfig(path=os.path.join(args.ckpt, "models_t5_umt5-xxl-enc-bf16.pth"))],
        tokenizer_config=ModelConfig(path=os.path.join(args.ckpt, "google/umt5-xxl")),
        redirect_common_files=False,
    )
    pipe.load_models_to_device(["text_encoder"])

    table, t0 = {}, time.time()
    for i in range(0, len(keys), args.batch):
        chunk = keys[i:i + args.batch]
        ids, mask = pipe.tokenizer(chunk, return_mask=True, add_special_tokens=True)
        ids, mask = ids.to(args.device), mask.to(args.device)
        seq_lens = mask.gt(0).sum(dim=1).long()
        with torch.no_grad():
            emb = pipe.text_encoder(ids, mask)
        # 与 WanVideoUnit_PromptEmbedder 一致：有效长度之后一律清零
        for j, v in enumerate(seq_lens):
            emb[j, v:] = 0
        for k, e, v in zip(chunk, emb, seq_lens):
            table[k] = e.to(torch.bfloat16).cpu().clone()
        print(f"  {min(i + args.batch, len(keys)):4d}/{len(keys)} · {time.time() - t0:.0f}s", flush=True)

    shape = next(iter(table.values())).shape
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    torch.save(table, args.out)
    print(f"DONE  {len(table):,} 键 · 每键 {tuple(shape)} bf16 · "
          f"{os.path.getsize(args.out) / 1e6:.0f} MB -> {args.out}")


if __name__ == "__main__":
    main()