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

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

SEL.json: {"selections": [{"object": NAME, ...}, ...]}
Images:   DIR/<object>_front.png  RGBA crops (Cupid uses the existing alpha as-is;
          _has_alpha() -> True so NO rembg re-segmentation is triggered).
Output:   OUTDIR/<object>.glb  -- Cupid's mesh in its OWN canonical model frame,
          exported exactly like Cupid's save_mesh():
          postprocessing_utils.to_glb(gaussian, mesh, simplify=0.95, texture_size=1024).
          The estimated camera pose is intentionally NOT applied: downstream
          metrics/align_baselines.py registers each prediction to the GT canonical
          frame (best-of-24 cube rotation + rigid ICP), identical to RVG/Pixal3D/
          Amodal3R/Hunyuan. So the raw model-frame mesh is the fair input.

Cupid only exposes single-image reconstruction (pipeline.run takes ONE image),
so only --views 1 is supported here.

Idempotent: existing OUTDIR/<object>.glb is skipped. Per-object try/except.
Re-execs itself inside the cupid conda env if not already there.
"""
import os
import sys

ENV = "/lp-dev/jonghoon/mv-mesh/baselines/cupid/env"
ENV_PY = os.path.join(ENV, "bin", "python")
REPO = "/lp-dev/jonghoon/mv-mesh/baselines/cupid/repo"
HF = "/lp-dev/jonghoon/mv-mesh/hf_cache"

_ENV_VARS = {
    "PYTHONUNBUFFERED": "1",
    "OMP_NUM_THREADS": "4",
    "MKL_NUM_THREADS": "4",
    "SPCONV_ALGO": "native",
    "ATTN_BACKEND": "flash_attn",
    "MKL_THREADING_LAYER": "GNU",
    "MKL_SERVICE_FORCE_INTEL": "0",
    "HF_HOME": HF,
    "HUGGINGFACE_HUB_CACHE": HF,
    "HF_HUB_CACHE": HF,
    "TORCH_HOME": "/lp-dev/jonghoon/mv-mesh/torch_hub",
    "CUDA_HOME": "/usr/local/cuda-12.8",
}


def _reexec_in_env():
    env = dict(os.environ)
    for k, v in _ENV_VARS.items():
        env[k] = v
    for k in ("OMP_NUM_THREADS", "MKL_NUM_THREADS"):
        if os.environ.get(k):
            env[k] = os.environ[k]
    env["CONDA_PREFIX"] = ENV
    env["PATH"] = os.path.join(ENV, "bin") + os.pathsep + "/usr/local/cuda-12.8/bin" + os.pathsep + env.get("PATH", "")
    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["_CUPID_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("_CUPID_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)
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, default=1, choices=[1])
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--limit", type=int, default=None)
    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")
    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)
    inputs_dir = os.path.join(outdir, "inputs")
    os.makedirs(inputs_dir, exist_ok=True)

    jobs = []
    for name in objects:
        out = os.path.join(outdir, f"{name}.glb")
        img = os.path.join(args.inputs, f"{name}_front.png")
        jobs.append((name, img, 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)

    import torch
    from PIL import Image
    from cupid.pipelines import Cupid3DPipeline
    from cupid.utils import sample_utils
    from cupid.utils import postprocessing_utils

    n_ok = n_fail = n_skip = 0
    pipeline = None

    for i, (name, img_path, 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)", flush=True)
                continue
            if not os.path.isfile(img_path):
                raise FileNotFoundError(img_path)

            if pipeline is None:
                tl = time.time()
                pipeline = Cupid3DPipeline.from_pretrained("hbb1/Cupid")
                pipeline.cuda()
                print(f"[batch] pipeline ready in {time.time()-tl:.1f}s", flush=True)
                t0 = time.time()

            image = sample_utils.load_image(img_path)  # pad_to_square, keeps RGBA alpha
            outputs = pipeline.run(image, seed=args.seed,
                                   formats=["mesh", "gaussian"])
            # Save the FED (preprocessed, cropped) image Cupid actually consumed
            try:
                proc = pipeline.crop_image(pipeline.preprocess_image(image))
                proc.image.save(os.path.join(inputs_dir, f"FED_{name}_front.png"))
            except Exception:
                pass

            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, image
            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}", flush=True)


if __name__ == "__main__":
    main()