# -*- coding: utf-8 -*-
import json
import os
import argparse
from tqdm import tqdm
import chardet

# 預定義標籤類別與對應編號
PREDEFINED_CLASSES = {
    "Edge bruises": 0
}

def read_json_file(file_path):
    with open(file_path, "rb") as f:
        raw_data = f.read()
    detected_encoding = chardet.detect(raw_data)['encoding']
    with open(file_path, "r", encoding=detected_encoding, errors="ignore") as f:
        return json.load(f)

def convert_label_json(json_dir, classes):
    json_paths = [p for p in os.listdir(json_dir) if p.endswith('.json')]
    for json_path in tqdm(json_paths, desc="Converting JSON files"):
        json_path_full = os.path.join(json_dir, json_path)
        try:
            json_dict = read_json_file(json_path_full)
            h, w = json_dict['imageHeight'], json_dict['imageWidth']
        except (KeyError, json.JSONDecodeError) as e:
            print(f"Error reading file '{json_path}': {e}")
            continue

        # Save txt path in the same directory as JSON file
        txt_path = os.path.join(json_dir, json_path.replace('.json', '.txt'))
        with open(txt_path, 'w') as txt_file:
            for shape_dict in json_dict['shapes']:
                try:
                    label = shape_dict['label']
                    label_index = classes.get(label, None)  # 查找標籤對應的編號
                    if label_index is None:
                        print(f"Label '{label}' not in predefined classes.")
                        continue
                    points = shape_dict['points']
                    if shape_dict['shape_type'] == 'rectangle':
                        x1, y1 = points[0]
                        x2, y2 = points[1]
                        center_x = (x1 + x2) / 2 / w
                        center_y = (y1 + y2) / 2 / h
                        width = abs(x2 - x1) / w
                        height = abs(y2 - y1) / h
                    else:
                        print(f"Unsupported shape type: {shape_dict['shape_type']}")
                        continue
                except (ValueError, KeyError) as e:
                    print(f"Error processing shape in file '{json_path}': {e}")
                    continue

                label_str = f"{label_index} {center_x} {center_y} {width} {height}\n"
                txt_file.writelines(label_str)

def extract_labels(json_dir):
    labels = set()
    json_paths = [p for p in os.listdir(json_dir) if p.endswith('.json')]
    print(f"Found {len(json_paths)} JSON files to process.")
    for json_path in json_paths:
        json_path_full = os.path.join(json_dir, json_path)
        try:
            json_dict = read_json_file(json_path_full)
            for shape_dict in json_dict['shapes']:
                labels.add(shape_dict['label'])
        except (KeyError, json.JSONDecodeError) as e:
            print(f"Error reading file '{json_path}': {e}")
    print(f"Extracted labels: {labels}")
    return list(labels)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Convert JSON annotations to TXT format')
    parser.add_argument('--json-dir', type=str, default='./img0204/P/', help='Directory containing JSON files')
    args = parser.parse_args()

    json_dir = args.json_dir

    # Use predefined classes
    predefined_classes = PREDEFINED_CLASSES

    # Perform conversion
    convert_label_json(json_dir, predefined_classes)