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/tency/metadata/sub_paths"

if not os.path.exists(os.path.join(metadata_path, "filepaths_63.npy")):
    print("process metadata...")
    print("loading metadata...")
    with open(os.path.join(metadata_path, "filelist.txt"), "r") as f:
        lines = f.readlines()
    filelist = [f"s3://suno-data/{line.strip()}" for line in lines]

    # 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(filelist) // num_partial
        if index == num_partial:
            p_filelist = filelist[hop * index - 1 :]
        else:
            p_filelist = filelist[hop * index : (index + 1) * hop]

        np.save(
            open(os.path.join(metadata_path, f"filepaths_{index}.npy"), "wb"),
            p_filelist,
        )

print("load precomputed files...")
index = int(args.index)
file_paths = np.load(os.path.join(metadata_path, f"filepaths_{index}.npy"))
print("loaded!")

# resample audio
print("start resampling...")
audio_path = "/app/suno/data/audio_mono_24khz/tency/audio/"
for i in tqdm.tqdm(range(len(file_paths))):
    _id = file_paths[i].split("/")[-2]
    _fn = os.path.basename(file_paths[i])
    save_path = os.path.join(audio_path, _id + "_" + _fn[:-3] + "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