import os
import soundfile as sf
import csv
from datasets import load_dataset

# Load the Italian subset of the Multilingual LibriSpeech dataset
dataset = load_dataset("facebook/multilingual_librispeech", "italian")

# Define the output directory
output_dir = "multilingual_librispeech_italian"
os.makedirs(output_dir, exist_ok=True)

def save_split(split_name, dry_run=False):
    split = dataset[split_name]
    split_dir = os.path.join(output_dir, split_name)
    os.makedirs(split_dir, exist_ok=True)

    wavs_dir = os.path.join(split_dir, "wavs")
    os.makedirs(wavs_dir, exist_ok=True)

    COLUMNS_TO_KEEP = ["transcript", "audio", "sampling_rate"]
    all_columns = split.column_names

    if dry_run:
        print(split)
        columns_to_remove = set(all_columns) - set(COLUMNS_TO_KEEP)
        split = split.remove_columns(columns_to_remove)
        print(split[0])
        return

    columns_to_remove = set(all_columns) - set(COLUMNS_TO_KEEP)
    split = split.remove_columns(columns_to_remove)

    metadata_path = os.path.join(split_dir, "metadata.csv")

    with open(metadata_path, mode='w', newline='', encoding='utf-8') as file:
        writer = csv.writer(file, delimiter='|')

        for i, example in enumerate(split):
            # Extract audio data and sampling rate
            audio = example["audio"]
            audio_array = audio["array"]
            sampling_rate = audio["sampling_rate"]

            # Define file paths
            audio_path = os.path.join(wavs_dir, f"{i}.wav")

            # Save audio file in WAV format
            sf.write(audio_path, audio_array, sampling_rate)

            # Save transcription
            # transcription_path = os.path.join(split_dir, f"{i}.txt")
            # with open(transcription_path, "w", encoding="utf-8") as f:
            #     f.write(example["transcript"])

            # Save metadata
            writer.writerow([audio_path, example["transcript"]])

# save_split("1_hours", dry_run=True)
save_split("9_hours")