#!/usr/bin/env python3
"""Fill missing Blender multi-view renders under work_dirs/eval_results/<save>/flux_edit."""

from __future__ import annotations

import argparse
import importlib.util
import os
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path


def load_render_fn(repo_root: Path):
    repo_root = repo_root.resolve()
    repo_str = str(repo_root)
    if repo_str not in sys.path:
        sys.path.insert(0, repo_str)
    os.chdir(repo_root)
    target = repo_root / "eval_3d_editing.py"
    spec = importlib.util.spec_from_file_location("eval_3d_editing", str(target))
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Cannot load {target}")
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod._render


def collect_tasks(flux_root: Path) -> list[Path]:
    out: list[Path] = []
    if not flux_root.exists():
        return out
    for task_dir in sorted(p for p in flux_root.iterdir() if p.is_dir()):
        glb = task_dir / "edit.glb"
        transforms = task_dir / "render" / "transforms.json"
        if glb.exists() and not transforms.exists():
            out.append(glb)
    return out


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo_root", required=True)
    parser.add_argument("--eval_results_dir", required=True)
    parser.add_argument("--blender_path", default="/opt/blender-4.2.19-linux-x64/blender")
    parser.add_argument("--engine", default="CYCLES", choices=["CYCLES", "BLENDER_EEVEE_NEXT"])
    parser.add_argument("--total_render_view_num", type=int, default=16)
    parser.add_argument("--eval_view_num", type=int, default=16)
    parser.add_argument("--render_timeout_seconds", type=int, default=600)
    parser.add_argument("--debug", action="store_true")
    parser.add_argument(
        "--gpu_ids",
        nargs="+",
        type=int,
        default=[0],
        help="Physical GPU ids to round-robin for CYCLES CUDA_VISIBLE_DEVICES",
    )
    parser.add_argument("--max_workers", type=int, default=4)
    args = parser.parse_args()

    repo_root = Path(args.repo_root).resolve()
    eval_dir = Path(args.eval_results_dir).resolve()
    flux_root = eval_dir / "flux_edit"
    render_fn = load_render_fn(repo_root)
    tasks = collect_tasks(flux_root)
    print(f"{eval_dir}: missing renders {len(tasks)} / flux_edit tasks", flush=True)
    if not tasks:
        return

    def work(idx: int, glb: Path) -> tuple[str, bool]:
        gpu = args.gpu_ids[idx % len(args.gpu_ids)]
        try:
            render_fn(
                str(glb),
                engine=args.engine,
                cuda_idx=int(gpu),
                total_render_view_num=args.total_render_view_num,
                eval_view_num=args.eval_view_num,
                BLENDER_EEVEE_NEXT_on_0_device=False,
                timeout_seconds=args.render_timeout_seconds,
                no_norm_scene=True,
                debug=args.debug,
                blender_path=args.blender_path,
            )
            ok = (glb.parent / "render" / "transforms.json").exists()
            return str(glb), ok
        except Exception as exc:
            return f"{glb} ({exc})", False

    failed = 0
    with ThreadPoolExecutor(max_workers=args.max_workers) as ex:
        futs = [ex.submit(work, i, glb) for i, glb in enumerate(tasks)]
        for fut in as_completed(futs):
            msg, ok = fut.result()
            if not ok:
                failed += 1
                print(f"FAIL {msg}", flush=True)
    print(f"done failed={failed}/{len(tasks)}", flush=True)
    if failed:
        raise SystemExit(1)


if __name__ == "__main__":
    main()