import os
import json
from tqdm import tqdm
from suno_utils.utils.text import read_jsonl, write_jsonl

if __name__ == "__main__":
    # read ids from the file
    with open(
        "/home/christian/code/christian/metadata/auk_id_to_artist_ids.json", "r"
    ) as f:
        auk_id_to_artist_ids = json.load(f)

    print(len(auk_id_to_artist_ids))

    base_dir = "/app2/suno/data/diffusion/v0"

    for subset in ["val", "tr"]:
        metas_filename = f"metas_{subset}_aligned_v3.jsonl"
        print(os.path.join(base_dir, metas_filename))
        metas = read_jsonl(os.path.join(base_dir, metas_filename))
        print(len(metas))

        new_metas = []
        artist_count = 0
        for meta in tqdm(metas):
            # check if this meta id is in the auk_id_to_artist_ids map
            if meta["id"] in auk_id_to_artist_ids:
                meta["artist_ids"] = auk_id_to_artist_ids[meta["id"]]
                artist_count += 1
            new_metas.append(meta)

        print(len(new_metas), artist_count)

        write_jsonl(
            new_metas,
            os.path.join(base_dir, f"metas_{subset}_aligned_v3_artist_ids.jsonl"),
        )
        print(
            f"Wrote {len(new_metas)} metas to {os.path.join(base_dir, f'metas_{subset}_aligned_v3_artist_ids.jsonl')}"
        )