import json
import argparse
import os


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--result_fold', type=str, required=True)
    parser.add_argument('--harmful_save_file', type=str, required=True)
    parser.add_argument('--safe_save_file', type=str, required=True)
    args = parser.parse_args()

    unsafe_count, total = 0, 0
    result_files = os.listdir(args.result_fold)
    harmful_out = open(args.harmful_save_file, 'w')
    safe_out = open(args.safe_save_file, 'w')
    for result_file in result_files:
        img_path = os.path.join("/data//mm-safety/val2017", result_file.split('_')[0]+".jpg")
        with open(os.path.join(args.result_fold, result_file), 'r', encoding='utf-8') as f:
            for line in f:
                item = json.loads(line)
                harmful_type = item['harmful_type']
                if harmful_type != '':
                    harmful_out.write(json.dumps({
                        'img_path': img_path, 'prompt': item['prompt'], 'response': item['response']
                    }) + '\n')
                    print(item['response'], "|", harmful_type)
                    unsafe_count += 1
                else:
                    safe_out.write(json.dumps({
                        'img_path': img_path, 'prompt': item['prompt']
                    }) + '\n')
                total += 1
    print(f"Unsafe responses: {unsafe_count} / {total} = {unsafe_count / total * 100:.2f}%")


if __name__ == '__main__':
    main()
