import os
import math
import numpy as np
import imageio
import trimesh
import pyrender
from tqdm import tqdm
# os.environ["CUDA_VISIBLE_DEVICES"] = "7"
os.environ['PYOPENGL_PLATFORM'] = 'egl'  # 设置渲染环境为 EGL（无头模式）

def render_video_from_obj(input_obj_path, output_video_path, fps=15, frame_count=60, resolution=(512, 512)):
    """
    Render a rotating 3D model (OBJ file) to a video with RGB and normal map side-by-side.
    
    Args:
        input_obj_path (str): Path to the input OBJ file.
        output_video_path (str): Path to save the output video.
        fps (int): Frames per second for the video.
        frame_count (int): Number of frames in the video.
        resolution (tuple): Resolution of the rendered video (width, height).
        
    Returns:
        str: Path to the output video.
    """
    # 检查输入文件是否存在
    if not os.path.exists(input_obj_path):
        raise FileNotFoundError(f"Input OBJ file not found: {input_obj_path}")

    # 加载3D模型
    scene_data = trimesh.load(input_obj_path)

    # 提取或合并网格
    if isinstance(scene_data, trimesh.Scene):
        mesh_data = trimesh.util.concatenate([geom for geom in scene_data.geometry.values()])
    else:
        mesh_data = scene_data

    # 确保顶点法线存在
    if not hasattr(mesh_data, 'vertex_normals') or mesh_data.vertex_normals is None:
        mesh_data.compute_vertex_normals()

    # 创建 Pyrender 场景并设置背景为白色
    render_scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0])
    mesh = pyrender.Mesh.from_trimesh(mesh_data, smooth=True)
    mesh_node = render_scene.add(mesh)

    # 设置摄像机参数
    camera = pyrender.PerspectiveCamera(yfov=np.deg2rad(30), znear=0.0001, zfar=100000.0)
    camera_pose = np.eye(4)
    camera_pose[2, 3] = 4.0  # 距离模型 20 个单位
    render_scene.add(camera, pose=camera_pose)

    # 添加全局环境光
    ambient_light = np.array([1.0, 1.0, 1.0]) * 2.0
    render_scene.ambient_light = ambient_light

    # 准备法线渲染场景
    normals = mesh_data.vertex_normals.copy()

    # 将法线映射到颜色范围 [0, 255]
    normal_colors = ((normals + 1) / 2 * 255)

    # 创建用于法线渲染的独立网格
    normal_mesh_data = mesh_data.copy()
    normal_mesh_data.visual.vertex_colors = np.hstack(
        [normal_colors, np.full((normals.shape[0], 1), 255, dtype=np.uint8)]  # 添加 Alpha 通道
    )

    # 创建法线渲染场景
    normal_scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0])
    normal_mesh = pyrender.Mesh.from_trimesh(normal_mesh_data, smooth=True)
    normal_mesh_node = normal_scene.add(normal_mesh)
    normal_scene.add(camera, pose=camera_pose)
    normal_scene.ambient_light = ambient_light

    # 初始化渲染器
    r = pyrender.OffscreenRenderer(*resolution)

    # 创建视频写入器
    writer = imageio.get_writer(output_video_path, fps=fps)

    # 渲染每一帧
    try:
        for frame_idx in tqdm(range(frame_count)):
            # 计算旋转角度
            angle = 2 * np.pi * frame_idx / frame_count
            rotation_matrix = np.array([
                [math.cos(angle), 0, math.sin(angle), 0],
                [0, 1, 0, 0],
                [-math.sin(angle), 0, math.cos(angle), 0],
                [0, 0, 0, 1]
            ])

            # 更新模型的姿态
            render_scene.set_pose(mesh_node, rotation_matrix)

            # 渲染 RGB 图像
            color, _ = r.render(render_scene)

            # 更新法线场景的姿态
            normal_scene.set_pose(normal_mesh_node, rotation_matrix)

            # 渲染法线图像
            normal, _ = r.render(normal_scene, flags=pyrender.RenderFlags.FLAT)
        
            # 拼接左右图像
            combined_frame = np.concatenate((color, normal), axis=1)

            # 写入视频帧
            writer.append_data(combined_frame)
    finally:
        # 释放资源
        writer.close()
        r.delete()

    print(f"Rendered video saved to {output_video_path}")
    return output_video_path

if __name__ == '__main__':
    # 示例调用
    input_obj_path = "output/gradio_cache/text_3D/_超级赛亚人_10/rgb_projected.obj"
    output_video_path = "output.mp4"
    render_video_from_obj(input_obj_path, output_video_path)