import json
import os

# Function to convert a single JSON file to COCO format
def convert_to_coco(json_data):
    # Initialize COCO data structure
    coco_data = {
        "info": {},
        "licenses": [],
        "images": [],
        "annotations": [],
        "categories": []
    }

    # Populate 'info' field (optional)
    coco_data["info"] = {
        "description": "Converted from custom format to COCO format",
        "url": "http://cocodataset.org",
        "version": "1.0",
        "year": 2024,
        "contributor": "Anonymous",
        "date_created": "2024-06-29"
    }

    # Populate 'images' field
    image_data = {
        "id": 1,
        "file_name": json_data["filename"],
        "width": int(json_data["size"]["width"]),
        "height": int(json_data["size"]["height"]),
        "license": None,
        "flickr_url": None,
        "coco_url": None,
        "date_captured": None
    }
    coco_data["images"].append(image_data)

    # Populate 'annotations' field
    annotation_id = 1
    for obj in json_data["object"]:
        annotation = {
            "id": annotation_id,
            "image_id": image_data["id"],
            "category_id": None,  # This needs to be mapped to COCO categories
            "segmentation": [],
            "area": 0,
            "bbox": [
                float(obj["bndbox"]["xmin"]),
                float(obj["bndbox"]["ymin"]),
                float(obj["bndbox"]["xmax"]) - float(obj["bndbox"]["xmin"]),
                float(obj["bndbox"]["ymax"]) - float(obj["bndbox"]["ymin"])
            ],
            "iscrowd": 0,
            "attributes": obj.get("attributes", [])
        }
        coco_data["annotations"].append(annotation)
        annotation_id += 1

    return coco_data

# Directory containing JSON files
json_dir = "/path/to/your/json/files/"

# Initialize a list to hold all COCO data dictionaries
all_coco_data = []

# Iterate over each JSON file in the directory
for filename in os.listdir(json_dir):
    if filename.endswith(".json"):
        json_path = os.path.join(json_dir, filename)
        with open(json_path, 'r') as f:
            json_data = json.load(f)
        
        # Convert JSON data to COCO format
        coco_data = convert_to_coco(json_data)
        
        # Append to the list of all COCO data
        all_coco_data.append(coco_data)

# Optionally, merge all COCO data into one dictionary
merged_coco_data = {
    "info": all_coco_data[0]["info"],  # Assuming info is the same across all files
    "licenses": all_coco_data[0]["licenses"],  # Assuming licenses are the same across all files
    "images": [],
    "annotations": [],
    "categories": all_coco_data[0]["categories"]  # Assuming categories are the same across all files
}

# Merge images and annotations
for coco_data in all_coco_data:
    merged_coco_data["images"].extend(coco_data["images"])
    merged_coco_data["annotations"].extend(coco_data["annotations"])

# Convert to JSON format
merged_json = json.dumps(merged_coco_data, indent=4)
print(merged_json)

# Optionally, save merged JSON to a file
output_path = os.path.join(json_dir, "merged_coco_data.json")
with open(output_path, 'w') as f:
    json.dump(merged_coco_data, f, indent=4)

print(f"Merged COCO data saved to {output_path}")