import json
import os

class SceneEnricher:
    def __init__(self, scene_file, scene_ref_file):
        self.scene_file = scene_file
        self.scene_ref_file = scene_ref_file
        self.scene_data = self.load_json(scene_file)
        self.scene_ref_data = self.load_json(scene_ref_file)

    def load_json(self, file_path):
        if os.path.exists(file_path):
            with open(file_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        else:
            print(f"警告: 文件 {file_path} 不存在")
            return {}

    def save_scene(self, output_file=None):
        if not output_file:
            output_file = self.scene_file
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(self.scene_data, f, indent=2, ensure_ascii=False)
        print(f"已保存补全后的scene到: {output_file}")

    def update_points_robots(self):
        if 'points' not in self.scene_data or not self.scene_data['points']:
            print("scene_data['points'] 为空，无法更新robots字段")
            return
        scene_points = {p['name']: p for p in self.scene_ref_data.get('points', [])}
        update_count = 0
        for point in self.scene_data['points']:
            name = point.get('name')
            if name in scene_points:
                robots = scene_points[name].get('robots')
                if robots:
                    point['robots'] = robots
                    update_count += 1
        print(f"已同步 {update_count} 个点的robots字段到输出points中")

    def update_area_points(self):
        if 'areas' not in self.scene_data or not self.scene_data['areas']:
            print("scene_data['areas'] 为空，无法更新区域points")
            return
        scene_areas = {a['id']: a for a in self.scene_ref_data.get('areas', [])}
        update_count = 0
        for area in self.scene_data['areas']:
            area_id = area.get('id')
            area_type = area.get('type')
            if area_type != 1 and area_id in scene_areas:
                scene_points = scene_areas[area_id].get('points')
                if scene_points is not None:
                    area['points'] = scene_points
                    update_count += 1
        print(f"已为 {update_count} 个区域同步scene_file中的points")

    def update_points_with_brand_type(self):
        bin_locations_file = os.path.join('converted', 'binLocationsList.json')
        if not os.path.exists(bin_locations_file):
            print(f"未找到 {bin_locations_file}，跳过brand/type补全")
            return
        with open(bin_locations_file, 'r', encoding='utf-8') as f:
            bin_data = json.load(f)
        scene_points = {p['name']: p for p in self.scene_ref_data.get('points', [])}
        point_brand_type = {}
        for group in bin_data.get('binLocationsList', []):
            for loc in group.get('binLocationList', []):
                point_name = loc.get('pointName')
                scene_point = scene_points.get(point_name)
                if scene_point:
                    brand = scene_point.get('brand')
                    type_ = scene_point.get('type')
                    if brand is not None and type_ is not None:
                        point_brand_type[point_name] = {'brand': brand, 'type': type_}
        update_count = 0
        for point in self.scene_data['points']:
            name = point.get('name')
            if name in point_brand_type:
                brand = point_brand_type[name]['brand']
                type_ = point_brand_type[name]['type']
                point['brand'] = brand
                point['type'] = type_
                update_count += 1
        print(f"已为 {update_count} 个点增加brand/type字段")

    def update_point_types(self):
        if 'points' not in self.scene_data or not self.scene_data['points']:
            print("scene_data['points'] 为空，无法更新点类型")
            return
        scene_points = {p['name']: p for p in self.scene_ref_data.get('points', [])}
        update_count = 0
        for point in self.scene_data['points']:
            if point.get('type') == 1:
                name = point.get('name')
                if name in scene_points:
                    scene_type = scene_points[name].get('type')
                    if scene_type is not None and point['type'] != scene_type:
                        point['type'] = scene_type
                        update_count += 1
        print(f"已同步 {update_count} 个点的type到scene_file中的值")

    def update_routes_pass(self):
        """
        遍历scene_file中routes字段下的路径，将scene_ref_file中的对应路径下的pass的值映射到最终生成的output_file中去
        """
        if 'routes' not in self.scene_data or not self.scene_data['routes']:
            print("scene_data['routes'] 为空，无法更新路径pass字段")
            return
        
        # 建立参考文件中路径的映射关系（根据desc字段）
        scene_ref_routes_by_desc = {}
        for route in self.scene_ref_data.get('routes', []):
            desc = route.get('desc')
            if desc:
                scene_ref_routes_by_desc[desc] = route
        
        print(f"参考文件中找到 {len(self.scene_ref_data.get('routes', []))} 条路径")
        print(f"当前文件中有 {len(self.scene_data['routes'])} 条路径")
        print(f"建立了 {len(scene_ref_routes_by_desc)} 个描述映射关系")
        
        update_count = 0
        for route in self.scene_data['routes']:
            desc = route.get('desc')
            if desc and desc in scene_ref_routes_by_desc:
                ref_route = scene_ref_routes_by_desc[desc]
                ref_pass = ref_route.get('pass')
                if ref_pass is not None:
                    route['pass'] = ref_pass
                    update_count += 1
                    if update_count <= 5:  # 只显示前5个匹配
                        print(f"匹配路径描述: {desc}, pass值: {ref_pass}")
        
        print(f"已同步 {update_count} 条路径的pass字段到输出routes中")

    def enrich(self, output_file=None):
        # 若 areas 为空，则从参考文件补全
        if ('areas' not in self.scene_data or not self.scene_data['areas']) and \
           ('areas' in self.scene_ref_data and self.scene_ref_data['areas']):
            print("scene_data['areas'] 为空，自动从参考文件补全areas字段")
            self.scene_data['areas'] = self.scene_ref_data['areas'].copy()


        # 优先从 scene_file 读取 robotGroups 和 robots 字段，如果为空则从 scene_ref_file 读取补全
        robotGroups = []
        robots = []
        if os.path.exists(self.scene_file):
            with open(self.scene_file, 'r', encoding='utf-8') as f:
                scene_origin = json.load(f)
            robotGroups = scene_origin.get('robotGroups', [])
            robots = scene_origin.get('robots', [])
        # 如果为空则从参考文件补全
        if not robotGroups and 'robotGroups' in self.scene_ref_data:
            robotGroups = self.scene_ref_data['robotGroups']
            print("robotGroups 字段为空，已从参考文件补全")
        if not robots and 'robots' in self.scene_ref_data:
            robots = self.scene_ref_data['robots']
            print("robots 字段为空，已从参考文件补全")
        self.scene_data['robotGroups'] = robotGroups
        self.scene_data['robots'] = robots

        self.update_points_robots()
        self.update_area_points()
        self.update_points_with_brand_type()
        self.update_point_types()
        self.update_routes_pass()
        self.save_scene(output_file)

def main():
    scene_file = "converted_test.scene"
    scene_ref_file = "现场测试地图-0819 (1).scene"
    output_file = "converted_scene.scene"  # 可自定义输出
    enricher = SceneEnricher(scene_file, scene_ref_file)
    print("=== 场景补全器 ===")
    print("执行字段补全...")
    print("-" * 50)
    enricher.enrich(output_file)
    print("补全完成！")

if __name__ == "__main__":
    main()
