#!/usr/bin/env python3
"""End-to-end image -> textured GLB, fully local on Apple Silicon (MLX).

Chains Hunyuan3D 2.1 Stage 1 (shape) + Stage 2 (PBR texture) into one command.
Distilled from tests/test_stage1_to_stage2.py, plus CLI flags, stage timings,
and MLX peak-memory reporting for MODELBEAST job logs.

    python generate_e2e.py IMAGE --output STEM [flags]

Produces (STEM given without extension):
    STEM_shape.glb   Stage 1 mesh (untextured)   [omitted with --no-texture]
    STEM.glb         final PBR-textured mesh      [or the shape mesh if --no-texture]
"""
import argparse
import gc
import os
import sys
import time
from pathlib import Path

REPO = Path(__file__).resolve().parent


def _mlx_mem(fn_names):
    import mlx.core as mx
    candidates = []
    for name in fn_names:
        candidates.append(getattr(mx, name, None))
        metal = getattr(mx, "metal", None)
        if metal is not None:
            candidates.append(getattr(metal, name, None))
    for fn in candidates:
        if fn is None:
            continue
        try:
            return fn()
        except Exception:
            continue
    return None


def peak_gb():
    v = _mlx_mem(["get_peak_memory"])
    return v / 1e9 if v is not None else float("nan")


def reset_peak():
    _mlx_mem(["reset_peak_memory"])


def main():
    ap = argparse.ArgumentParser(description="Hunyuan3D 2.1 MLX image->textured GLB")
    ap.add_argument("image", help="input reference image (ideally a bg-removed cutout)")
    ap.add_argument("--output", required=True, help="output path stem (extension ignored)")
    # Stage 1 (shape)
    ap.add_argument("--steps", type=int, default=50, help="Stage 1 denoising steps")
    ap.add_argument("--guidance", type=float, default=7.5, help="Stage 1 CFG scale")
    ap.add_argument("--octree-resolution", type=int, default=256, help="Stage 1 marching-cubes grid")
    # Stage 2 (texture)
    ap.add_argument("--max-num-view", type=int, default=6, help="Stage 2 baked views")
    ap.add_argument("--view-resolution", type=int, default=512, help="Stage 2 per-view diffusion res")
    ap.add_argument("--texture-size", type=int, default=2048, help="Stage 2 UV atlas res (2048|4096)")
    ap.add_argument("--remesh-faces", type=int, default=40000, help="Stage 2 bake-mesh decimation target")
    # Common
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--no-texture", action="store_true", help="Stage 1 only (geometry, much faster)")
    ap.add_argument("--weights", default="dgrauet/hunyuan3d-2.1-mlx", help="HF repo id or local dir")
    a = ap.parse_args()

    # Resolve I/O to absolute BEFORE chdir; the pipeline reads relative config
    # paths (hy3dpaint/cfgs/...), so it must run from the repo root.
    image_path = str(Path(a.image).resolve())
    out_stem = Path(a.output).resolve()
    out_stem.parent.mkdir(parents=True, exist_ok=True)
    if not os.path.exists(image_path):
        print(f"ERROR: image not found: {image_path}", flush=True)
        sys.exit(1)

    os.environ["HY3D_REMESH_FACES"] = str(a.remesh_faces)
    os.environ.setdefault("HF_HUB_DISABLE_XET", "1")

    os.chdir(REPO)
    sys.path.insert(0, str(REPO / "hy3dshape"))
    sys.path.insert(0, str(REPO / "hy3dpaint"))

    # ---- Stage 1: image -> mesh ----
    print(f"[Stage 1] shape generation from {image_path}", flush=True)
    from hy3dshape.pipeline_mlx import ShapePipeline

    reset_peak()
    t0 = time.time()
    shape_pipe = ShapePipeline.from_pretrained(a.weights)
    print(f"  pipeline ready in {time.time() - t0:.1f}s", flush=True)

    t0 = time.time()
    mesh = shape_pipe(
        image_path,
        num_inference_steps=a.steps,
        guidance_scale=a.guidance,
        octree_resolution=a.octree_resolution,
        seed=a.seed,
    )
    dt1 = time.time() - t0
    nverts, nfaces = len(mesh.vertices), len(mesh.faces)
    print(f"  shape in {dt1:.1f}s ({nverts} verts, {nfaces} faces), "
          f"peak {peak_gb():.1f}GB", flush=True)

    if a.no_texture:
        final = out_stem.with_suffix(".glb")
        mesh.export(str(final))
        print(f"[done] {final}  (geometry only, {dt1:.0f}s)", flush=True)
        return

    shape_glb = str(out_stem.parent / (out_stem.name + "_shape.glb"))
    mesh.export(shape_glb)
    print(f"  saved shape: {shape_glb}", flush=True)

    # Free the DiT before loading the paint stack (both are memory-heavy)
    del shape_pipe, mesh
    gc.collect()
    reset_peak()

    # ---- Stage 2: mesh + reference -> textured GLB ----
    print("[Stage 2] PBR texture synthesis", flush=True)
    from textureGenPipeline_mlx import (
        Hunyuan3DPaintConfigMLX,
        Hunyuan3DPaintPipelineMLX,
    )

    cfg = Hunyuan3DPaintConfigMLX(max_num_view=a.max_num_view, resolution=a.view_resolution)
    cfg.texture_size = a.texture_size
    cfg.mlx_seed = a.seed

    t0 = time.time()
    paint_pipe = Hunyuan3DPaintPipelineMLX(cfg)
    print(f"  pipeline ready in {time.time() - t0:.1f}s", flush=True)

    tex_obj = str(out_stem.with_suffix(".obj"))
    t0 = time.time()
    paint_pipe(
        mesh_path=shape_glb,
        image_path=image_path,
        output_mesh_path=tex_obj,
        use_remesh=True,
        save_glb=True,
    )
    dt2 = time.time() - t0

    final = out_stem.with_suffix(".glb")
    if not final.exists():
        print(f"ERROR: expected textured GLB not produced at {final}", flush=True)
        sys.exit(2)
    size_mb = final.stat().st_size / 1024 / 1024
    print(f"  texture in {dt2:.1f}s, peak {peak_gb():.1f}GB", flush=True)
    print(f"[done] {final}  ({size_mb:.1f} MB)  "
          f"shape={dt1:.0f}s tex={dt2:.0f}s total={dt1 + dt2:.0f}s", flush=True)


if __name__ == "__main__":
    main()