import argparse
import json
import os
import soundfile as sf


def preprocess_one_dir(in_dir, out_dir, out_filename):
    """Create .json file for one condition."""
    file_infos = []
    in_dir = os.path.abspath(in_dir)
    wav_list = os.listdir(in_dir)
    wav_list.sort()
    for wav_file in wav_list:
        if not wav_file.endswith(".wav"):
            continue
        wav_path = os.path.join(in_dir, wav_file)
        samples = sf.SoundFile(wav_path)
        file_infos.append((wav_path, len(samples)))
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    with open(os.path.join(out_dir, out_filename + ".json"), "w") as f:
        json.dump(file_infos, f, indent=4)


def preprocess(inp_args):
    """Create .json files for all conditions."""
    speaker_list = ["mix"] + [f"s{n+1}" for n in range(inp_args.n_src)]
    for data_type in ["tr", "cv", "tt"]:
        for spk in speaker_list:
            preprocess_one_dir(
                os.path.join(inp_args.in_dir, data_type, spk),
                os.path.join(inp_args.out_dir, data_type),
                spk,
            )


if __name__ == "__main__":
    parser = argparse.ArgumentParser("WSJ0-MIX data preprocessing")
    parser.add_argument(
        "--in_dir", type=str, default=None, help="Directory path of wham including tr, cv and tt"
    )
    parser.add_argument("--n_src", type=int, default=2, help="Number of sources in wsj0-mix")
    parser.add_argument(
        "--out_dir", type=str, default=None, help="Directory path to put output files"
    )
    args = parser.parse_args()
    print(args)
    preprocess(args)
