import os
import torch
from tqdm import tqdm
from suno_utils.utils.text import (
    write_jsonl,
    read_jsonl,
    write_json,
    read_json,
    normalize_whitespace,
)

bundle_jsonl_filepaths = [
    "s3://suno-data/datasets/bundles/v1/youtube_music/metas.jsonl",
    "s3://suno-data/datasets/bundles/v1/genius_hq/metas_plus.jsonl",
    "s3://suno-data/datasets/bundles/v1/jamendo/metas.jsonl",
    "s3://suno-data/datasets/bundles/v1/imslp/metas.jsonl",
    "s3://suno-data/datasets/bundles/v2/pond5_music/metas.jsonl",
    "s3://suno-data/datasets/bundles/v2/deezer/metas.jsonl",
    "s3://suno-data/datasets/bundles/v2/ytm_tagged/metas.jsonl",
    "s3://suno-data/datasets/bundles/v3/discogs/metas_plus.jsonl",
]

# where to search for s3 filepath
dataset_key_to_dataset_name = {
    "youtube_music_lyrics_foreign": "youtube_music",
    "youtube_music_lyrics": "youtube_music",
    "youtube_music": "youtube_music",
    "genius_hq_lyrics": "genius_hq",
    "genius_hq_lyrics_foreign": "genius_hq",
    "jamendo": "jamendo",
    "imslp": "imslp",
    "pond5_music": "pond5_music",
    "deezer_lyrics": "deezer",
    "deezer_lyrics_foreign": "deezer",
    "deezer": "deezer",
    "ytm_tagged": "ytm_tagged",
    "discogs": "discogs",
    "discogs_lyrics": "discogs",
    "discogs_lyrics_foreign": "discogs",
    "discogs_covers": "discogs",
}

if __name__ == "__main__":
    # configuration
    meta_dir = "metadata"  # local dir to store bundle metadata
    os.makedirs(meta_dir, exist_ok=True)
    use_val = True

    if use_val:
        metas_filepath = "/app/suno/data/chirp_v4/multi/metas_val.jsonl"
        out_metas_filepath = "/app/suno/data/chirp_v4_ft/multi/metas_val.jsonl"
    else:
        metas_filepath = "/app/suno/data/chirp_v4/multi/metas_tr.jsonl"
        out_metas_filepath = "/app/suno/data/chirp_v4_ft/multi/metas_tr.jsonl"

    # load the memmap aligned jsonl file
    print(f"Reading {metas_filepath} from disk...")
    main_metas = read_jsonl(metas_filepath)

    master_meta_map = {}
    for bundle_jsonl_filepath in bundle_jsonl_filepaths:
        dataset_name = os.path.basename(os.path.dirname(bundle_jsonl_filepath))
        output_filepath = os.path.join(meta_dir, f"{dataset_name}_metas.jsonl")
        if not os.path.isfile(output_filepath):
            os.system(f"aws s3 cp {bundle_jsonl_filepath} {output_filepath}")

        # open the metadata file
        print(f"Reading {output_filepath} from disk...")
        metas = read_jsonl(output_filepath)

        # create a map indexed by the id (might not be unique)
        metas_map = {}
        for meta in metas:
            metas_map[meta["id"]] = meta

        # add to master map
        master_meta_map[dataset_name] = metas_map

    new_metas = []
    count = 0
    total = len(main_metas)
    # iterate over metas
    print("Connecting existing metadata with s3_filepaths...")
    pbar = tqdm(main_metas)
    for idx, meta in enumerate(pbar):
        dataset = meta.get("dataset")
        s3_filepath = meta.get("s3_filepath")

        if dataset is None and s3_filepath is None:
            print("No 'dataset' or 's3_filepath' specified. Skipping lookup...")
            print(meta)
        else:
            if s3_filepath is not None:
                pass  # do nothing if we already have the s3_filepath
            else:  # try to attempt lookup
                # connect this example to the correct bundle
                dataset_name = dataset_key_to_dataset_name[dataset]
                id_key = meta["id"]

                # now search in the master map, this can fail
                full_meta = master_meta_map[dataset_name].get(id_key)

                if full_meta is not None:
                    # insert s3_filepath into meta
                    if "s3_filepath" in full_meta:
                        meta["s3_filepath"] = full_meta["s3_filepath"]
                        count += 1
                    elif "audio_filepath" in full_meta:
                        meta["s3_filepath"] = full_meta["audio_filepath"]
                        count += 1
                    else:
                        print("No s3_filepath info found...")
                else:
                    print(dataset_name, meta)

        pbar.set_description(
            desc=f"{count}/{total} {(count/total)*100:0.2f}% have s3_filepath"
        )

        # store meta with s3_filepath if available
        new_metas.append(meta)

    write_jsonl(new_metas, out_metas_filepath)