import omni.usd
import omni.replicator.core as rep
from pxr import UsdGeom
from omni.isaac.sensor import Camera
from PIL import Image
import json, os, asyncio

# ==== 配置 ====
sensor_cam_path = "/World/MySensorCamera"
save_dir = "/home/sig/qianluo/omni-3dgs-extension/assets/scripts/test_cap/app_capture/839920/test"
json_path = "/home/sig/qianluo/omni-3dgs-extension/assets/scripts/test_cap/839920_test_gt.json"
os.makedirs(save_dir, exist_ok=True)

stage = omni.usd.get_context().get_stage()

# 创建 Sensor Camera
if not stage.GetPrimAtPath(sensor_cam_path).IsValid():
    cam = Camera(prim_path=sensor_cam_path, frequency=30, resolution=(640, 480))
    cam.initialize()
else:
    cam = Camera(prim_path=sensor_cam_path)
    cam.initialize()

# 让相机进入渲染管线
rp = rep.create.render_product(sensor_cam_path, (640, 480))

# 初始化 JSON
if os.path.exists(json_path):
    with open(json_path, "r") as f:
        dataset = json.load(f)
else:
    dataset = {
        "dataset_metadata": {
            "name": "GVLN",
            "dataset_type": "dataset_type",
            "dataset_description": "dataset_description"
        },
        "scenes": [
            {
                "scene_id": 0,
                "scene_name": "0",
                "samples": [
                    {
                        "trajectory_id": "0",
                        "instructions": [],
                        "points": []
                    }
                ]
            }
        ]
    }

scene_entry = dataset["scenes"][0]
sample_entry = scene_entry["samples"][0]
points_list = sample_entry["points"]


async def capture_pose():
    """GUI + RTX Sensor 相机的异步抓取"""
    # orchestrator 保证真的渲染一帧
    await rep.orchestrator.step_async()

    img = cam.get_rgba()
    if img is None or img.size == 0:
        print("❌ 没有图像")
        return

    # 保存 PNG
    idx = len(points_list)
    img_path = os.path.join(save_dir, f"pose_{idx}.png")
    Image.fromarray(img[:, :, :3]).save(img_path)

    # 获取姿态 & 光学参数
    pos, quat = cam.get_world_pose()
    prim = stage.GetPrimAtPath(sensor_cam_path)
    cam_geom = UsdGeom.Camera(prim)

    focal_length = cam_geom.GetFocalLengthAttr().Get()
    h_aperture = cam_geom.GetHorizontalApertureAttr().Get()
    v_aperture = cam_geom.GetVerticalApertureAttr().Get()
    focus_distance = cam_geom.GetFocusDistanceAttr().Get()
    clipping_range = cam_geom.GetClippingRangeAttr().Get()

    # 更新 JSON 数据
    points_list.append({
        "point": str(idx),
        "position": pos.tolist(),
        "rotation": quat.tolist(),
        "action": [],
        "camera_images": [],
        "focal_length": focal_length,
        "horizontal_aperture": h_aperture,
        "vertical_aperture": v_aperture,
        "focus_distance": focus_distance,
        "clipping_range": list(clipping_range)
    })
    with open(json_path, "w") as f:
        json.dump(dataset, f, indent=2)

    # 控制台日志输出（和之前同步版本一致）
    print(f"✅ 已保存姿态 {idx}")
    print(" Position:", pos.tolist())
    print(" Quaternion:", quat.tolist())
    print(f" 光学参数: focal={focal_length}, Ha={h_aperture}, Va={v_aperture}, "
          f"focusDist={focus_distance}, clipping={list(clipping_range)}")
    print(" 图片文件:", img_path)


# 在 Script Editor 触发
asyncio.ensure_future(capture_pose())