import gradio as gr
import tempfile
from transformers import pipeline
from gtts import gTTS

# โหลดโมเดล ASR สำหรับฟังเสียงภาษาไทยให้แม่นยำยิ่งขึ้น
asr_pipeline = pipeline(
    "automatic-speech-recognition",
    model="biodatlab/whisper-th-medium-combined",  # หรือ distill-whisper-th-small
    chunk_length_s=30,
    device=0  # ใช้ GPU ถ้ามี
)

# โหลดโมเดลแปลภาษา (Multilingual)
translator_pipeline = pipeline(
    "translation",
    model="facebook/nllb-200-distilled-600M"
)

# ฟังก์ชันหลัก: รับเสียง + ทิศทางการแปล
def translate_voice(audio_input, direction):
    if audio_input is None:
        return "กรุณาอัปโหลดหรืออัดเสียงก่อน", None

    # จัดการไฟล์เสียงจาก Gradio
    if isinstance(audio_input, tuple):
        audio_path = audio_input[0]
    else:
        audio_path = audio_input

    # แปลงเสียงเป็นข้อความ
    asr_result = asr_pipeline(audio_path)
    original_text = asr_result["text"]

    # กำหนดภาษาต้นทาง/ปลายทาง
    if direction == "มลายู → ไทย":
        src_lang = "zsm_Latn"   # มลายูมาตรฐาน
        tgt_lang = "tha_Thai"
        tts_lang = "th"
    else:
        src_lang = "tha_Thai"
        tgt_lang = "zsm_Latn"
        tts_lang = "ms"

    # แปลข้อความ
    translation = translator_pipeline(
        original_text,
        src_lang=src_lang,
        tgt_lang=tgt_lang
    )[0]["translation_text"]

    # สร้างเสียงจากข้อความแปลด้วย gTTS
    tts = gTTS(text=translation, lang=tts_lang)
    temp_audio_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3").name
    tts.save(temp_audio_path)

    return translation, temp_audio_path

# UI ด้วย Gradio
demo = gr.Interface(
    fn=translate_voice,
    inputs=[
        gr.Audio(sources=["microphone", "upload"], type="filepath", label="พูดหรืออัปโหลดเสียง"),
        gr.Radio(["มลายู → ไทย", "ไทย → มลายู"], value="มลายู → ไทย", label="เลือกทิศทางแปล")
    ],
    outputs=[
        gr.Textbox(label="ข้อความที่แปลได้"),
        gr.Audio(label="เสียงที่แปลกลับ")
    ],
    title="🩺 ระบบแปลเสียง มลายู ↔ ไทย สำหรับการแพทย์",
    theme="default"
)

demo.launch()