import random
import subprocess
from multiprocessing import Pool, cpu_count
from pathlib import Path

from tqdm import tqdm


def convert_to_flac(src_file_path):
    dst_file_path = src_file_path.with_suffix(".flac")
    dst_file_path.parent.mkdir(parents=True, exist_ok=True)

    try:
        subprocess.check_call(
            [
                "ffmpeg",
                "-y",
                "-i",
                str(src_file_path),
                "-acodec",
                "flac",
                "-threads",
                "0",
                str(dst_file_path),
            ],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )

        # remove the input file
        src_file_path.unlink()
        return True
    except subprocess.CalledProcessError:
        return False


if __name__ == "__main__":
    src_dir = Path("dataset/tts/WenetSpeech/cleaned")

    wav_files = list(src_dir.rglob("*.wav"))
    random.shuffle(wav_files)
    print(f"Found {len(wav_files)} wav files")

    success_counter = 0
    fail_counter = 0

    with Pool(processes=cpu_count(), maxtasksperchild=100) as pool:
        with tqdm(
            pool.imap_unordered(convert_to_flac, wav_files), total=len(wav_files)
        ) as pbar:
            for success in pbar:
                if success:
                    success_counter += 1
                else:
                    fail_counter += 1

            pbar.set_description(f"Success: {success_counter}, Fail: {fail_counter}")

    print(f"Successfully converted: {success_counter}")
    print(f"Failed conversions: {fail_counter}")
