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/ytm_tagged/metadata/"

if not os.path.exists(os.path.join(metadata_path, "ids_63.npy")):
    print("process metadata...")
    print("loading metadata...")
    ytm_metas = read_jsonl("/app/suno/data/chirp_v3/metadata/ytm_tagged.jsonl")
    original_metas = read_from_s3(
        "s3://suno-data/datasets/bundles/v2/ytm_tagged/metas.jsonl",
        read_f=read_jsonl,
    )
    original_ids_to_ix = {m["id"]: i for i, m in enumerate(original_metas)}
    print("loaded!")

    for index in range(64):
        print("processing %d..." % index)
        # split data
        hop = len(ytm_metas) // num_partial
        if index == num_partial:
            p_ytm_metas = ytm_metas[hop * index - 1 :]
        else:
            p_ytm_metas = ytm_metas[hop * index : (index + 1) * hop]

        ids = [m["id"] for m in p_ytm_metas]
        file_paths = [
            original_metas[original_ids_to_ix[m["id"]]]["s3_filepath"]
            for m in p_ytm_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/ytm_tagged/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