"""Generate real gallery assets by driving the live Hugging Face Space.

The whole pipeline (DeepSeek-R1 -> FLUX -> TRELLIS) runs on the Space's ZeroGPU;
this script just orchestrates the calls over the gradio_client API and saves the
produced image / video / GLB locally, then writes a manifest the app reads.

Usage:
    python scripts/generate_gallery.py --smoke      # one prompt, prints shapes
    python scripts/generate_gallery.py              # full list

Auth: uses the locally cached HF token (huggingface-cli login).
"""
import argparse
import json
import os
import re
import shutil

from gradio_client import Client, handle_file
from huggingface_hub import get_token

SPACE_ID = "yonnel/text-to-3d_flux_trellis"
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
GALLERY_DIR = os.path.join(ROOT, "assets", "gallery")

PROMPTS = [
    "a backpack for kids, flower style",
    "medieval flip flops",
    "cat shaped cake mold",
]


def slugify(text: str) -> str:
    return re.sub(r"[^a-z0-9]+", "_", text.lower()).strip("_")[:40]


def as_path(x):
    """Extract a local filepath from a gradio_client result item."""
    if x is None:
        return None
    if isinstance(x, str):
        return x
    if isinstance(x, dict):
        if "video" in x:
            v = x["video"]
            return v if isinstance(v, str) else (v.get("path") or v.get("url"))
        return x.get("path") or x.get("url")
    if isinstance(x, (list, tuple)) and x:
        return as_path(x[0])
    return None


def generate_one(client: Client, prompt: str, verbose: bool = True) -> dict:
    # Step 1 - DeepSeek-R1: prompt refinement
    r1 = client.predict(prompt, api_name="/refine_prompt")
    reasoning, clean_prompt, status1 = r1
    if verbose:
        print(f"  [1] status: {status1!r}")
        print(f"  [1] clean_prompt: {clean_prompt[:160]!r}")
    if not clean_prompt:
        raise RuntimeError(f"empty refined prompt (status: {status1})")

    # Step 2 - FLUX: image (fixed seed so the gallery is reproducible)
    r2 = client.predict(
        clean_prompt, 42, False, 512, 512, 6, api_name="/generate_image"
    )
    image_path = as_path(r2[0])
    if verbose:
        print(f"  [2] image: {image_path}")
    if not image_path or not os.path.exists(image_path):
        raise RuntimeError(f"no image produced (status: {r2[-1]})")

    # Step 3 - TRELLIS: video + GLB
    r3 = client.predict(
        handle_file(image_path), 0, 7.5, 12, 3.0, 12, api_name="/image_to_3d"
    )
    video_path = as_path(r3[0])
    glb_path = as_path(r3[1]) or as_path(r3[3])
    if verbose:
        print(f"  [3] video: {video_path}")
        print(f"  [3] glb:   {glb_path}")

    return {
        "prompt": prompt,
        "refined_prompt": clean_prompt,
        "reasoning": reasoning,
        "_image": image_path,
        "_video": video_path,
        "_glb": glb_path,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--smoke", action="store_true", help="only the first prompt")
    args = ap.parse_args()

    prompts = PROMPTS[:1] if args.smoke else PROMPTS
    os.makedirs(GALLERY_DIR, exist_ok=True)

    print(f"Connecting to {SPACE_ID} ...")
    client = Client(SPACE_ID, token=get_token(), verbose=False)
    # The per-session temp dir is created by the Space's `load` handler, which
    # never fires over the API -> call it explicitly so image_to_3d can write.
    client.predict(api_name="/start_session")

    manifest = []
    for i, prompt in enumerate(prompts, 1):
        print(f"[{i}/{len(prompts)}] {prompt}")
        try:
            res = generate_one(client, prompt)
        except Exception as e:
            print(f"  !! failed: {e}")
            continue

        slug = slugify(prompt)
        item_dir = os.path.join(GALLERY_DIR, slug)
        os.makedirs(item_dir, exist_ok=True)
        entry = {"prompt": prompt, "refined_prompt": res["refined_prompt"]}
        for kind, key, ext in [
            ("image", "_image", ".png"),
            ("video", "_video", ".mp4"),
            ("glb", "_glb", ".glb"),
        ]:
            src = res[key]
            if src and os.path.exists(src):
                dst = os.path.join(item_dir, kind + ext)
                shutil.copyfile(src, dst)
                entry[kind] = os.path.relpath(dst, ROOT).replace("\\", "/")
        manifest.append(entry)
        print(f"  -> saved to {item_dir}")

    if args.smoke:
        print("\nSMOKE RESULT:\n" + json.dumps(manifest, indent=2, ensure_ascii=False))
        return

    manifest_path = os.path.join(GALLERY_DIR, "manifest.json")
    with open(manifest_path, "w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2, ensure_ascii=False)
    print(f"\nWrote {manifest_path} ({len(manifest)} items)")


if __name__ == "__main__":
    main()