app.py

import gradio as gr
import openai, os, io, tempfile
from dotenv import load_dotenv

# ============== 환경 변수 & OpenAI 초기화 ===========================
load_dotenv()
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
    raise RuntimeError("OPENAI_API_KEY를 .env 파일에 설정하세요!")
client = openai.OpenAI(api_key=api_key)

# ============== 언어 설정 ===========================================
LANGUAGES = [
    "Korean", "English", "Japanese", "Chinese",
    "Thai", "Russian", "Vietnamese",
    "Spanish", "French"
]
LANG_CODE = {
    "Korean": "ko", "English": "en", "Japanese": "ja", "Chinese": "zh",
    "Thai": "th", "Russian": "ru", "Vietnamese": "vi",
    "Spanish": "es", "French": "fr"
}
VOICE = {
    lang: ("nova" if lang in ["Korean", "Japanese", "Chinese"] else "alloy")
    for lang in LANGUAGES
}

# ============== 공통 함수 ===========================================
def _gpt_translate(text: str, src: str, tgt: str) -> str:
    """GPT-3.5 번역"""
    rsp = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": (
                    f"You are a professional translator. Translate the following {src} text to {tgt}. "
                    f"Only provide the translation without additional commentary."
                )
            },
            {"role": "user", "content": text}
        ],
        temperature=0.3,
        max_tokens=2048
    )
    return rsp.choices[0].message.content.strip()

def _tts(text: str, lang: str) -> str:
    """TTS-1 음성 합성 → 임시 mp3 경로 반환"""
    out = client.audio.speech.create(
        model="tts-1",
        voice=VOICE.get(lang, "alloy"),
        input=text[:4096]
    )
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
    tmp.write(out.content)
    tmp.close()
    return tmp.name

def translate_audio(audio_path, src, tgt):
    """단일 음성 파일 번역(TTS 포함)"""
    with open(audio_path, "rb") as f:
        stt = client.audio.transcriptions.create(
            model="whisper-1",
            file=f,
            language=LANG_CODE.get(src)
        )
    original = stt.text.strip()
    if not original:
        return "⚠️ 음성 인식 실패", "", None

    translated = _gpt_translate(original, src, tgt)
    tts_path = _tts(translated, tgt)
    return original, translated, tts_path

# ============== 실시간 스트리밍 전사/번역 ============================
STREAM_CHUNK_SEC = 4  # Whisper 호출 주기(초)

def stream_generator(mic_stream, src, tgt):
    """마이크 스트림 -> 주기적 청크 번역(누적 출력)"""
    buffer = io.BytesIO()
    wav_header = None
    original_acc, translated_acc = "", ""

    while True:
        chunk = mic_stream.recv()
        if chunk is None:                 # 스트림 종료
            break

        if not wav_header:
            wav_header = chunk[:44]       # WAV 헤더(16kHz 16-bit mono)
        buffer.write(chunk)

        # 지정 시간만큼 쌓이면 Whisper 호출
        if buffer.getbuffer().nbytes > 16000 * 2 * STREAM_CHUNK_SEC:
            wav_bytes = wav_header + buffer.getvalue()
            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
                tmp.write(wav_bytes)
                tmp.close()
                o, t, _ = translate_audio(tmp.name, src, tgt)

            original_acc += " " + o
            translated_acc += " " + t
            yield original_acc.strip(), translated_acc.strip()
            buffer = io.BytesIO()         # 버퍼 초기화

    # 마지막 남은 버퍼 처리
    if buffer.getbuffer().nbytes:
        wav_bytes = wav_header + buffer.getvalue()
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
            tmp.write(wav_bytes)
            tmp.close()
            o, t, _ = translate_audio(tmp.name, src, tgt)

        yield (original_acc + " " + o).strip(), (translated_acc + " " + t).strip()

# ============== 4개국 동시 번역 =====================================
FOUR_LANGS = ["English", "Chinese", "Thai", "Russian"]

def translate_audio_four(audio_path, src):
    """음성 파일 → 원문 + 4개 언어 동시 번역"""
    with open(audio_path, "rb") as f:
        stt = client.audio.transcriptions.create(
            model="whisper-1",
            file=f,
            language=LANG_CODE.get(src)
        )
    original = stt.text.strip()
    if not original:
        return ["⚠️ 음성 인식 실패"] + [""] * 4

    outs = [original]
    for lang in FOUR_LANGS:
        outs.append(_gpt_translate(original, src, lang))
    return outs  # 총 5개(원문+4언어)

# ============== Gradio UI ===========================================
with gr.Blocks(title="SMARTok Demo", theme=gr.themes.Soft()) as app:
    with gr.Tabs():
        # ① 마이크/파일 번역
        with gr.TabItem("🎙️ 마이크/파일 번역"):
            src1 = gr.Dropdown(LANGUAGES, value="Korean", label="입력")
            tgt1 = gr.Dropdown(LANGUAGES, value="English", label="출력")
            mic1 = gr.Audio(
                sources=["microphone", "upload"],
                type="filepath",
                label="🎤 녹음 또는 오디오 파일 업로드"
            )
            btn1 = gr.Button("번역")
            stt1 = gr.Textbox(label="원문", lines=5)
            tlt1 = gr.Textbox(label="번역", lines=5)
            out1 = gr.Audio(label="TTS", type="filepath", autoplay=True)

            btn1.click(
                translate_audio,
                inputs=[mic1, src1, tgt1],
                outputs=[stt1, tlt1, out1]
            )

        # ② 파일 전용 번역
        with gr.TabItem("🎧 파일 전용 번역"):
            src2 = gr.Dropdown(LANGUAGES, value="Korean", label="입력")
            tgt2 = gr.Dropdown(LANGUAGES, value="English", label="출력")
            file2 = gr.Audio(
                sources=["upload"],
                type="filepath",
                label="오디오 파일 업로드"
            )
            btn2 = gr.Button("번역")
            stt2 = gr.Textbox(label="원문", lines=5)
            tlt2 = gr.Textbox(label="번역", lines=5)
            out2 = gr.Audio(label="TTS", type="filepath", autoplay=True)

            btn2.click(
                translate_audio,
                inputs=[file2, src2, tgt2],
                outputs=[stt2, tlt2, out2]
            )

        # ③ 실시간 스트리밍 번역(Beta)
        with gr.TabItem("⏱️ 실시간 번역 (Beta)"):
            gr.Markdown("마이크를 켜면 3~4초 간격으로 자막이 갱신됩니다.")
            src3 = gr.Dropdown(LANGUAGES, value="Korean", label="입력")
            tgt3 = gr.Dropdown(LANGUAGES, value="English", label="출력")
            mic3 = gr.Audio(
                sources=["microphone"],
                streaming=True,
                label="🎤 실시간 마이크 입력"
            )
            stt3 = gr.Textbox(label="원문(실시간)", lines=8)
            tlt3 = gr.Textbox(label="번역(실시간)", lines=8)

            def gen(audio, src_lang, tgt_lang):
                yield from stream_generator(audio, src_lang, tgt_lang)

            mic3.stream(gen, inputs=[src3, tgt3], outputs=[stt3, tlt3])

        # ④ 4개 언어 동시 번역
        with gr.TabItem("🌏 4개 언어 동시"):
            gr.Markdown("입력 음성을 **English / Chinese(简体) / Thai / Russian** 4개 언어로 동시에 번역합니다.")
            src4 = gr.Dropdown(LANGUAGES, value="Korean", label="입력 언어")
            aud4 = gr.Audio(
                sources=["microphone", "upload"],
                type="filepath",
                label="🎤 녹음 또는 오디오 파일 업로드"
            )
            btn4 = gr.Button("번역")

            with gr.Row():
                org4 = gr.Textbox(label="원문", lines=4)
                en4  = gr.Textbox(label="English", lines=4)
                zh4  = gr.Textbox(label="Chinese (简体)", lines=4)
                th4  = gr.Textbox(label="Thai", lines=4)
                ru4  = gr.Textbox(label="Russian", lines=4)

            btn4.click(
                translate_audio_four,
                inputs=[aud4, src4],
                outputs=[org4, en4, zh4, th4, ru4]
            )

# ============== 앱 실행 =============================================
if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860, share=False, debug=True)