#!/usr/bin/env python
"""Batch ReconViaGen driver: load the pipeline ONCE, process an object list.

Usage:
  python batch_reconviagen.py --selection SEL.json --inputs DIR --out OUTDIR --views 1|2 [--seed 42]

SEL.json: {"selections": [{"object": NAME, ...}, ...]}
Images:   DIR/<object>_front.png  (+ DIR/<object>_side.png when --views 2)  RGBA crops.
Output:   OUTDIR/<object>.glb  -- byte-for-byte the same pipeline/settings as
          models/ReconViaGen/run_infer.sh -> ReconViaGen/infer.py (seed 42,
          multidiffusion, ss_steps=25, slat_steps=12, simplify=0.95, tex=1024).

Idempotent: existing OUTDIR/<object>.glb is skipped. Per-object try/except.
The script re-execs itself inside the reconviagen conda env (same env as
run_infer.sh) if it is not already running there, so it can be launched with
any python.
"""
import os
import sys

REPO = "/lp-dev/jonghoon/mv-mesh/ReconViaGen"
ENV = "/lp-dev/jonghoon/mv-mesh/envs/reconviagen"
ENV_PY = os.path.join(ENV, "bin", "python")
HF = "/lp-dev/jonghoon/mv-mesh/hf_cache"

# ---- env identical to run_infer.sh -----------------------------------------
_ENV_VARS = {
    "PYTHONUNBUFFERED": "1",
    "OMP_NUM_THREADS": "4",
    "MKL_NUM_THREADS": "4",
    "SPCONV_ALGO": "native",
    "ATTN_BACKEND": "flash-attn",
    # a python parent that already imported numpy/MKL leaks
    # MKL_THREADING_LAYER=INTEL which clashes with libgomp -> force GNU
    "MKL_THREADING_LAYER": "GNU",
    "MKL_SERVICE_FORCE_INTEL": "0",
    "HF_HOME": HF,
    "HUGGINGFACE_HUB_CACHE": HF,
    "HF_HUB_CACHE": HF,
    "TORCH_HOME": os.path.join(HF, "torch"),
}


def _reexec_in_env():
    """Re-exec this script with the reconviagen env's python + env vars."""
    env = dict(os.environ)
    for k, v in _ENV_VARS.items():
        env[k] = v  # these are fixed by run_infer.sh, not user-overridable
    # OMP/MKL thread caps: honour an explicit non-empty user setting
    for k in ("OMP_NUM_THREADS", "MKL_NUM_THREADS"):
        if os.environ.get(k):
            env[k] = os.environ[k]
    # conda activate equivalent (prefix env)
    env["CONDA_PREFIX"] = ENV
    env["PATH"] = os.path.join(ENV, "bin") + os.pathsep + env.get("PATH", "")
    # CUDA_VISIBLE_DEVICES: respect if set, else pick the GPU with most free mem
    if not env.get("CUDA_VISIBLE_DEVICES"):
        try:
            import subprocess
            out = subprocess.check_output(
                ["nvidia-smi", "--query-gpu=index,memory.free",
                 "--format=csv,noheader,nounits"], text=True)
            best = max((l.split(",") for l in out.strip().splitlines()),
                       key=lambda r: int(r[1]))
            env["CUDA_VISIBLE_DEVICES"] = best[0].strip()
        except Exception as e:
            print(f"[batch] (non-fatal) GPU autopick failed: {e}", flush=True)
    env["_RVG_BATCH_INENV"] = "1"
    print(f"[batch] re-exec in {ENV_PY} "
          f"(CUDA_VISIBLE_DEVICES={env.get('CUDA_VISIBLE_DEVICES')})", flush=True)
    os.execve(ENV_PY, [ENV_PY, os.path.abspath(__file__)] + sys.argv[1:], env)


if not os.environ.get("_RVG_BATCH_INENV"):
    if not os.path.isfile(ENV_PY):
        sys.exit(f"ERROR: env python not found: {ENV_PY}")
    _reexec_in_env()

# ---------------------------------------------------------------------------
for _k, _v in _ENV_VARS.items():
    os.environ.setdefault(_k, _v)

import argparse
import json
import time
import traceback

os.chdir(REPO)          # trellis does sys.path.append("wheels/vggt") relative to CWD
sys.path.insert(0, REPO)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--selection", required=True)
    ap.add_argument("--inputs", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--views", type=int, required=True, choices=[1, 2, 4, 8])
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--limit", type=int, default=None,
                    help="only the first N objects (debug)")
    # inference settings -- defaults match infer.py exactly; do not change
    ap.add_argument("--mode", default="multidiffusion",
                    choices=["multidiffusion", "stochastic"])
    ap.add_argument("--ss-steps", type=int, default=25)
    ap.add_argument("--slat-steps", type=int, default=12)
    ap.add_argument("--simplify", type=float, default=0.95)
    ap.add_argument("--texture-size", type=int, default=1024)
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--nshards", type=int, default=1)
    ap.add_argument("--reverse", action="store_true",
                    help="process objects in reverse order (for a convergent "
                         "co-located helper against a forward-order run; both "
                         "skip-existing + atomic-write, so they safely meet in "
                         "the middle)")
    args = ap.parse_args()

    with open(args.selection) as f:
        sel = json.load(f)["selections"]
    objects = [s["object"] for s in sel]
    if args.limit:
        objects = objects[: args.limit]
    if args.nshards > 1:
        objects = objects[args.shard::args.nshards]
    if args.reverse:
        objects = objects[::-1]
    outdir = os.path.abspath(args.out)
    os.makedirs(outdir, exist_ok=True)

    # ---- resolve jobs up front -------------------------------------------
    jobs = []
    for name in objects:
        out = os.path.join(outdir, f"{name}.glb")
        _vtags = {1: ["front"], 2: ["front", "side"],
                  4: ["front", "side", "back", "oside"],
             8: ["front", "side", "back", "oside", "top", "bottom", "top2", "bottom2"]}[args.views]
        imgs = [os.path.join(args.inputs, f"{name}_{t}.png") for t in _vtags]
        jobs.append((name, imgs, out))

    todo = [j for j in jobs if not os.path.isfile(j[2])]
    print(f"[batch] {len(jobs)} object(s), {len(jobs) - len(todo)} already done, "
          f"{len(todo)} to run | views={args.views} seed={args.seed}", flush=True)

    n_ok = n_fail = n_skip = 0
    pipeline = None
    t_load = None

    import torch
    from PIL import Image
    from trellis.utils import postprocessing_utils

    for i, (name, imgs, out) in enumerate(jobs, 1):
        t0 = time.time()
        try:
            if os.path.isfile(out):
                n_skip += 1
                print(f"[{i}/{len(jobs)}] {name} SKIP (exists) {out}", flush=True)
                continue
            for p in imgs:
                if not os.path.isfile(p):
                    raise FileNotFoundError(p)

            if pipeline is None:
                from infer import load_pipeline  # ReconViaGen/infer.py
                tl = time.time()
                pipeline = load_pipeline()
                t_load = time.time() - tl
                print(f"[batch] pipeline ready in {t_load:.1f}s", flush=True)
                t0 = time.time()  # don't bill the load to this object

            inputs_dir = os.path.join(outdir, "inputs")
            os.makedirs(inputs_dir, exist_ok=True)

            images = []
            for p in imgs:
                fed = pipeline.preprocess_image(Image.open(p))
                images.append(fed)
                stem = os.path.splitext(os.path.basename(p))[0]
                fed.save(os.path.join(inputs_dir, f"FED_{stem}.png"))

            outputs, _, _ = pipeline.run(
                image=images,
                seed=args.seed,
                formats=["gaussian", "mesh"],
                preprocess_image=False,
                sparse_structure_sampler_params={
                    "steps": args.ss_steps, "cfg_strength": 7.5,
                    "cfg_interval": [0.6, 1.0], "guidance_rescale": 0.7,
                    "rescale_t": 5.0,
                },
                slat_sampler_params={
                    "steps": args.slat_steps, "cfg_strength": 7.5,
                    "cfg_interval": [0.6, 1.0], "guidance_rescale": 0.5,
                    "rescale_t": 3.0,
                },
                mode=args.mode,
            )
            torch.cuda.empty_cache()
            glb = postprocessing_utils.to_glb(
                outputs["gaussian"][0], outputs["mesh"][0],
                simplify=args.simplify, texture_size=args.texture_size,
                verbose=False,
            )
            tmp = out + ".tmp.glb"
            glb.export(tmp)
            if not os.path.isfile(tmp):
                raise RuntimeError("export produced no file")
            os.replace(tmp, out)
            del outputs, glb, images
            n_ok += 1
            print(f"[{i}/{len(jobs)}] {name} OK {time.time() - t0:.1f}s -> {out}",
                  flush=True)
        except Exception:
            n_fail += 1
            traceback.print_exc()
            print(f"[{i}/{len(jobs)}] {name} FAIL {time.time() - t0:.1f}s", flush=True)
        finally:
            try:
                torch.cuda.empty_cache()
            except Exception:
                pass

    print(f"BATCH DONE ok={n_ok} fail={n_fail} skip={n_skip} total={len(jobs)} "
          f"load={'%.1fs' % t_load if t_load else 'n/a'}", flush=True)
    return 1 if n_fail else 0


if __name__ == "__main__":
    sys.exit(main())