import os
import glob
import json

from typing import List

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


def find_all_audio_files(audio_dir: str, exts: List[str] = ["wav"]):
    # find all audio files
    filepaths = []
    for ext in exts:
        filepaths += glob.glob(
            os.path.join(audio_dir, "**", f"*.{ext}"), recursive=True
        )
    return filepaths


if __name__ == "__main__":

    SCORE_THRESHOLD = -0.5
    METAS_DIR = (
        "/home/christian/code/christian/outputs/quality_metas/audio_2ch_48khz_lg/train"
    )

    metas_filepaths = glob.glob(os.path.join(METAS_DIR, "*.jsonl"))
    audio_filepaths = []

    total_num_metas = 0
    for metas_filepath in metas_filepaths:
        # load info from metas jsonl
        metas = read_jsonl(metas_filepath)
        num_metas = len(metas)
        subset_audio_filepaths = []

        for meta in metas:
            # check the audio quality
            score = float(meta["audio_quality"]["score"])
            if score > SCORE_THRESHOLD:
                subset_audio_filepaths.append(meta["audio_file"])

        audio_filepaths.extend(subset_audio_filepaths)
        total_num_metas += num_metas

        keep_percent = (len(subset_audio_filepaths) / num_metas) * 100
        print(f"{os.path.basename(metas_filepath)}: {keep_percent:.2f}%")

    total_keep_percent = (len(audio_filepaths) / total_num_metas) * 100
    print(
        f"Total: {len(audio_filepaths)}/{total_num_metas} ({total_keep_percent:.2f}%)"
    )