import os
import bpy
import sys
import json
import time
import mathutils
import numpy as np
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from utils.blender_util import BlenderUtils
from utils.material_util import MaterialUtil

class DataRenderer:
    def __init__(self):
        config = {
            "renderer": {      
                "generate":
                    {
                        "object_dir": "/media/hofee/data/data/scaled_object_meshes",
                        "table_model_path": "/media/hofee/data/data/others/table.obj",
                        "plane_size": 10,
                        "binocular_vision": True,
                        "light_and_camera_config":{
                            "Camera": {
                                "near_plane": 0.01,
                                "far_plane": 5,
                                "fov_vertical": 25,
                                "resolution": [640, 400],
                                "eye_distance": 0.10,
                                "eye_angle": 25
                            },
                            "Light": {
                                "location": (0,0,3.5),
                                "orientation": (0, 0, 0),
                                "power": 150
                            }
                        }
                    }
            }
        }

        self.plane_size = config["renderer"]["generate"]["plane_size"] 
        self.table_model_path = config["renderer"]["generate"]["table_model_path"]
        self.light_and_camera_config = config["renderer"]["generate"]["light_and_camera_config"]
        self.obj_dir = config["renderer"]["generate"]["object_dir"]
        self.binocular_vision = config["renderer"]["generate"]["binocular_vision"]
        self.target_obj = None
        self.random_obj_list = []

        BlenderUtils.setup_scene(self.light_and_camera_config, self.table_model_path, binocular_vision=self.binocular_vision)
        self.table = BlenderUtils.get_obj(BlenderUtils.TABLE_NAME)
        print(BlenderUtils.get_obj(BlenderUtils.CAMERA_OBJECT_NAME))
    
    def reset(self):
        self.target_obj = None
        self.random_obj_list = []
        BlenderUtils.reset_objects_and_platform()
    
    def do_render(self, cam_pose, restore_info, temp_dir):
        
        self.reset()
        start_time = time.time()
        self.restore_scene(restore_info=restore_info)
        end_time = time.time()
        print(f"Time taken for restoring scene: {end_time - start_time} seconds")
        object_name = self.target_obj.name
        temp_file_name = f"tmp"
        if "." in object_name:
            object_name = object_name.split(".")[0]
        start_time = time.time()
        BlenderUtils.set_camera_at(cam_pose)
        BlenderUtils.render_mask(temp_dir, temp_file_name, binocular_vision=self.binocular_vision)
        MaterialUtil.change_object_material(self.target_obj, MaterialUtil.create_normal_material())
        BlenderUtils.render_normal_and_depth(temp_dir, temp_file_name, binocular_vision=self.binocular_vision)
        end_time = time.time()
        print(f"Time taken for rendering: {end_time - start_time} seconds")
        
        BlenderUtils.save_cam_params(temp_dir, temp_file_name, binocular_vision=self.binocular_vision)
        
        depth_dir = os.path.join(temp_dir, "depth")
        for depth_file in os.listdir(depth_dir):
            if not depth_file.endswith(".png"):
                name, _ = os.path.splitext(depth_file)
                file_path = os.path.join(depth_dir, depth_file)
                new_file_path = os.path.join(depth_dir, f"{name}.png")
                os.rename(file_path,new_file_path)
                
    def restore_scene(self, restore_info):

        for obj_name, obj_info in restore_info.items():
            print("restoring", obj_name)
            if obj_name == BlenderUtils.DISPLAY_TABLE_NAME:
                self.restore_display_platform(obj_info)
            else:
                if obj_name == "target_name":
                    continue
                if obj_name == restore_info["target_name"]:
                    obj_mesh_path = BlenderUtils.get_obj_path(self.obj_dir, obj_name)
                    obj = BlenderUtils.load_obj(obj_name, obj_mesh_path)
                    obj.location = mathutils.Vector(obj_info["location"])
                    obj.rotation_euler = mathutils.Vector(obj_info["rotation_euler"])
                    obj.scale = mathutils.Vector(obj_info["scale"])
                    MaterialUtil.change_object_material(obj, MaterialUtil.create_mask_material(color=(0, 1.0, 0)))
                    self.target_obj = obj

    def restore_display_platform(self, platform_info):
        bpy.ops.mesh.primitive_cylinder_add(radius=platform_info["radius"], depth=platform_info["height"])
        platform = bpy.context.selected_objects[-1]
        platform.name = BlenderUtils.DISPLAY_TABLE_NAME
        platform.location = mathutils.Vector(platform_info["location"])

        bpy.ops.rigidbody.object_add()
        bpy.context.object.rigid_body.type = 'PASSIVE'
        MaterialUtil.change_object_material(platform, MaterialUtil.create_mask_material(color=(1.0, 0, 0)))

def main(temp_dir):
    params_data_path = os.path.join(temp_dir, "params.json")
    with open(params_data_path, 'r') as f:
        params_data = json.load(f)
    cam_pose = np.array(params_data["cam_pose"])
    print(cam_pose)
    scene_info_path = os.path.join(params_data["scene_path"], "scene_info.json")
    with open(scene_info_path, 'r') as f:
        scene_info = json.load(f)
    start_time = time.time()
    data_renderer = DataRenderer()
    end_time = time.time()
    print(f"Time taken for initialization: {end_time - start_time} seconds")
    start_time = time.time()
    data_renderer.do_render(cam_pose, scene_info, temp_dir)
    end_time = time.time()
    print(f"Time taken for rendering: {end_time - start_time} seconds")
    depth_dir = os.path.join(temp_dir, "depth")
    for depth_file in os.listdir(depth_dir):
        if not depth_file.endswith(".png"):
            name, _ = os.path.splitext(depth_file)
            file_path = os.path.join(depth_dir, depth_file)
            new_file_path = os.path.join(depth_dir, f"{name}.png")
            os.rename(file_path,new_file_path)
    BlenderUtils.save_cam_params(temp_dir, "tmp", binocular_vision=data_renderer.binocular_vision)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: blender -b -P data_renderer.py -- <temp_dir>")
    else:
        temp_dir = sys.argv[-1]
        main(temp_dir)