"""Run a registered YAM task: build scene -> solve -> evaluate -> video.

    python scripts/yam_task.py --list
    python scripts/yam_task.py --task grape_box --seed 0
    python scripts/yam_task.py --task pot_dual_lift --seed 4 --no-randomize

One Python file per task under source/bimanual/yam/tasks/, in the shape ManiSkill uses: a
registered class with _load_scene / _initialize_episode / solve / evaluate.
"""
import argparse, json, os, sys
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO/"source"))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--task")
    ap.add_argument("--list", action="store_true")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--no-randomize", action="store_true")
    ap.add_argument("--episode", type=int, default=-1)
    ap.add_argument("--video", default="")
    ap.add_argument("--set", action="append", default=[], metavar="KEY=VALUE",
                    help="override a task parameter, on top of tasks/configs/<task>.yaml")
    ap.add_argument("--first-frame", default="",
                    help="build the scene, save ONE frame here, and exit without solving")
    args, unknown = ap.parse_known_args()

    from bimanual.yam import registry
    import bimanual.yam.tasks  # noqa: F401  (registers every task)

    if args.list or not args.task:
        rows = registry.list_tasks()
        w = max([len(r["name"]) for r in rows]+[4])
        print(f"{'name'.ljust(w)}  {'class'.ljust(20)}  tags")
        for r in rows:
            print(f"{r['name'].ljust(w)}  {r['class'].ljust(20)}  {r['tags']}")
        print(f"\n{len(rows)} task(s) registered")
        return 0

    cls = registry.REGISTRY.get(args.task)
    if cls is None:
        raise SystemExit(f"unknown task {args.task!r}. Registered: {sorted(registry.REGISTRY)}")

    # YAML + --set have to be resolved BEFORE env_vars(): which USDs get registered and how hard
    # the gripper clamps are decided before the simulator starts, so a config that changed
    # rw_objects would otherwise be read too late to have any effect.
    from bimanual.yam.envs import config as taskcfg
    overrides = dict(taskcfg.load(args.task))
    overrides.update(taskcfg.parse_cli(args.set))

    # asset registration and clamp force must be set BEFORE the simulator starts
    for k, v in cls.env_vars(overrides).items():
        os.environ[k] = v
        print(f"[run] {k}={v}", flush=True)
    os.environ.setdefault("OMNI_KIT_ACCEPT_EULA", "YES")

    from isaaclab.app import AppLauncher
    ap2 = argparse.ArgumentParser(); AppLauncher.add_app_launcher_args(ap2)
    largs = ap2.parse_args(unknown); largs.headless = True; largs.enable_cameras = True
    app = AppLauncher(largs).app

    import gymnasium as gym
    import bimanual.tasks.manager_based.yam  # noqa: F401
    from isaaclab_tasks.utils import parse_env_cfg

    cfg = parse_env_cfg("Template-YAM-Play-v0", device="cuda:0", num_envs=1)
    # one long scripted sequence: the 12 s episode limit would auto-reset mid-run and snap the
    # arm back to its home joints
    cfg.episode_length_s = 1.0e6
    try:
        cfg.terminations.time_out = None
    except Exception as e:
        print("[run] time_out disable failed:", e)
    try:
        cfg.viewer.eye = tuple(cls.viewer_eye); cfg.viewer.lookat = tuple(cls.viewer_lookat)
        cfg.viewer.resolution = (720, 540)
    except Exception as e:
        print("[run] viewer cfg:", e)

    env = gym.make("Template-YAM-Play-v0", cfg=cfg, render_mode="rgb_array")
    env.reset()
    origin = env.unwrapped.scene.env_origins[0].cpu().numpy()

    task = cls(seed=args.seed, randomize=not args.no_randomize, episode=args.episode,
               video=args.video, overrides=taskcfg.parse_cli(args.set))
    task.build(env, origin)
    if args.first_frame:
        # Scene-only render: building is the expensive-but-short part, solving is the long part.
        # This makes a 30-task overview minutes rather than hours.
        import imageio.v2 as _iio
        # The renderer needs warm-up: the frames right after a reset come back empty or with the
        # wrong camera, so a single render() straight after build() writes a black image (all 29
        # tiles of the first overview were black). Step AND render repeatedly, keep the last.
        frame = None
        for _ in range(60):
            task.step()
            img = env.render()
            if img is not None:
                frame = img
        os.makedirs(os.path.dirname(args.first_frame) or ".", exist_ok=True)
        _iio.imwrite(args.first_frame, frame)
        print(f"[run] first frame -> {args.first_frame} {frame.shape}", flush=True)
        env.close(); app.close()
        print("YAM_TASK_DONE", flush=True)
        return 0
    # Re-baseline the joints RIGHT BEFORE solving. Recording them at build time counted the
    # fixture's own settling as progress: switch_toggle "passed" on 0.44 rad of gravity and
    # 0.015 rad of actual robot work.
    task._start_joints = task._joint_state()
    info = task.solve()
    print(f"[run] solver: {json.dumps(info, default=str)}", flush=True)
    result = task.evaluate()
    task.save_video(args.video or None)
    env.close(); app.close()
    print("YAM_TASK_DONE", flush=True)
    return 0 if result.get("success") else 1


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