app.py

"""
Open3DForge — Image-to-game-ready 3D asset pipeline
====================================================

Milestone 1: Foundation
-----------------------
This is the scaffold. It validates that:
  - The HF Space builds with the right Gradio + ZeroGPU configuration
  - `@spaces.GPU` allocates and releases an H200 successfully
  - `gr.Model3D` renders GLB files inline
  - The 5-tab UI structure is sound
  - Workspace folder management works
  - Quota tracking persists across page reloads

Subsequent milestones fill in the actual pipeline stages.

Deployed exclusively on HF Spaces ZeroGPU — no local execution path.
"""

from __future__ import annotations

import os
import subprocess
import sys
import time
from pathlib import Path

# nvdiffrast is compiled against CUDA 13.0 and its .so links to libcudart.so.13.
# os.environ["LD_LIBRARY_PATH"] only affects child processes, not the running
# Python interpreter.  Use ctypes.CDLL with RTLD_GLOBAL to load libcudart.so.13
# into the current process before nvdiffrast is imported.
import ctypes, glob as _glob

def _preload_cudart() -> None:
    cuda_lib_dir = "/cuda-image/usr/local/cuda-13.0/lib64"
    # Try versioned name first, then unversioned fallback.
    for pattern in (
        f"{cuda_lib_dir}/libcudart.so.13*",
        f"{cuda_lib_dir}/libcudart.so",
    ):
        matches = _glob.glob(pattern)
        if matches:
            try:
                ctypes.CDLL(matches[0], mode=ctypes.RTLD_GLOBAL)
                print(f"Preloaded {matches[0]}", flush=True)
            except OSError as e:
                print(f"Preload failed for {matches[0]}: {e}", flush=True)
            return

_preload_cudart()

import gradio as gr
import spaces  # HF ZeroGPU


# GPU arch list used to compile nvdiffrast.
# +PTX on the last entry generates portable PTX code that the CUDA runtime
# JIT-compiles for any newer GPU (e.g. Blackwell SM 10.0) at first use.
_NVDR_ARCH = "8.0;8.9;9.0+PTX"
_NVDR_MARKER = Path("/tmp/nvdiffrast_arch.txt")


def _install_nvdiffrast() -> None:
    """Install nvdiffrast from source at startup.

    Must run here (not requirements.txt) because pip's build isolation hides
    torch from nvdiffrast's build system.

    Shim strategy: ZeroGPU machine has CUDA 13.0 toolkit but torch 2.9.1 was
    compiled with CUDA 12.8. PyTorch's _check_cuda_version raises on mismatch.
    We put a shim nvcc on PATH that reports 12.8 for --version but delegates
    real compilation to the actual CUDA 13.0 nvcc.

    Arch strategy: compile native kernels for SM 8.0 (A100) and 8.9 (L4/L40),
    plus PTX for SM 9.0 (H100/H200). The +PTX entry lets the CUDA 13.0 runtime
    JIT-compile for any newer architecture (e.g. SM 10.0 Blackwell) at first use.

    Marker file: /tmp/nvdiffrast_arch.txt stores the arch string used.  If it
    matches _NVDR_ARCH, skip rebuild on warm container restarts.
    """
    import os
    import shutil
    import tempfile
    import torch

    # Check if already installed with the current arch — skip rebuild if so.
    try:
        import nvdiffrast  # noqa: F401
        if _NVDR_MARKER.exists() and _NVDR_MARKER.read_text().strip() == _NVDR_ARCH:
            print("nvdiffrast: already installed with correct arch.", flush=True)
            return
        print("nvdiffrast: arch mismatch or missing marker — forcing reinstall.", flush=True)
        subprocess.run(
            [sys.executable, "-m", "pip", "uninstall", "nvdiffrast", "-y"],
            check=False, capture_output=True,
        )
    except (ImportError, OSError):
        pass

    torch_cuda = torch.version.cuda or "12.8"
    print(f"nvdiffrast: torch.version.cuda={torch_cuda}", flush=True)

    # Locate real nvcc.
    real_nvcc: str | None = None
    real_cuda_home: str | None = None
    for candidate in [
        os.environ.get("CUDA_HOME", ""),
        "/cuda-image/usr/local/cuda-13.0",
        "/cuda-image/usr/local/cuda-12.8",
        "/cuda-image/usr/local/cuda",
        "/usr/local/cuda-13.0",
        "/usr/local/cuda",
    ]:
        if candidate and os.path.isfile(os.path.join(candidate, "bin", "nvcc")):
            real_cuda_home = candidate
            real_nvcc = os.path.join(candidate, "bin", "nvcc")
            break
    if real_nvcc is None:
        real_nvcc = shutil.which("nvcc")
    print(f"nvdiffrast: real_nvcc={real_nvcc}", flush=True)

    env = os.environ.copy()

    if real_nvcc:
        fake_cuda = Path(tempfile.mkdtemp(prefix="cuda_shim_"))
        (fake_cuda / "bin").mkdir()
        if real_cuda_home:
            for sub in ("include", "lib", "lib64"):
                src = Path(real_cuda_home) / sub
                if src.exists():
                    (fake_cuda / sub).symlink_to(src)
        shim = fake_cuda / "bin" / "nvcc"
        shim.write_text(
            "#!/bin/sh\n"
            'case "$1" in\n'
            '  --version)\n'
            f'    echo "Cuda compilation tools, release {torch_cuda}, V{torch_cuda}.0"\n'
            '    ;;\n'
            '  *)\n'
            f'    exec {real_nvcc} "$@"\n'
            '    ;;\n'
            'esac\n'
        )
        shim.chmod(0o755)
        env["CUDA_HOME"] = str(fake_cuda)
        env["CUDA_PATH"] = str(fake_cuda)
        print(f"nvdiffrast: shim CUDA_HOME={fake_cuda}", flush=True)

    # PTX in the arch list lets the runtime JIT-compile for any newer GPU
    # (e.g. Blackwell SM 10.0) that lacks a pre-compiled cubin.
    env["TORCH_CUDA_ARCH_LIST"] = _NVDR_ARCH

    print(f"Building nvdiffrast (arch={_NVDR_ARCH}, ~2 min)...", flush=True)
    result = subprocess.run(
        [sys.executable, "-m", "pip", "install", "--no-build-isolation",
         "git+https://github.com/NVlabs/nvdiffrast.git"],
        env=env, check=False, capture_output=True, text=True,
    )
    print(result.stdout[-3000:] if result.stdout else "", flush=True)
    if result.returncode != 0:
        print(f"nvdiffrast build FAILED (rc={result.returncode}):", flush=True)
        print(result.stderr[-2000:] if result.stderr else "", flush=True)
    else:
        print("nvdiffrast build succeeded.", flush=True)
        _NVDR_MARKER.write_text(_NVDR_ARCH)


_install_nvdiffrast()

from src import quota, ui_helpers, workspace
from src.stages.stage1_generate import generate_trellis, generate_hunyuan

# ---------------------------------------------------------------------------
# ZeroGPU test function — proves that GPU allocation works.
# Replaced in Milestone 2 with the actual TRELLIS.2 generator.
# ---------------------------------------------------------------------------

@spaces.GPU(duration=15)
def zerogpu_smoke_test() -> str:
    """Allocate a GPU, run a trivial torch op, release. Reports timing."""
    start = time.time()
    try:
        import torch
        if not torch.cuda.is_available():
            return "❌ GPU not available inside @spaces.GPU. Something is wrong."

        device = torch.device("cuda")
        # Trivial GPU work
        a = torch.randn(1024, 1024, device=device)
        b = torch.randn(1024, 1024, device=device)
        c = a @ b
        torch.cuda.synchronize()
        result_sum = float(c.sum().item())

        gpu_name = torch.cuda.get_device_name(0)
        vram_total = torch.cuda.get_device_properties(0).total_memory / 1e9
        elapsed = time.time() - start

        quota.record_usage("smoke_test", elapsed)

        return (
            f"✅ **GPU allocation successful**\n\n"
            f"- Device: `{gpu_name}`\n"
            f"- VRAM: `{vram_total:.1f} GB`\n"
            f"- Test op (1024×1024 matmul): `{elapsed:.2f}s`\n"
            f"- Result sum: `{result_sum:.2f}` (sanity check, non-zero = ✓)\n\n"
            f"ZeroGPU integration is working. Ready for Milestone 2."
        )
    except Exception as e:
        elapsed = time.time() - start
        return f"❌ GPU test failed after {elapsed:.2f}s:\n```\n{type(e).__name__}: {e}\n```"


# ---------------------------------------------------------------------------
# Stage stubs — return placeholder messages until milestones implement them.
# ---------------------------------------------------------------------------

def _stub(stage: str) -> str:
    return (
        f"🚧 **{stage}** — not implemented yet.\n\n"
        f"This stub will be replaced in a future milestone. "
        f"See `PLAN.md` for the full pipeline spec."
    )


def handle_generate(images, model, quality, seed, _steps, _octree, tex_size, _symmetry, do_rembg):
    """Dispatch to the correct generation backend.

    Yields (status_markdown, viewer_path_or_None) tuples for streaming.
    The viewer path is yielded only in the final tuple so Gradio can serve
    the GLB directly without a separate state lookup.
    """
    if "TRELLIS" in model:
        yield f"⚙️ **TRELLIS.2** · {quality}\n\nContacting remote Space...", None
        result = generate_trellis(images, quality, int(seed), int(tex_size))
        viewer_path = ui_helpers.get_viewer_model_path()
        yield result, viewer_path
        return
    yield from generate_hunyuan(images, quality, int(seed), int(tex_size), do_rembg=bool(do_rembg))


@spaces.GPU(duration=600)
def run_post_process(
    do_repair, do_cleanup, do_decimate, target_faces,
    do_symmetry, do_unwrap, do_normal_bake, normal_format,
    do_albedo, do_material, do_ao, ao_quality,
    do_inpaint, do_lods, do_collision, pivot, scale_m,
):
    """Run post-processing pipeline. Yields cumulative status markdown for streaming."""
    state = workspace.get_state()
    current_glb = state.raw_gen_glb or state.high_poly_glb
    if not current_glb or not current_glb.exists():
        yield "❌ No generated asset found. Run Stage 1 (Generate) first."
        return

    log = []

    def _emit(line: str):
        log.append(line)
        return "\n".join(log)

    if do_repair:
        yield _emit("⏳ Repairing mesh (pymeshfix)...")
        try:
            from src.stages.stage2_repair import repair_mesh
            current_glb, msg = repair_mesh(current_glb)
            yield _emit(f"✅ {msg}")
        except Exception as e:
            yield _emit(f"⚠️ Repair error: {e}")

    if do_cleanup:
        yield _emit("⏳ Cleaning geometry (trimesh)...")
        try:
            from src.stages.stage2_cleanup import cleanup_mesh
            current_glb, msg = cleanup_mesh(current_glb)
            yield _emit(f"✅ {msg}")
        except Exception as e:
            yield _emit(f"⚠️ Cleanup error: {e}")

    if do_decimate:
        yield _emit(f"⏳ Decimating to {int(target_faces):,} faces...")
        try:
            from src.stages.stage2_decimate import decimate_mesh_final
            current_glb, msg = decimate_mesh_final(current_glb, int(target_faces))
            yield _emit(f"✅ {msg}")
        except Exception as e:
            yield _emit(f"⚠️ Decimation error: {e}")

    if do_symmetry:
        yield _emit("⏳ Enforcing bilateral symmetry...")
        try:
            from src.stages.stage2_symmetry import apply_symmetry
            current_glb, msg = apply_symmetry(current_glb, "bilateral-X")
            yield _emit(f"✅ {msg}")
        except Exception as e:
            yield _emit(f"⚠️ Symmetry error: {e}")

    if do_unwrap:
        yield _emit("⏳ UV unwrapping (xatlas)...")
        try:
            from src.stages.stage2_uv import unwrap_uvs
            current_glb, msg = unwrap_uvs(current_glb)
            yield _emit(f"✅ {msg}")
        except Exception as e:
            yield _emit(f"⚠️ UV unwrap error: {e}")

    if do_normal_bake:
        yield _emit("⏳ Baking normal map (GPU)...")
        try:
            from src.stages.stage2_bake_normal import bake_normal_map
            st = workspace.get_state()
            hp = st.high_poly_glb
            lo = st.unwrapped_glb or current_glb
            if not hp or not hp.exists():
                yield _emit("⚠️ Normal bake: no high-poly GLB. Generate first.")
            else:
                _gl, _dx, msg = bake_normal_map(hp, lo, map_size=2048, dx_format=True)
                yield _emit(f"✅ {msg}")
        except Exception as e:
            yield _emit(f"⚠️ Normal bake error: {e}")

    st = workspace.get_state()
    hp = st.high_poly_glb
    lo = st.unwrapped_glb or current_glb

    if do_albedo:
        yield _emit("⏳ Baking albedo map (GPU)...")
        try:
            from src.stages.stage2_bake_albedo import bake_albedo
            if not hp or not hp.exists():
                yield _emit("⚠️ Albedo bake: no high-poly. Generate first.")
            else:
                _, msg = bake_albedo(hp, lo, map_size=2048)
                yield _emit(f"✅ {msg}")
        except Exception as e:
            yield _emit(f"⚠️ Albedo bake error: {e}")

    if do_material:
        yield _emit("⏳ Baking material maps (GPU)...")
        try:
            from src.stages.stage2_bake_albedo import bake_material
            if not hp or not hp.exists():
                yield _emit("⚠️ Material bake: no high-poly. Generate first.")
            else:
                _, _, msg = bake_material(hp, lo, map_size=2048)
                yield _emit(f"✅ {msg}")
        except Exception as e:
            yield _emit(f"⚠️ Material bake error: {e}")

    if do_ao:
        yield _emit(f"⏳ Baking AO ({ao_quality}, ray casting)...")
        try:
            from src.stages.stage2_bake_ao import bake_ao
            _, msg = bake_ao(current_glb, lo, map_size=2048, quality=ao_quality)
            yield _emit(f"✅ {msg}")
        except Exception as e:
            yield _emit(f"⚠️ AO bake error: {e}")

    st2 = workspace.get_state()

    if do_albedo or do_material or do_ao:
        yield _emit("⏳ Packing ORM texture...")
        try:
            from src.stages.stage2_finalize import pack_orm
            _, msg = pack_orm(st2.ao_png, st2.roughness_png, st2.metallic_png)
            yield _emit(f"✅ {msg}")
        except Exception as e:
            yield _emit(f"⚠️ ORM pack error: {e}")

    if do_lods:
        yield _emit("⏳ Generating LODs...")
        try:
            from src.stages.stage2_finalize import generate_lods
            lod_src = st2.final_glb or st2.low_poly_glb or current_glb
            _, msg = generate_lods(lod_src)
            yield _emit(f"✅ {msg}")
        except Exception as e:
            yield _emit(f"⚠️ LOD error: {e}")

    if do_collision:
        yield _emit("⏳ Generating collision mesh (CoACD)...")
        try:
            from src.stages.stage2_finalize import generate_collision
            col_src = st2.low_poly_glb or current_glb
            _, msg = generate_collision(col_src)
            yield _emit(f"✅ {msg}")
        except Exception as e:
            yield _emit(f"⚠️ Collision error: {e}")

    yield _emit("⏳ Setting pivot and scale...")
    try:
        from src.stages.stage2_finalize import set_pivot, validate_scale
        piv_src = st2.low_poly_glb or current_glb
        piv_src, msg = set_pivot(piv_src, pivot)
        yield _emit(f"✅ {msg}")
        _, msg = validate_scale(piv_src, float(scale_m))
        yield _emit(f"✅ {msg}")
    except Exception as e:
        yield _emit(f"⚠️ Pivot/scale error: {e}")

    if do_inpaint:
        yield _emit("🚧 SDXL inpaint — not yet implemented")

    final = workspace.get_state()
    out = final.final_glb or final.low_poly_glb or final.cleaned_glb or final.repaired_glb
    if out and out.exists():
        yield _emit(f"\n**Output:** `{out.name}` · ready for Stage 3 or Export.")


def handle_auto_rig(rig_type, seed, spring, fmt):
    from src.stages.stage3_rig import auto_rig
    _, msg = auto_rig(rig_type=rig_type, seed=int(seed))
    return msg


def handle_export(engine, asset_name, asset_type, include_lods, include_collision):
    from src.stages.stage4_export import export_ue5, _checklist
    state = workspace.get_state()

    if engine != "UE5":
        return f"🚧 {engine} export — not implemented. UE5 is the only supported engine.", None

    # Game-ready checklist
    issues = _checklist(state)
    checklist_md = ""
    if issues:
        checklist_md = "\n\n⚠️ **Checklist warnings:**\n" + "\n".join(f"- {i}" for i in issues)

    try:
        zip_path, msg, _ = export_ue5(
            asset_name=asset_name.strip() or "Asset",
            asset_type=asset_type,
            include_lods=include_lods,
            include_collision=include_collision,
        )
        return msg + checklist_md, str(zip_path)
    except Exception as e:
        return f"❌ Export failed: {e}{checklist_md}", None


# ---------------------------------------------------------------------------
# Presets tab logic (M11 — real save/load wired)
# ---------------------------------------------------------------------------

def ui_refresh_presets() -> gr.Dropdown:
    return gr.Dropdown(choices=workspace.list_presets(), label="Saved presets")


def ui_save_preset(
    name: str,
    gen_model, gen_quality, gen_seed, gen_tex_size, gen_rembg,
    pp_repair, pp_cleanup, pp_decimate, pp_target_faces,
    pp_symmetry, pp_unwrap, pp_normal_bake, pp_normal_format,
    pp_albedo, pp_material, pp_ao, pp_ao_quality,
    pp_inpaint, pp_lods, pp_collision, pp_pivot, pp_scale_m,
    rig_type, rig_seed, rig_format,
    ex_engine, ex_type,
) -> tuple[str, gr.Dropdown]:
    if not name or not name.strip():
        return "❌ Preset name cannot be empty.", ui_refresh_presets()
    config = {
        "name": name, "version": 1, "created_at": time.time(),
        "stage1": {"model": gen_model, "quality": gen_quality, "seed": int(gen_seed), "tex_size": int(gen_tex_size), "rembg": bool(gen_rembg)},
        "stage2": {
            "repair": pp_repair, "cleanup": pp_cleanup, "decimate": pp_decimate,
            "target_faces": int(pp_target_faces), "symmetry": pp_symmetry,
            "unwrap": pp_unwrap, "normal_bake": pp_normal_bake, "normal_format": pp_normal_format,
            "albedo_bake": pp_albedo, "material_bake": pp_material, "ao": pp_ao, "ao_quality": pp_ao_quality,
            "inpaint": pp_inpaint, "lods": pp_lods, "collision": pp_collision,
            "pivot": pp_pivot, "scale_m": float(pp_scale_m),
        },
        "stage3": {"rig_type": rig_type, "seed": int(rig_seed), "format": rig_format},
        "stage4": {"engine": ex_engine, "asset_type": ex_type},
    }
    workspace.save_preset(name, config)
    return f"✅ Saved preset: `{name}`", ui_refresh_presets()


def ui_load_preset(name: str):
    """Apply a saved preset to all UI components."""
    _no_op = gr.update()
    n = 27  # number of component outputs after pr_status
    if not name:
        return ("❌ Select a preset first.",) + (_no_op,) * n
    try:
        cfg = workspace.load_preset(name)
    except FileNotFoundError:
        return (f"❌ Preset '{name}' not found.",) + (_no_op,) * n

    s1 = cfg.get("stage1", {})
    s2 = cfg.get("stage2", {})
    s3 = cfg.get("stage3", {})
    s4 = cfg.get("stage4", {})

    return (
        f"✅ Loaded preset: `{name}`",
        # Stage 1
        gr.update(value=s1.get("model", "Hunyuan3D-2.1 (Organic / Characters)")),
        gr.update(value=s1.get("quality", "Balanced (~60s)")),
        gr.update(value=s1.get("seed", 42)),
        gr.update(value=s1.get("tex_size", 2048)),
        gr.update(value=s1.get("rembg", True)),
        # Stage 2
        gr.update(value=s2.get("repair", True)),
        gr.update(value=s2.get("cleanup", True)),
        gr.update(value=s2.get("decimate", True)),
        gr.update(value=s2.get("target_faces", 25000)),
        gr.update(value=s2.get("symmetry", False)),
        gr.update(value=s2.get("unwrap", True)),
        gr.update(value=s2.get("normal_bake", True)),
        gr.update(value=s2.get("normal_format", "DirectX (UE5)")),
        gr.update(value=s2.get("albedo_bake", True)),
        gr.update(value=s2.get("material_bake", True)),
        gr.update(value=s2.get("ao", True)),
        gr.update(value=s2.get("ao_quality", "Standard")),
        gr.update(value=s2.get("inpaint", False)),
        gr.update(value=s2.get("lods", True)),
        gr.update(value=s2.get("collision", True)),
        gr.update(value=s2.get("pivot", "bottom_center")),
        gr.update(value=s2.get("scale_m", 1.8)),
        # Stage 3
        gr.update(value=s3.get("rig_type", "Humanoid")),
        gr.update(value=s3.get("seed", 0)),
        gr.update(value=s3.get("format", "FBX (UE5 recommended)")),
        # Stage 4
        gr.update(value=s4.get("engine", "UE5")),
        gr.update(value=s4.get("asset_type", "Prop (SM_)")),
    )


def ui_delete_preset(name: str) -> tuple[str, gr.Dropdown]:
    if not name:
        return "❌ Select a preset to delete.", ui_refresh_presets()
    if workspace.delete_preset(name):
        return f"🗑️ Deleted preset: `{name}`", ui_refresh_presets()
    return f"❌ Preset not found: `{name}`", ui_refresh_presets()


# ---------------------------------------------------------------------------
# Layout
# ---------------------------------------------------------------------------

CUSTOM_CSS = """
.status-bar {
    font-size: 0.85em;
    color: #888;
    padding: 8px 12px;
    border-top: 1px solid #333;
    margin-top: 12px;
}
.asset-summary {
    font-size: 0.9em;
    background: rgba(255,255,255,0.03);
    padding: 12px;
    border-radius: 6px;
    border: 1px solid rgba(255,255,255,0.08);
}
"""


def build_ui() -> gr.Blocks:
    with gr.Blocks(
        title="Open3DForge",
    ) as demo:

        # --- Header --------------------------------------------------------
        gr.Markdown(
            "# 🛠️ Open3DForge\n"
            "*Personal image-to-game-ready 3D asset pipeline · UE5-first · "
            "Built on HF ZeroGPU*"
        )

        # --- Tabs ----------------------------------------------------------
        with gr.Tabs() as tabs:

            # ============ Tab 1: Generate =================================
            with gr.Tab("1. Generate", id=1):
                gr.Markdown(
                    "### Stage 1 — Image to 3D\n"
                    "Upload 1–4 reference images. Multi-view dramatically "
                    "improves quality for characters (front / 3-quarter / side / back)."
                )
                with gr.Row():
                    with gr.Column(scale=1):
                        gen_images = gr.File(
                            label="Reference images (1–4)",
                            file_count="multiple",
                            file_types=["image"],
                        )
                        gen_model = gr.Radio(
                            choices=[
                                "Hunyuan3D-2.1 (Organic / Characters)",
                                "TRELLIS.2 (Hard Surface)",
                            ],
                            value="Hunyuan3D-2.1 (Organic / Characters)",
                            label="Generation model",
                        )
                        gen_quality = gr.Radio(
                            choices=["Fast (~30s)", "Balanced (~60s)", "Hero (~90s)"],
                            value="Balanced (~60s)",
                            label="Quality preset",
                        )
                        gen_rembg = gr.Checkbox(
                            value=True,
                            label="Remove background automatically",
                            info="Uses rembg IS-Net to strip background before generation. "
                                 "Disable if your image already has a transparent background.",
                        )
                        with gr.Accordion("Advanced", open=False):
                            gen_seed = gr.Number(value=42, label="Seed", precision=0)
                            gen_steps = gr.Slider(
                                20, 50, value=35, step=5,
                                label="Inference steps",
                            )
                            gen_octree = gr.Dropdown(
                                choices=[256, 384, 512], value=384,
                                label="Octree resolution",
                            )
                            gen_tex_size = gr.Dropdown(
                                choices=[1024, 2048, 4096], value=2048,
                                label="Texture size",
                            )
                            gen_symmetry = gr.Radio(
                                choices=["off", "bilateral", "radial"],
                                value="off",
                                label="Symmetry hint",
                            )
                        gen_btn = gr.Button("Generate", variant="primary")
                    with gr.Column(scale=1):
                        gen_status = gr.Markdown("*Awaiting input.*")
                # Click wired below, after viewer component is defined.

            # ============ Tab 2: Post-Process =============================
            with gr.Tab("2. Post-Process", id=2):
                gr.Markdown(
                    "### Stage 2 — Mesh cleanup, UV unwrap, texture baking\n"
                    "Toggle steps on/off. Decimation has a live preview, the "
                    "rest run on confirm."
                )
                with gr.Row():
                    with gr.Column(scale=1):
                        pp_repair = gr.Checkbox(value=True, label="Mesh repair (pymeshfix)")
                        pp_cleanup = gr.Checkbox(value=True, label="Geometry cleanup (PyMeshLab)")
                        pp_decimate = gr.Checkbox(value=True, label="Decimation")
                        pp_target_faces = gr.Slider(
                            1000, 200000, value=25000, step=1000,
                            label="Target faces",
                        )
                        pp_symmetry = gr.Checkbox(value=False, label="Enforce bilateral symmetry")
                        pp_unwrap = gr.Checkbox(value=True, label="UV unwrap (xatlas)")
                        pp_normal_bake = gr.Checkbox(value=True, label="Normal bake (nvdiffrast)")
                        pp_normal_format = gr.Radio(
                            choices=["DirectX (UE5)", "OpenGL (Unity/Godot)"],
                            value="DirectX (UE5)",
                            label="Normal format",
                        )
                        pp_albedo_bake = gr.Checkbox(value=True, label="Albedo bake")
                        pp_material_bake = gr.Checkbox(value=True, label="Material bake (TRELLIS.2 attrs)")
                        pp_ao = gr.Checkbox(value=True, label="AO bake")
                        pp_ao_quality = gr.Radio(
                            choices=["Fast", "Standard", "High"],
                            value="Standard",
                            label="AO quality",
                        )
                        pp_inpaint = gr.Checkbox(value=False, label="SDXL inpaint hidden UVs (~30s GPU)")
                        pp_lods = gr.Checkbox(value=True, label="Generate LODs (LOD0/1/2)")
                        pp_collision = gr.Checkbox(value=True, label="Collision mesh (CoACD)")
                        pp_pivot = gr.Radio(
                            choices=["bottom_center", "geometric_center", "custom"],
                            value="bottom_center",
                            label="Pivot point",
                        )
                        pp_scale_m = gr.Number(value=1.8, label="Real-world height (meters)")
                        pp_btn = gr.Button("Run Post-Processing", variant="primary")
                    with gr.Column(scale=1):
                        pp_face_preview = gr.Markdown("*Face count preview: move the slider.*")
                        pp_status = gr.Markdown("*No asset to process. Generate one first.*")

                def _decimate_preview(target_faces):
                    state = workspace.get_state()
                    src = state.low_poly_glb or state.cleaned_glb or state.repaired_glb or state.raw_gen_glb
                    if not src or not src.exists():
                        return "*Generate an asset first.*"
                    try:
                        from src.stages.stage2_decimate import decimate_preview
                        fc, vc = decimate_preview(src, int(target_faces))
                        return f"**Preview:** ~{fc:,} faces · ~{vc:,} vertices at target {int(target_faces):,}"
                    except Exception as e:
                        return f"Preview error: {e}"

                pp_target_faces.change(
                    fn=_decimate_preview,
                    inputs=pp_target_faces,
                    outputs=pp_face_preview,
                )
                _pp_event = pp_btn.click(
                    fn=run_post_process,
                    inputs=[pp_repair, pp_cleanup, pp_decimate, pp_target_faces,
                            pp_symmetry, pp_unwrap, pp_normal_bake, pp_normal_format,
                            pp_albedo_bake, pp_material_bake, pp_ao, pp_ao_quality,
                            pp_inpaint, pp_lods, pp_collision, pp_pivot, pp_scale_m],
                    outputs=pp_status,
                )

            # ============ Tab 3: Auto-Rig =================================
            with gr.Tab("3. Auto-Rig", id=3):
                gr.Markdown(
                    "### Stage 3 — Auto-rigging (optional)\n"
                    "Uses UniRig (VAST-AI). For characters and creatures. "
                    "After rigging, drop the FBX into [Mixamo](https://mixamo.com) "
                    "for free animation presets."
                )
                with gr.Row():
                    with gr.Column(scale=1):
                        rig_type = gr.Dropdown(
                            choices=["Humanoid", "Quadruped", "Bird", "Insect", "Custom"],
                            value="Humanoid",
                            label="Character type",
                        )
                        rig_seed = gr.Number(value=0, label="Skeleton seed", precision=0)
                        rig_spring = gr.Checkbox(value=False, label="Spring bones (hair/cloth/tail)")
                        rig_format = gr.Radio(
                            choices=["FBX (UE5 recommended)", "GLB"],
                            value="FBX (UE5 recommended)",
                            label="Export format",
                        )
                        rig_btn = gr.Button("Auto-Rig", variant="primary")
                    with gr.Column(scale=1):
                        rig_status = gr.Markdown("*Process an asset in Stage 2 first.*")
                _rig_event = rig_btn.click(
                    fn=handle_auto_rig,
                    inputs=[rig_type, rig_seed, rig_spring, rig_format],
                    outputs=rig_status,
                )

            # ============ Tab 4: Export ===================================
            with gr.Tab("4. Export", id=4):
                gr.Markdown(
                    "### Stage 4 — Engine-ready export\n"
                    "UE5 default: FBX with DirectX normals + ORM-packed textures."
                )
                with gr.Row():
                    with gr.Column(scale=1):
                        ex_engine = gr.Dropdown(
                            choices=["UE5", "Unity (HDRP)", "Godot 4", "Blender", "Web (Three.js)"],
                            value="UE5",
                            label="Target engine",
                        )
                        ex_name = gr.Textbox(value="Asset_01", label="Asset name")
                        ex_type = gr.Radio(
                            choices=["Character (SK_)", "Prop (SM_)", "Environment (SM_)"],
                            value="Prop (SM_)",
                            label="Asset type",
                        )
                        ex_include_lods = gr.Checkbox(value=True, label="Include LODs")
                        ex_include_collision = gr.Checkbox(value=True, label="Include collision mesh")
                        ex_btn = gr.Button("Export", variant="primary")
                    with gr.Column(scale=1):
                        ex_status = gr.Markdown("*Nothing to export yet.*")
                        ex_file = gr.File(label="Download", visible=True)
                _ex_event = ex_btn.click(
                    fn=handle_export,
                    inputs=[ex_engine, ex_name, ex_type, ex_include_lods, ex_include_collision],
                    outputs=[ex_status, ex_file],
                )

            # ============ Tab 5: Presets ==================================
            with gr.Tab("5. Presets", id=5):
                gr.Markdown(
                    "### Saved configurations\n"
                    "Save the current settings across all tabs as a named preset. "
                    "Five defaults ship with the app: `character_UE5_hero`, "
                    "`character_UE5_npc`, `prop_UE5_hero`, `prop_UE5_standard`, "
                    "`environment_UE5_background`."
                )
                with gr.Row():
                    with gr.Column():
                        pr_list = gr.Dropdown(
                            choices=workspace.list_presets(),
                            label="Saved presets",
                        )
                        pr_refresh = gr.Button("Refresh list", size="sm")
                        with gr.Row():
                            pr_name = gr.Textbox(label="New preset name", scale=2)
                            pr_save = gr.Button("Save current settings", variant="primary", scale=1)
                        with gr.Row():
                            pr_load = gr.Button("Load selected", variant="secondary", scale=1)
                            pr_delete = gr.Button("Delete selected", variant="stop", scale=1)
                        pr_status = gr.Markdown()
                pr_refresh.click(fn=ui_refresh_presets, outputs=pr_list)
                pr_load.click(
                    fn=ui_load_preset,
                    inputs=pr_list,
                    outputs=[
                        pr_status,
                        gen_model, gen_quality, gen_seed, gen_tex_size, gen_rembg,
                        pp_repair, pp_cleanup, pp_decimate, pp_target_faces,
                        pp_symmetry, pp_unwrap, pp_normal_bake, pp_normal_format,
                        pp_albedo_bake, pp_material_bake, pp_ao, pp_ao_quality,
                        pp_inpaint, pp_lods, pp_collision, pp_pivot, pp_scale_m,
                        rig_type, rig_seed, rig_format,
                        ex_engine, ex_type,
                    ],
                )
                pr_save.click(
                    fn=ui_save_preset,
                    inputs=[
                        pr_name,
                        gen_model, gen_quality, gen_seed, gen_tex_size, gen_rembg,
                        pp_repair, pp_cleanup, pp_decimate, pp_target_faces,
                        pp_symmetry, pp_unwrap, pp_normal_bake, pp_normal_format,
                        pp_albedo_bake, pp_material_bake, pp_ao, pp_ao_quality,
                        pp_inpaint, pp_lods, pp_collision, pp_pivot, pp_scale_m,
                        rig_type, rig_seed, rig_format,
                        ex_engine, ex_type,
                    ],
                    outputs=[pr_status, pr_list],
                )
                pr_delete.click(fn=ui_delete_preset, inputs=pr_list, outputs=[pr_status, pr_list])

            # ============ Tab 6: Diagnostics (hidden in prod, useful now) =
            with gr.Tab("Diagnostics", id=99):
                gr.Markdown(
                    "### Milestone 1 — Foundation check\n"
                    "Verify the Space environment is working correctly before "
                    "building out the pipeline."
                )
                with gr.Row():
                    with gr.Column():
                        diag_btn = gr.Button("🧪 Run GPU smoke test", variant="primary")
                        diag_out = gr.Markdown()
                    with gr.Column():
                        gr.Markdown("**Workspace state:**")
                        diag_state = gr.JSON(value=workspace.get_state().to_dict())
                        diag_refresh = gr.Button("Refresh state", size="sm")
                diag_btn.click(fn=zerogpu_smoke_test, outputs=diag_out)
                diag_refresh.click(
                    fn=lambda: workspace.get_state().to_dict(),
                    outputs=diag_state,
                )

        # --- Persistent right-side viewer + asset summary ------------------
        gr.Markdown("---")
        with gr.Row():
            with gr.Column(scale=2):
                viewer = gr.Model3D(
                    label="3D viewer",
                    value=ui_helpers.get_viewer_model_path(),
                    clear_color=[0.1, 0.1, 0.12, 1.0],
                    height=500,
                )
            with gr.Column(scale=1):
                summary = gr.Markdown(
                    value=ui_helpers.get_asset_summary(),
                    elem_classes=["asset-summary"],
                )
                refresh_summary = gr.Button("🔄 Refresh viewer", size="sm")
                refresh_summary.click(
                    fn=lambda: (
                        ui_helpers.get_viewer_model_path(),
                        ui_helpers.get_asset_summary(),
                    ),
                    outputs=[viewer, summary],
                )

        # --- Global status bar --------------------------------------------
        status_bar = gr.Markdown(
            value=ui_helpers.get_status_bar(),
            elem_classes=["status-bar"],
        )

        # --- Wire generate button now that viewer is in scope ---------------
        # handle_generate yields (status_str, glb_path_or_None) tuples so that
        # the viewer updates the instant the final mesh is ready — no separate
        # state lookup needed.
        _gen_event = gen_btn.click(
            fn=handle_generate,
            inputs=[gen_images, gen_model, gen_quality, gen_seed,
                    gen_steps, gen_octree, gen_tex_size, gen_symmetry,
                    gen_rembg],
            outputs=[gen_status, viewer],
        )

        # --- Global refresh: every pipeline action updates summary + status bar.
        # For gen_btn the viewer is already updated directly above; for the rest
        # we include the viewer in _global_refresh so post-process / rig / export
        # results also appear.
        def _global_refresh():
            return (
                ui_helpers.get_viewer_model_path(),
                ui_helpers.get_asset_summary(),
                ui_helpers.get_status_bar(),
            )

        def _summary_refresh():
            return ui_helpers.get_asset_summary(), ui_helpers.get_status_bar()

        # gen_btn: viewer already updated by direct output; only refresh metadata.
        _gen_event.then(fn=_summary_refresh, outputs=[summary, status_bar])

        # Other processing buttons: refresh viewer + metadata after completion.
        for _ev in (_pp_event, _rig_event, _ex_event):
            _ev.then(fn=_global_refresh, outputs=[viewer, summary, status_bar])

        # Utility buttons: refresh immediately on click.
        for _btn in (diag_btn, diag_refresh, refresh_summary):
            _btn.click(fn=_global_refresh, outputs=[viewer, summary, status_bar])

    return demo


# ---------------------------------------------------------------------------
# Entrypoint
# ---------------------------------------------------------------------------
# On HF Spaces, app.py is executed directly. We construct the demo and call
# .launch() at module level. The HF Space runtime handles all networking;
# we just need to bind to 0.0.0.0:7860.

workspace.ensure_dirs()
demo = build_ui()
demo.queue(default_concurrency_limit=1).launch(
    server_name="0.0.0.0",
    server_port=7860,
    show_error=True,
    # Allow Gradio to serve files from the workspace directory so that
    # gr.Model3D can display generated GLBs without a 403 error.
    allowed_paths=[str(workspace.WORKSPACE)],
    # Gradio 6: show_api removed. Use footer_links instead.
    footer_links=["gradio", "settings"],
    theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"),
    css=CUSTOM_CSS,
)