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



METAS_DIR = "/app/suno/data/audio_mono_24khz/pond5/"

# get track metadata
metas_fn = os.path.join(METAS_DIR, "pond5_metadata.jsonl")
track_paths = []
if os.path.exists(metas_fn):
    print("files exsist")
    metas = read_jsonl(metas_fn)
    track_paths = np.load(os.path.join(METAS_DIR, "file_paths.npy"))
    print("files loaded")
else:
    def _clean_tag(s):
        if s is None or not isinstance(s, str):
            return ""
        # remove symbold with special meaning
        s = re.sub(r"[\{\}\_\[\]]", " ", s)
        # squash whitespace
        return normalize_whitespace(s)

    pd5_base_metas = read_from_s3("s3://suno-data/datasets/bundles/v2/pond5_music/metas.jsonl", read_f=read_jsonl)
    pd5_extra_metas = read_from_s3("s3://suno-data/datasets/harvest/pond5_music/pond5_metas.jsonl", read_f=read_jsonl)

    extra_info = {}
    seen_ids = set()
    genre_counter = collections.Counter()
    for m in tqdm.tqdm(pd5_extra_metas):
        if "s3_filepath" not in m:
            continue
        if m["id"] in seen_ids:
            continue
        if m["duration"] < 30:
            continue
        if m["duration"] > 60 * 4:
            continue
        name = m.get("name", "")
        description = m.get("description", "")
        tags = m.get("tags", [])
        genre = m.get("genre", "")
        if not genre:
            continue
        genre_counter[genre.lower().strip()] += 1
        if not description or not tags or not name:
            continue
        if len(description.strip().split()) < 10:
            continue
        if len(description.strip().split()) > 100:
            continue
        if not isinstance(tags, list) or len(tags) < 5 or len(tags) > 50:
            continue
        # if there are weird tags
        if max(len(tag) for tag in tags) > 30:
            continue
        extra_info[m["id"]] = {
            "name": name if isinstance(name, str) else "",
            "description": (description if isinstance(description, str) else ""),
            "tags": ([genre] if genre else []) + tags if isinstance(tags, list) else "",
        }
        seen_ids.add(m["id"])

    metas = []
    total_duration = 0
    n_c =0
    for m in tqdm.tqdm(pd5_base_metas):
    #     if n_c > 5:
    #         break
        if m["id"] not in seen_ids:
            continue
        n_c += 1
        info = extra_info[m["id"]]
        new_m = {"id": m["id"]}
        # print(m, info)
        tags = [
            info["name"],
            info["description"],
        ]
        tags.extend(info["tags"])
        tags = [_clean_tag(t) for t in tags]
        tags = [t for t in tags if len(t) > 0]
        if len(tags) > 10:
            new_m["tags"] = list(set(tags))
            metas.append(new_m)
            total_duration += m["duration_s"]
            track_paths.append(m.get("audio_filepath", m.get("s3_filepath", m.get("filepath"))))
    print(f"{len(metas):,} entries, total duration {total_duration/60/60:,.0f} hr")

    write_jsonl(metas, os.path.join(METAS_DIR, "pond5_metadata.jsonl"))
    np.save(open(os.path.join(METAS_DIR, "file_paths.npy"), "wb"), track_paths)

# 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)

# split data
hop = len(metas) // num_partial
if index == num_partial:
    metas = metas[hop * index - 1:]
    track_paths = track_paths[hop * index - 1:]
else:
    metas = metas[hop * index : (index + 1) * hop]
    track_paths = track_paths[hop * index : (index + 1) * hop]

# resample audio
print("start resampling...")
audio_path = "/app/suno/data/audio_mono_24khz/pond5/audio/"
for i in tqdm.tqdm(range(len(metas))):
    _id = metas[i]["id"]
    save_path = os.path.join(audio_path, _id + ".wav")
    if not os.path.exists(save_path):
        audio = Audio.from_s3(track_paths[i], sample_rate=24000, n_channels=1)
        audio.write_wav(save_path)