import gradio as gr
from gtts import gTTS
from pydub import AudioSegment
import pandas as pd
import io

csv_files = {"A1":"OF3KA1.csv","A2":"OF3KA2.csv","B1":"OF3KB1.csv","B2":"OF3KB2.csv","C1":"OF3KC1.csv","5K":"OF5K.csv"}

# Function to load the DataFrame based on level selection
def load_data(level):
    # Use the csv_files dictionary to get the correct file name for the given level
    csv_file_path = f"./{csv_files[level]}"  # This correctly uses the dictionary
    data = pd.read_csv(csv_file_path)
    return data

def generate_speech(level, x, y):
    data = load_data(level)  # Load data for the selected level
    
    # Ensure x and y are integers and within range
    x, y = int(x), int(y)
    filtered_df = data[(data['SID'] >= x) & (data['SID'] <= y)]

    combined_audio = AudioSegment.silent(duration=1000)  # Start with silence for padding

    for _, row in filtered_df.iterrows():
        sentence = f"Number {row['SID']}. {row['WORD']}! {row['WORD']} is {row['POS']}."
        tts = gTTS(text=sentence, lang='en')
        mp3_fp = io.BytesIO()
        tts.write_to_fp(mp3_fp)
        mp3_fp.seek(0)
        sentence_audio = AudioSegment.from_file(mp3_fp, format="mp3")
        combined_audio += sentence_audio + AudioSegment.silent(duration=1500)

    mp3_io = io.BytesIO()
    combined_audio.export(mp3_io, format='mp3')
    mp3_io.seek(0)
    return mp3_io.getvalue()

# Interface with updated level selection including all specified options
iface = gr.Interface(
    fn=generate_speech,
    inputs=[
        gr.Dropdown(label="Select Level (3K: A1, A2, B1, B2, C1; 5K: additional B2 and C1)", choices=['A1', 'A2', 'B1', 'B2', 'C1', '5K']),  # Updated choices
        gr.Number(label="Start Number (x)"),
        gr.Number(label="End Number (y)")
    ],
    outputs=gr.Audio(label="Generated Speech"),
    title="Oxford Learner Vocabulary by CEFR levels: Learn with Sound",
    description="Choose a level and define the starting and ending numbers for that level. The system will create a single audio file formatted as 'Number 1. Word is a noun.' After submission, you have the option to download the audio file. Additionally, you can download the numbered lists for each level from the 'My Apps' section at https://mrkim21.github.io (web address)"
)

iface.launch(share=True, debug=True)