"""Text-to-Image — hybrid MLX + PyTorch.

Phase 1 (MLX): block-diffusion VQ token generation with CFG.
Phase 2 (PyTorch): SigVQ + ZImageTransformer2DModel + VAE → pixel image.

The MLX backbone is released before the PyTorch decoder loads to fit in 64 GB
unified memory (decoder is ~12 GB, backbone is ~32 GB).
"""
import argparse
import gc
import json
import os
import sys
import time
from pathlib import Path

import mlx.core as mx
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer

REPO_ROOT = Path(__file__).resolve().parent.parent / "llada2-uni-repo"
sys.path.insert(0, str(REPO_ROOT))

# Stub out flash_attn (not available on Apple Silicon). The decoder has a
# dispatch_attention_fn fallback via diffusers that we use instead.
import types as _types, importlib.machinery as _im
if "flash_attn" not in sys.modules:
    _stub = _types.ModuleType("flash_attn")
    _stub.__spec__ = _im.ModuleSpec(name="flash_attn", loader=None)
    _stub.__version__ = "0.0.0-stub"
    _stub.flash_attn_func = lambda *a, **k: (_ for _ in ()).throw(
        RuntimeError("flash_attn unavailable"))
    sys.modules["flash_attn"] = _stub

from llada2.model import LLaDA2Config, LLaDA2Model
from llada2.weights import load_weights_into_model
from llada2.generate_image import generate_image_tokens, extract_vq_tokens


def build_t2i_prompt(tokenizer, prompt_text: str, image_h: int, image_w: int):
    """Return (cond_ids, uncond_ids) — prompt id lists for CFG."""
    sys_tmpl = "You are a text-to-image generation assistant."
    # _build_chat equivalent
    sys_ids = tokenizer(f"<role>SYSTEM</role> {sys_tmpl} <role>HUMAN</role>").input_ids
    asst_ids = tokenizer("<role>ASSISTANT</role>").input_ids

    soi = tokenizer("<|image|>").input_ids
    boi = tokenizer("<boi>").input_ids
    h_tok = tokenizer(f"<|reserved_token_{image_h}|>").input_ids
    w_tok = tokenizer(f"<|reserved_token_{image_w}|>").input_ids
    img_header = soi + h_tok + w_tok + boi

    cond_ids = sys_ids + tokenizer(prompt_text).input_ids + asst_ids + img_header
    uncond_ids = sys_ids + tokenizer("<uncondition>").input_ids + asst_ids + img_header
    return cond_ids, uncond_ids


def decode_to_pixels(token_ids: list[int], h: int, w: int, model_path: Path,
                     decoder_steps: int, resolution_multiplier: int,
                     decode_mode: str = "decoder-turbo"):
    """Call the official decoder to render pixels."""
    import torch
    from decoder import decode_vq_tokens
    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    return decode_vq_tokens(
        token_ids, h, w, str(model_path), device,
        resolution_multiplier=resolution_multiplier,
        num_steps=decoder_steps, decode_mode=decode_mode,
    )


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompt", required=True, type=str)
    ap.add_argument("--image-h", default=512, type=int)
    ap.add_argument("--image-w", default=512, type=int)
    ap.add_argument("--steps", default=16, type=int)
    ap.add_argument("--block-length", default=32, type=int)
    ap.add_argument("--cfg-scale", default=4.0, type=float)
    ap.add_argument("--decoder-steps", default=50, type=int)
    ap.add_argument("--decode-mode", default="normal",
                    choices=["decoder-turbo", "normal"],
                    help="'normal' = full 50-step decoder (cleaner, ~8 min), "
                         "'decoder-turbo' = 8-step distilled (faster but brittle ≈ striping)")
    ap.add_argument("--resolution-multiplier", default=2, type=int)
    ap.add_argument("--output", default="t2i_output.png", type=str)
    ap.add_argument("--repo-id", default="inclusionAI/LLaDA2.0-Uni", type=str)
    ap.add_argument("--save-vq", default=None, type=str, help="Save intermediate VQ tokens to .json")
    ap.add_argument("--load-vq", default=None, type=str, help="Skip phase 1, load VQ tokens from .json")
    args = ap.parse_args()

    print("[t2i] fetching model files…")
    snap = Path(snapshot_download(
        args.repo_id,
        allow_patterns=[
            "model-*.safetensors", "model.safetensors.index.json",
            "config.json", "tokenizer*", "special_tokens_map.json",
            "decoder-turbo/*", "decoder/*", "image_tokenizer/*", "vae/*",
        ],
    ))

    # Generate image: LLaDA2 divides H and W by 2 internally before computing grid.
    # Net result: grid = (image_h // 2 // 16) x (image_w // 2 // 16)
    grid_h = args.image_h // 2 // 16
    grid_w = args.image_w // 2 // 16
    gen_length = grid_h * grid_w

    if args.load_vq:
        with open(args.load_vq) as f:
            cached = json.load(f)
        vq_tokens = cached["token_ids"]
        grid_h, grid_w = cached["h"], cached["w"]
        print(f"[t2i] loaded {len(vq_tokens)} VQ tokens from {args.load_vq}")
    else:
        # ---------- Phase 1: MLX VQ-token generation ----------
        tokenizer = AutoTokenizer.from_pretrained(str(snap), trust_remote_code=True)
        config = LLaDA2Config.from_hf(json.loads((snap / "config.json").read_text()))

        cond_ids, uncond_ids = build_t2i_prompt(tokenizer, args.prompt, grid_h, grid_w)
        print(f"[t2i] prompt tokens: {len(cond_ids)} | grid: {grid_h}x{grid_w} ({gen_length} VQ tokens)")

        print("[t2i] building model + loading backbone…")
        model = LLaDA2Model(config)
        t0 = time.time()
        load_weights_into_model(model, snap, dtype=mx.bfloat16, verbose=False)
        print(f"[t2i] backbone loaded in {time.time()-t0:.1f}s")

        prompt_ids = mx.array([cond_ids], dtype=mx.int32)
        uc_ids = mx.array([uncond_ids], dtype=mx.int32)

        t0 = time.time()
        out = generate_image_tokens(
            model, prompt_ids, uc_ids,
            gen_length=gen_length,
            block_length=args.block_length,
            steps_per_block=args.steps,
            cfg_scale=args.cfg_scale,
            mask_token_id=config.mask_token_id,
            image_token_offset=config.image_token_offset,
            vocab_size=config.vocab_size,
        )
        mx.eval(out)
        vq_tokens = (out[0, len(cond_ids):len(cond_ids) + gen_length] - config.image_token_offset).tolist()
        print(f"[t2i] VQ generation in {time.time()-t0:.1f}s, {len(vq_tokens)} tokens, "
              f"range [{min(vq_tokens)}, {max(vq_tokens)}]")

        if args.save_vq:
            with open(args.save_vq, "w") as f:
                json.dump({"token_ids": vq_tokens, "h": grid_h, "w": grid_w,
                           "prompt": args.prompt}, f)
            print(f"[t2i] saved VQ tokens → {args.save_vq}")

        # ---------- Free MLX backbone before PyTorch decoder loads ----------
        del model, out
        gc.collect()
        mx.clear_cache()

    # ---------- Phase 2: PyTorch decode → pixels ----------
    print(f"[t2i] decoding VQ tokens → pixels ({args.decoder_steps} steps)…")
    t0 = time.time()
    img = decode_to_pixels(
        vq_tokens, grid_h, grid_w, snap,
        decoder_steps=args.decoder_steps,
        resolution_multiplier=args.resolution_multiplier,
        decode_mode=args.decode_mode,
    )
    print(f"[t2i] decoded in {time.time()-t0:.1f}s")

    img.save(args.output)
    print(f"[t2i] wrote {args.output}")


if __name__ == "__main__":
    main()