import argparse
import os
import json
import labelme
import base64
import pandas as pd

def arg_directory(path):
    if os.path.isdir(path):
        return path
    else:
        raise argparse.ArgumentTypeError(f'`{path}` is not a valid path')

def process_json_files(json_dir, img_dir, output_dir):
    json_files = [f for f in os.listdir(json_dir) if f.endswith('.json')]
    
    for json_file in json_files:
        if json_file.startswith('.'):
            continue

        json_path = os.path.join(json_dir, json_file)
        img_path = os.path.join(img_dir, json_file.replace('.json', '.jpg'))

        if not os.path.exists(img_path):
            continue

        try:
            with open(img_path, 'rb') as img_file:
                image_data = base64.b64encode(img_file.read()).decode('utf-8')
        except FileNotFoundError:
            continue

        with open(json_path, 'r') as f:
            json_data = json.load(f)

        json_data['imageData'] = image_data
        json_data['imagePath'] = img_path

        output_path = os.path.join(output_dir, json_file)
        with open(output_path, 'w') as f:
            json.dump(json_data, f, indent=4)

def main():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter, description='Add image data info to .json file')
    parser.add_argument('path_to_json', type=arg_directory, help='Path to the folder containing the .json files', required=True)
    parser.add_argument('path_to_img', type=arg_directory, help='Path to the folder containing the .jpg images', required=True)
    parser.add_argument('-d', '--directory', type=arg_directory, help='Directory to which modified .json files will be stored'
                                                                'if no argument, directory will be same as path_to_json', required=False, 
                                                                default=None)
    args = parser.parse_args()

    if args.directory == None:
        process_json_files(args.path_to_json, args.path_to_img, args.path_to_json)
    else:    
        process_json_files(args.path_to_json, args.path_to_img, args.directory)

if __name__ == "__main__":
    main()