from argparse import ArgumentParser
import json

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "--record_json",
        type=str,
    )
    parser.add_argument(
        "--ori_json",
        type=str,
    )
    parser.add_argument(
        "--target_json",
        type=str,
    )
    args = parser.parse_args()

    with open(args.record_json, "r") as f:
        records = json.load(f)

    with open(args.ori_json, "r") as f:
        ori_data = json.load(f)

    for item in records:
        item_id = item["id"]
        ori_item = ori_data[item_id]

        if "image_url" in ori_item:
            caption_keys = [
                "sharegpt4v_long_cap",
                "cogvlm_long",
                "blip2_short_cap",
                "llava13b_long_cap",
                "spatial_caption",
                "coca_caption",
                "user_prompt",
                "tags_prompt",
                "gpt4v_concise_elements",
                "gpt4v_regions_detailed_description",
                "gpt4v_detailed_description",
                "gpt4v_concise_description",
            ]
            all_caps = {key: ori_item[key] for key in caption_keys if key in ori_item and ori_item[key]}
            item["caps"] = all_caps
            item["image"] = item["image_url"]
        elif "path" in item:
            item["image"] = item["path"]
            item["caps"] = {"prompt": ori_item["prompt"]}
        else:
            raise NotImplementedError

    with open(args.target_json, "w") as f:
        json.dump(records, f, indent=1)
