import json
import os
from pathlib import Path

import numpy as np
from simpler_env.evaluation.maniskill2_evaluator import maniskill2_evaluator

# from IPython import embed; embed()
from examples.SimplerEnv.eval_files.bridgert_ik_trace import install_ik_trace
from examples.SimplerEnv.eval_files.custom_argparse import get_args
from examples.SimplerEnv.eval_files.model2simpler_interface import ModelClient


def start_debugpy_once():
    import debugpy

    if getattr(start_debugpy_once, "_started", False):
        return
    debugpy.listen(("0.0.0.0", 10092))
    print("🔍 Waiting for VSCode attach on 0.0.0.0:10092 ...")
    debugpy.wait_for_client()
    start_debugpy_once._started = True


if __name__ == "__main__":
    args = get_args()

    os.environ["DISPLAY"] = ""
    # prevent a single jax process from taking up all the GPU memory
    os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

    install_ik_trace()
    if os.getenv("DEBUG", False):
        start_debugpy_once()
    model = ModelClient(
        policy_ckpt_path=args.ckpt_path,  # to get unnormalization stats
        policy_setup=args.policy_setup,
        port=args.port,
        action_scale=args.action_scale,
        cfg_scale=1.5,  # cfg from 1.5 to 7 also performs well
    )

    # policy model creation; update this if you are using a new policy model
    # run real-to-sim evaluation
    success_arr = maniskill2_evaluator(model, args)
    print(args)
    print(" " * 10, "Average success", np.mean(success_arr))
    if args.result_json:
        result = {
            "schema_version": 2,
            "env_name": args.env_name,
            "episode_start": args.obj_episode_range[0],
            "episode_end": args.obj_episode_range[1],
            "episode_indices": list(range(*args.obj_episode_range)),
            "checkpoint": str(Path(args.ckpt_path).resolve()),
            "server_metadata": getattr(model, "server_metadata", {}),
            "successes": np.asarray(success_arr, dtype=bool).tolist(),
        }
        destination = Path(args.result_json)
        destination.parent.mkdir(parents=True, exist_ok=True)
        temporary = destination.with_suffix(destination.suffix + ".tmp")
        temporary.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
        os.replace(temporary, destination)