import os
import re
import tqdm
import collections
import argparse
import numpy as np

from suno_utils.audio import Audio
from suno_utils.utils.text import write_jsonl, read_jsonl, write_json, read_json, normalize_whitespace
from suno_utils.utils.s3 import read_from_s3, check_s3_file_exists, open_from_s3


# parse
parser = argparse.ArgumentParser()
parser.add_argument("--index", required=True, help="partial index")
parser.add_argument("--num_partial", required=True, help="number of partial data")
args = parser.parse_args()
index = int(args.index)
num_partial = int(args.num_partial)

# load data
metadata_path = "/app/suno/data/audio_mono_24khz/genius_hq/metadata/"

if not os.path.exists(os.path.join(metadata_path, "ids_63.npy")):
    print("process metadata...")
    print("loading metadata...")
    genius_metas = read_jsonl("/app/suno/data/chirp_v3/metadata/genius_hq.jsonl")
    original_metas = read_from_s3("s3://suno-data/datasets/harvest/genius_hq/pairs_metadata.jsonl", read_f=read_jsonl)
    original_ids_to_ix = {m["genius_meta"]["slug"]: i for i, m in enumerate(original_metas)}
    print('loaded!')
    
    for index in range(64):
        print("processing %d..." % index)
        # split data
        hop = len(genius_metas) // num_partial
        if index == num_partial:
            p_genius_metas = genius_metas[hop * index - 1:]
        else:
            p_genius_metas = genius_metas[hop * index : (index + 1) * hop]

        ids = [m["id"] for m in p_genius_metas]
        file_paths = [original_metas[original_ids_to_ix[m["original_id"]]]["audio_filepath"] for m in p_genius_metas]
        np.save(open(os.path.join(metadata_path, "ids_%d.npy" % index), "wb"), ids)
        np.save(open(os.path.join(metadata_path, "filepaths_%d.npy" % index), "wb"), file_paths)

print("load precomputed files...")
index = int(args.index)
ids = np.load(os.path.join(metadata_path, "ids_%d.npy" % index))
file_paths = np.load(os.path.join(metadata_path, "filepaths_%d.npy" % index))
print("loaded!")

# resample audio
print("start resampling...")
audio_path = "/app/suno/data/audio_mono_24khz/genius_hq/audio/"
for i in tqdm.tqdm(range(len(ids))):
    _id = ids[i]
    save_path = os.path.join(audio_path, _id + ".wav")

    if not os.path.exists(save_path):
        try:
            audio = Audio.from_s3(file_paths[i], sample_rate=24000, n_channels=1)
            if (audio.duration_s) > (30 and audio.duration_s < 500):
                audio.write_wav(save_path)
        except:
            continue