import os
import json
import xml.etree.ElementTree as ET
from xml.dom import minidom
from tqdm import tqdm
import re  # 引入正则表达式库

# --- 配置参数 ---
# 1. 原始JSON文件所在的文件夹路径
json_folder = '../label/up_json'  # 示例路径，请修改为您的JSON文件夹

# 2. 原始图片文件所在的文件夹路径 (用于获取图片尺寸)
image_folder = '../label/up'  # 示例路径，请修改为您的图片文件夹

# 3. 生成的XML文件要保存的文件夹路径
output_xml_folder = '../label/up_xml'

# 4. 您要检测的目标类别名称 (对应 label "3")
class_name_for_label_3 = "Space weld workpiece"  # 这是您XML示例中的名称

# 5. 分组的大小
group_size = 5


# --- 配置结束 ---


def create_xml_annotation(image_info, objects_info):
    """
    根据传入的信息生成XML树对象
    :param image_info: 包含图片文件名、尺寸等信息的字典
    :param objects_info: 包含多个物体信息的列表，每个物体是一个字典
    :return: XML ElementTree对象
    """
    # 创建根节点
    annotation = ET.Element('annotation')

    # 子节点 - folder
    folder = ET.SubElement(annotation, 'folder')
    folder.text = 'JPEGImages'

    # 子节点 - filename
    filename_node = ET.SubElement(annotation, 'filename')
    filename_node.text = image_info['filename']

    # 子节点 - path (路径通常不那么重要，但最好有一个)
    path = ET.SubElement(annotation, 'path')
    # 路径指向JPEGImages文件夹
    image_path_in_voc = os.path.join('..', 'JPEGImages', image_info['filename'])
    path.text = image_path_in_voc

    # 子节点 - source
    source = ET.SubElement(annotation, 'source')
    database = ET.SubElement(source, 'database')
    database.text = 'Unknown'

    # 子节点 - size
    size = ET.SubElement(annotation, 'size')
    width = ET.SubElement(size, 'width')
    width.text = str(image_info['width'])
    height = ET.SubElement(size, 'height')
    height.text = str(image_info['height'])
    depth = ET.SubElement(size, 'depth')
    depth.text = str(image_info.get('depth', 3))

    # 子节点 - segmented
    segmented = ET.SubElement(annotation, 'segmented')
    segmented.text = '0'

    # 为每个物体添加 object 节点
    for obj in objects_info:
        object_node = ET.SubElement(annotation, 'object')
        name = ET.SubElement(object_node, 'name')
        name.text = obj['name']
        pose = ET.SubElement(object_node, 'pose')
        pose.text = 'Unspecified'
        truncated = ET.SubElement(object_node, 'truncated')
        truncated.text = '0'
        difficult = ET.SubElement(object_node, 'difficult')
        difficult.text = '0'
        bndbox = ET.SubElement(object_node, 'bndbox')
        xmin = ET.SubElement(bndbox, 'xmin')
        xmin.text = str(int(obj['xmin']))
        ymin = ET.SubElement(bndbox, 'ymin')
        ymin.text = str(int(obj['ymin']))
        xmax = ET.SubElement(bndbox, 'xmax')
        xmax.text = str(int(obj['xmax']))
        ymax = ET.SubElement(bndbox, 'ymax')
        ymax.text = str(int(obj['ymax']))

    return annotation


def prettify_xml(elem):
    """
    格式化XML输出，使其更易读
    """
    rough_string = ET.tostring(elem, 'utf-8')
    reparsed = minidom.parseString(rough_string)
    return reparsed.toprettyxml(indent="  ")


def main():
    if not os.path.exists(output_xml_folder):
        os.makedirs(output_xml_folder)
        print(f"创建输出文件夹: {output_xml_folder}")

    json_files = sorted([f for f in os.listdir(json_folder) if f.endswith('.json')])

    print(f"找到 {len(json_files)} 个JSON文件，开始转换...")

    for json_file in tqdm(json_files, desc="处理JSON文件"):
        base_name = os.path.splitext(json_file)[0]

        # 使用正则表达式匹配前缀和数字
        match = re.match(r'([a-zA-Z]+)(\d+)', base_name)

        # 1. 检查当前文件是否是一个分组的起始文件
        is_group_start_file = False
        if match:
            num = int(match.group(2))
            # 如果数字是 1, 6, 11, ... 这样的，就认为是起始文件
            if (num - 1) % group_size == 0:
                is_group_start_file = True
        else:
            # 如果文件名不符合 l1, r5 这种格式，我们认为它是“普通”文件，自己就是一个组
            is_group_start_file = True

        if not is_group_start_file:
            # 如果不是起始文件（如l2, l3...），则跳过，因为它的标注已由l1处理
            continue

        # --- 是起始文件，处理这个分组 ---
        json_path = os.path.join(json_folder, json_file)

        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)

        # 2. 从起始文件中提取所有符合条件的标注对象
        objects_to_write = []
        for shape in data.get('shapes', []):
            if shape.get('label') == '1' and shape.get('shape_type') == 'rectangle':
                points = shape.get('points', [])
                if len(points) == 2:
                    x_coords = sorted([p[0] for p in points])
                    y_coords = sorted([p[1] for p in points])
                    objects_to_write.append({
                        'name': class_name_for_label_3,
                        'xmin': x_coords[0], 'ymin': y_coords[0],
                        'xmax': x_coords[1], 'ymax': y_coords[1],
                    })

        if not objects_to_write:
            continue

        # 3. 确定该标注要应用到哪些图片上
        target_image_names = []
        if match:
            # 文件名符合 l1, r6 等格式
            prefix = match.group(1)
            start_num = int(match.group(2))
            for i in range(group_size):
                # 假设图片格式为 .jpg
                target_image_names.append(f"{prefix}{start_num + i}.jpg")
        else:
            # 普通文件，只应用到同名文件
            # 假设图片格式为 .jpg
            target_image_names.append(f"{base_name}.jpg")

        # 4. 为分组内的每个目标图片生成XML文件
        for image_name in target_image_names:
            image_path = os.path.join(image_folder, image_name)
            if not os.path.exists(image_path):
                print(f"\n警告：找不到图片 '{image_name}'，跳过生成其XML文件。")
                continue

            # 使用JSON中的尺寸信息
            image_info = {'filename': image_name, 'width': data['imageWidth'], 'height': data['imageHeight']}

            xml_tree = create_xml_annotation(image_info, objects_to_write)
            xml_string = prettify_xml(xml_tree)
            xml_filename = os.path.splitext(image_name)[0] + '.xml'
            output_path = os.path.join(output_xml_folder, xml_filename)

            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(xml_string)

    print("转换完成！所有XML文件已保存在: ", output_xml_folder)


if __name__ == '__main__':
    main()