import gradio as gr
from gtts import gTTS
from pydub import AudioSegment
import numpy as np
import os

def text_to_speech(prompt):
    try:
        # 1) gTTS로 mp3 생성
        tts = gTTS(text=prompt, lang="bg")
        audio_file = "output.mp3"
        tts.save(audio_file)
    except Exception as e:
        print("gTTS 생성 오류:", e)
        raise e  # 에러를 다시 발생시켜 Gradio에서 감지하도록 함
    
    try:
        # 2) mp3 -> numpy 변환 (pydub 사용)
        sound = AudioSegment.from_mp3(audio_file)
        samples = np.array(sound.get_array_of_samples())
        
        # 스테레오면 모노로 변환
        if sound.channels > 1:
            samples = samples.reshape((-1, sound.channels))
            samples = samples.mean(axis=1)
        
        # int16 -> float32 정규화
        samples = samples.astype(np.float32) / 32768.0
        sample_rate = sound.frame_rate
    except Exception as e:
        print("pydub 로딩/처리 오류:", e)
        raise e
    finally:
        if os.path.exists(audio_file):
            os.remove(audio_file)
    
    return samples, sample_rate

with gr.Blocks() as demo:
    gr.Markdown("## Bulgarian Text-to-Speech (TTS)")
    with gr.Row():
        input_prompt = gr.Textbox(label="Enter a prompt in Bulgarian:")
        output_audio = gr.Audio(label="Generated Speech", type="numpy")
    generate_button = gr.Button("Generate Speech")
    
    generate_button.click(text_to_speech, inputs=input_prompt, outputs=output_audio)

if __name__ == "__main__":
    demo.launch()