import argparse
import sys
import os
import time

import imageio
import numpy as np
import h5py
import torch

# 先引入 AppLauncher，并尽早实例化 SimulationApp，避免 pxr 未加载
from isaaclab.app import AppLauncher
import cli_args  # isort: skip

# CLI
parser = argparse.ArgumentParser(description="Play an RL agent with multi-cam recording.")
parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
parser.add_argument("--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations.")
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
parser.add_argument("--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point.")
parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
parser.add_argument("--use_pretrained_checkpoint", action="store_true", help="Use the pre-trained checkpoint from Nucleus.")
parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.")
parser.add_argument("--max_steps", type=int, default=None, help="最大步数，达到后提前退出")
cli_args.add_rsl_rl_args(parser)
AppLauncher.add_app_launcher_args(parser)
args_cli, hydra_args = parser.parse_known_args()
if args_cli.video:
    args_cli.enable_cameras = True
# 将 hydra 剩余参数传递
sys.argv = [sys.argv[0]] + hydra_args

# ==== 先实例化 SimulationApp ====
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app

# ==== 之后再 import 依赖 isaac/pxr 的模块 ====
import gymnasium as gym
from rsl_rl.runners import DistillationRunner, OnPolicyRunner
from isaaclab.envs import DirectMARLEnv, DirectMARLEnvCfg, DirectRLEnvCfg, ManagerBasedRLEnvCfg, multi_agent_to_single_agent
from isaaclab.utils.assets import retrieve_file_path
from isaaclab.utils.dict import print_dict
from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg, RslRlVecEnvWrapper, export_policy_as_jit, export_policy_as_onnx
from isaaclab_rl.utils.pretrained_checkpoint import get_published_pretrained_checkpoint

import isaaclab_tasks  # noqa: F401
from isaaclab_tasks.utils import get_checkpoint_path
from isaaclab_tasks.utils.hydra import hydra_task_config
import mindbot.tasks  # noqa: F401

CAM_NAMES = ["cam_head", "cam_chest", "cam_left_hand", "cam_right_hand", "cam_top", "cam_side"]


@hydra_task_config(args_cli.task, args_cli.agent)
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg):
    task_name = args_cli.task.split(":")[-1]
    train_task_name = task_name.replace("-Play", "")

    agent_cfg: RslRlBaseRunnerCfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
    env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs
    env_cfg.seed = agent_cfg.seed
    env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device

    log_root_path = os.path.abspath(os.path.join("logs", "rsl_rl", agent_cfg.experiment_name))
    print(f"[INFO] Loading experiment from directory: {log_root_path}")
    if args_cli.use_pretrained_checkpoint:
        resume_path = get_published_pretrained_checkpoint("rsl_rl", train_task_name)
        if not resume_path:
            print("[INFO] No pre-trained checkpoint for this task.")
            return
    elif args_cli.checkpoint:
        resume_path = retrieve_file_path(args_cli.checkpoint)
    else:
        resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)

    log_dir = os.path.dirname(resume_path)
    env_cfg.log_dir = log_dir

    env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
    if isinstance(env.unwrapped, DirectMARLEnv):
        env = multi_agent_to_single_agent(env)
    env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)

    print(f"[INFO]: Loading model checkpoint from: {resume_path}")
    if agent_cfg.class_name == "OnPolicyRunner":
        runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
    elif agent_cfg.class_name == "DistillationRunner":
        runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
    else:
        raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}")
    runner.load(resume_path)

    policy = runner.get_inference_policy(device=env.unwrapped.device)
    try:
        policy_nn = runner.alg.policy
    except AttributeError:
        policy_nn = runner.alg.actor_critic

    # 导出模型
    export_model_dir = os.path.join(os.path.dirname(resume_path), "exported")
    normalizer = getattr(policy_nn, "actor_obs_normalizer", getattr(policy_nn, "student_obs_normalizer", None))
    export_policy_as_jit(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.pt")
    export_policy_as_onnx(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.onnx")

    # dt = env.unwrapped.step_dt

    # # 录制缓冲
    # cam_buffers = {n: [] for n in CAM_NAMES}
    # joint_log, joint_vel_log, action_log, ts_log = [], [], [], []
    # t0 = time.time()

    # obs = env.get_observations()
    # timestep = 0
    # step_count = 0
    # try:
    #     while simulation_app.is_running():
    #         start_time = time.time()
    #         with torch.inference_mode():
    #             actions = policy(obs)
    #             obs, _, dones, _ = env.step(actions)
    #             policy_nn.reset(dones)

    #         # 相机帧（取 env0，如需全部 env 去掉 [0]）
    #         for name in CAM_NAMES:
    #             if name not in env.unwrapped.scene.sensors:
    #                 continue
    #             cam = env.unwrapped.scene.sensors[name]
    #             rgba = cam.data.output.get("rgba", cam.data.output.get("rgb"))
    #             if rgba is None:
    #                 continue
    #             frame = rgba[0].cpu().numpy()
    #             if frame.shape[-1] == 4:
    #                 frame = frame[..., :3]
    #             cam_buffers[name].append(frame)

    #         # 关节 / 速度 / 动作
    #         robot = env.unwrapped.scene["Mindbot"]
    #         joint_log.append(robot.data.joint_pos.cpu().numpy())
    #         joint_vel_log.append(robot.data.joint_vel.cpu().numpy())
    #         action_log.append(actions.cpu().numpy())
    #         ts_log.append(time.time() - t0)

    #         step_count += 1
    #         if args_cli.max_steps and step_count >= args_cli.max_steps:
    #             break

    #         if args_cli.video:
    #             timestep += 1
    #             if timestep == args_cli.video_length:
    #                 break

    #         sleep_time = dt - (time.time() - start_time)
    #         if args_cli.real_time and sleep_time > 0:
    #             time.sleep(sleep_time)
    # finally:
    #     # 保存 HDF5
    #     h5_path = os.path.join(log_dir, "rollout_multi_cam.h5")
    #     with h5py.File(h5_path, "w") as f:
    #         f.create_dataset("joint_pos", data=np.stack(joint_log), compression="gzip")
    #         f.create_dataset("joint_vel", data=np.stack(joint_vel_log), compression="gzip")
    #         f.create_dataset("actions", data=np.stack(action_log), compression="gzip")
    #         f.create_dataset("timestamps", data=np.array(ts_log))
    #         for name, frames in cam_buffers.items():
    #             if not frames:
    #                 continue
    #             dset = f.create_dataset(f"cams/{name}/rgb", data=np.stack(frames), compression="gzip")
    #             if name in ["cam_head", "cam_chest", "cam_left_hand", "cam_right_hand"]:
    #                 fx, fy, cx, cy = 911.77, 911.5, 624.07, 364.05
    #             else:
    #                 fx, fy, cx, cy = 458.7488, 458.8663, 323.3297, 240.6295
    #             dset.attrs["fx"] = fx
    #             dset.attrs["fy"] = fy
    #             dset.attrs["cx"] = cx
    #             dset.attrs["cy"] = cy
    #             dset.attrs["focal_length"] = 1.93
    #     print(f"[INFO] Saved HDF5 to {h5_path}")

    #     # 可选：单路 MP4
    #     head_frames = cam_buffers["cam_head"]
    #     if head_frames:
    #         fps = int(round(1.0 / dt))
    #         video_path = os.path.join(log_dir, "cam_head.mp4")
    #         imageio.mimsave(video_path, head_frames, fps=fps)
    #         print(f"[INFO] Saved video to {video_path}")

    #     env.close()
    dt = env.unwrapped.step_dt

    # 录制缓冲
    cam_buffers = {n: [] for n in CAM_NAMES}
    joint_log, joint_vel_log, action_log, ts_log = [], [], [], []
    t0 = time.time()

    obs = env.get_observations()
    timestep = 0
    while simulation_app.is_running():
        start_time = time.time()
        with torch.inference_mode():
            actions = policy(obs)
            obs, _, dones, _ = env.step(actions)
            policy_nn.reset(dones)

        # 相机帧（取 env0，如需全部 env 去掉 [0]）
        for name in CAM_NAMES:
            if name not in env.unwrapped.scene.sensors:
                continue
            cam = env.unwrapped.scene.sensors[name]
            rgba = cam.data.output.get("rgba", cam.data.output.get("rgb"))
            if rgba is None:
                continue
            frame = rgba[0].cpu().numpy()
            if frame.shape[-1] == 4:
                frame = frame[..., :3]
            cam_buffers[name].append(frame)

        # 关节 / 速度 / 动作
        robot = env.unwrapped.scene["Mindbot"]
        joint_log.append(robot.data.joint_pos.cpu().numpy())
        joint_vel_log.append(robot.data.joint_vel.cpu().numpy())
        action_log.append(actions.cpu().numpy())
        ts_log.append(time.time() - t0)

        if args_cli.video:
            timestep += 1
            if timestep == args_cli.video_length:
                break

        sleep_time = dt - (time.time() - start_time)
        if args_cli.real_time and sleep_time > 0:
            time.sleep(sleep_time)

    # 保存 HDF5
    h5_path = os.path.join(log_dir, "rollout_multi_cam.h5")
    with h5py.File(h5_path, "w") as f:
        f.create_dataset("joint_pos", data=np.stack(joint_log), compression="gzip")
        f.create_dataset("joint_vel", data=np.stack(joint_vel_log), compression="gzip")
        f.create_dataset("actions", data=np.stack(action_log), compression="gzip")
        f.create_dataset("timestamps", data=np.array(ts_log))
        for name, frames in cam_buffers.items():
            if not frames:
                continue
            dset = f.create_dataset(f"cams/{name}/rgb", data=np.stack(frames), compression="gzip")
            # 内参：按你相机 vector 设置
            if name in ["cam_head", "cam_chest", "cam_left_hand", "cam_right_hand"]:
                fx, fy, cx, cy = 911.77, 911.5, 624.07, 364.05
            else:
                fx, fy, cx, cy = 458.7488, 458.8663, 323.3297, 240.6295
            dset.attrs["fx"] = fx
            dset.attrs["fy"] = fy
            dset.attrs["cx"] = cx
            dset.attrs["cy"] = cy
            dset.attrs["focal_length"] = 1.93
    print(f"[INFO] Saved HDF5 to {h5_path}")

    # 可选：单路 MP4
    head_frames = cam_buffers["cam_head"]
    if head_frames:
        fps = int(round(1.0 / dt))
        video_path = os.path.join(log_dir, "cam_head.mp4")
        imageio.mimsave(video_path, head_frames, fps=fps)
        print(f"[INFO] Saved video to {video_path}")

    env.close()


if __name__ == "__main__":
    main()
    simulation_app.close()