import gradio as gr
from datasets import load_dataset

# --- Configuration ---
DATASET_NAME = "Cnam-LMSSC/vibravox-test"
SUBSETS = ["speech_clean", "speech_noisy", "speechless_clean", "speechless_noisy"]
SPLITS = ["train", "validation", "test"]
TEXT_COLUMN = "raw_text"
PHONEMIZED_TEXT_COLUMN = "phonemized_text"
GENDER_COLUMN = "gender"
AUDIO_COLUMNS = [
    "audio.headset_microphone",
    "audio.throat_microphone",
    "audio.soft_in_ear_microphone",
    "audio.rigid_in_ear_microphone",
    "audio.forehead_accelerometer",
    "audio.temple_vibration_pickup"
]

# --- Main Application Logic ---

def load_and_update_all(subset, split):
    """
    Loads a new dataset and returns updates for the entire UI.
    """
    try:
        dataset = load_dataset(DATASET_NAME, name=subset, split=split)
        has_text_fields = TEXT_COLUMN in dataset.features

        sample = dataset[0]
        sentence = sample.get(TEXT_COLUMN)
        phonemized_text = sample.get(PHONEMIZED_TEXT_COLUMN)
        gender = sample.get(GENDER_COLUMN)
        
        raw_audio_data = [
            (sample[col]['sampling_rate'], sample[col]['array']) for col in AUDIO_COLUMNS
        ]
        
        # --- THE FIX IS HERE ---
        # We add a condition to handle datasets with only one row.
        dataset_len = len(dataset)
        if dataset_len <= 1:
            # If there's only one item, hide the slider as it's not needed.
            slider_update = gr.update(visible=False)
        else:
            # Otherwise, show and configure the slider as normal.
            slider_update = gr.update(maximum=dataset_len - 1, value=0, visible=True, interactive=True)
        # --------------------

        return (
            dataset,
            slider_update, # Use the new slider_update variable here
            gr.update(value=sentence, visible=has_text_fields),
            gr.update(value=phonemized_text, visible=has_text_fields),
            gr.update(value=gender, visible=has_text_fields),
            *raw_audio_data,
            gr.update(value="", visible=False)
        )
    except Exception as e:
        error_message = f"Failed to load {subset}/{split}. Error: {e}"
        empty_audio = (None, None)
        return (
            None,
            gr.update(visible=False),
            gr.update(visible=False), gr.update(visible=False), gr.update(visible=False),
            *[empty_audio] * len(AUDIO_COLUMNS),
            gr.update(value=error_message, visible=True)
        )

def get_audio_row(dataset, index):
    """
    Fetches a new row from the currently loaded dataset when the slider moves.
    """
    if dataset is None:
        return [None] * (3 + len(AUDIO_COLUMNS))

    index = int(index)
    sample = dataset[index]

    has_text_fields = TEXT_COLUMN in dataset.features
    sentence = sample.get(TEXT_COLUMN)
    phonemized_text = sample.get(PHONEMIZED_TEXT_COLUMN)
    gender = sample.get(GENDER_COLUMN)

    raw_audio_data = [
        (sample[col]['sampling_rate'], sample[col]['array']) for col in AUDIO_COLUMNS
    ]

    return [sentence, phonemized_text, gender] + raw_audio_data

# --- Build the Gradio Interface (No changes needed here) ---
with gr.Blocks(css="footer {display: none !important}") as demo:
    gr.Markdown("# Vibravox Viewer")
    loaded_dataset_state = gr.State(None)
    with gr.Row():
        subset_dropdown = gr.Dropdown(SUBSETS, value="speech_clean", label="Select Subset")
        split_dropdown = gr.Dropdown(SPLITS, value="train", label="Select Split")
    error_box = gr.Textbox(visible=False, interactive=False, container=False)
    with gr.Row():
        sentence_output = gr.Textbox(label="Raw Text", interactive=False)
        phonemized_output = gr.Textbox(label="Phonemized Text", interactive=False)
        gender_output = gr.Textbox(label="Gender", interactive=False)
    slider = gr.Slider(label="Select Data Row")
    with gr.Row():
        audio1 = gr.Audio(label="Headset Microphone")
        audio2 = gr.Audio(label="Laryngophone (Throat Mic)")
        audio3 = gr.Audio(label="Soft In-Ear Microphone")
    with gr.Row():
        audio4 = gr.Audio(label="Rigid In-Ear Microphone")
        audio5 = gr.Audio(label="Forehead Accelerometer")
        audio6 = gr.Audio(label="Temple Vibration Pickup")

    all_outputs = [loaded_dataset_state, slider, sentence_output, phonemized_output, gender_output, audio1, audio2, audio3, audio4, audio5, audio6, error_box]
    data_outputs = [sentence_output, phonemized_output, gender_output, audio1, audio2, audio3, audio4, audio5, audio6]
    
    demo.load(fn=load_and_update_all, inputs=[subset_dropdown, split_dropdown], outputs=all_outputs)
    subset_dropdown.change(fn=load_and_update_all, inputs=[subset_dropdown, split_dropdown], outputs=all_outputs)
    split_dropdown.change(fn=load_and_update_all, inputs=[subset_dropdown, split_dropdown], outputs=all_outputs)
    slider.change(fn=get_audio_row, inputs=[loaded_dataset_state, slider], outputs=data_outputs)

demo.launch()