app.py

import os, asyncio, json, tempfile, websockets, pdfplumber
import gradio as gr
import openai
from dotenv import load_dotenv

# ─── 0. 초기화 ───────────────────────────────────────────────
load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")
if not openai.api_key:
    raise RuntimeError("OPENAI_API_KEY 가 .env 에 없습니다!")

LANG = ["Korean","English","Japanese","Chinese",
        "Thai","Russian","Vietnamese","Spanish","French"]
VOICE = {l: ("nova" if l in ["Korean","Japanese","Chinese"] else "alloy")
         for l in LANG}
FOUR = ["English","Chinese","Thai","Russian"]
WS_URL = "wss://api.openai.com/v1/audio/transcriptions/stream"

# ─── 1. 공통 GPT 번역 / TTS ─────────────────────────────────
async def gpt_translate(text, src, tgt):
    rsp = await openai.AsyncClient().chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role":"system",
                   "content":f"Translate {src} → {tgt}. Return only the text."},
                  {"role":"user","content":text}],
        temperature=0.3,max_tokens=2048)
    return rsp.choices[0].message.content.strip()

async def gpt_tts(text, lang):
    rsp = await openai.AsyncClient().audio.speech.create(
        model="tts-1", voice=VOICE[lang], input=text[:4096])
    tmp = tempfile.NamedTemporaryFile(delete=False,suffix=".mp3")
    tmp.write(rsp.content); tmp.close(); return tmp.name

# ─── 2. PDF 번역 ────────────────────────────────────────────
def translate_pdf(file, src, tgt):
    if not file: return "⚠️ PDF 업로드 필요", ""
    with pdfplumber.open(file.name) as pdf:
        text = "\n".join(p.extract_text() or "" for p in pdf.pages[:5]).strip()
    if not text:
        return "⚠️ 텍스트 추출 실패", ""
    return text, asyncio.run(gpt_translate(text, src, tgt))

# ─── 3. WebSocket STT 헬퍼 ──────────────────────────────────
async def ws_stt_generator(audio_queue: asyncio.Queue):
    """
    백그라운드 태스크:
      - audio_queue 로부터 chunk(bytes) 수신
      - WS 로 전송, 서버 event 수신 → yield (partial text, final?)
    """
    async with websockets.connect(
        WS_URL,
        extra_headers={"Authorization": f"Bearer {openai.api_key}"},
        max_size=None
    ) as ws:
        async def sender():
            while True:
                chunk = await audio_queue.get()
                if chunk is None:          # 종료 플래그
                    await ws.send(json.dumps({"terminate": True}))
                    break
                await ws.send(chunk)
        asyncio.create_task(sender())
        async for msg in ws:
            data = json.loads(msg)
            yield data["text"], data.get("final", False)

# ─── 4. Gradio 스트림 핸들러 ─────────────────────────────────
async def realtime_single(mic, src, tgt, state):
    """
    mic: bytes chunk (Gradio 자동)
    state: {"queue": Queue, "task": Task, "orig": str, "trans": str}
    """
    if state is None:
        state = {"queue": asyncio.Queue(), "task": None, "orig":"", "trans":""}

    if mic is None:   # 스트림 종료
        await state["queue"].put(None)
        return state["orig"], state["trans"], state

    # 첫 호출이면 WS 태스크 시작
    if state["task"] is None:
        async def run_ws():
            async for text, final in ws_stt_generator(state["queue"]):
                state["orig"] += (" " if state["orig"] else "") + text
                add = await gpt_translate(text, src, tgt)
                state["trans"] += (" " if state["trans"] else "") + add
        state["task"] = asyncio.create_task(run_ws())

    # 마이크 chunk enqueue
    await state["queue"].put(mic)
    return state["orig"], state["trans"], state

async def realtime_four(mic, src, state):
    if state is None:
        state = {"queue": asyncio.Queue(), "task": None,
                 "orig":"", "English":"", "Chinese":"", "Thai":"", "Russian":""}

    if mic is None:
        await state["queue"].put(None)
        return tuple(state[k] for k in
                     ["orig","English","Chinese","Thai","Russian"]) + (state,)

    if state["task"] is None:
        async def run_ws():
            async for text, _ in ws_stt_generator(state["queue"]):
                state["orig"] += (" "+text)
                for lang in FOUR:
                    state[lang] += (" "+ await gpt_translate(text, src, lang))
        state["task"] = asyncio.create_task(run_ws())

    await state["queue"].put(mic)
    return tuple(state[k] for k in
                 ["orig","English","Chinese","Thai","Russian"]) + (state,)

# ─── 5. UI ──────────────────────────────────────────────────
with gr.Blocks(title="SMARTok Demo") as demo:
    with gr.Tabs():
        # 탭 1 – 오디오 번역
        with gr.TabItem("🎙️ 오디오"):
            src1=gr.Dropdown(LANG,value="Korean",label="입력")
            tgt1=gr.Dropdown(LANG,value="English",label="출력")
            aud1=gr.Audio(sources=["microphone","upload"],type="filepath")
            btn1=gr.Button("번역")
            o1=gr.Textbox(label="원문"); t1=gr.Textbox(label="번역")
            a1=gr.Audio(label="TTS",type="filepath",autoplay=True)
            btn1.click(lambda a,s,t: translate_pdf.__wrapped__ if False else translate_pdf,
                       [aud1,src1,tgt1],[o1,t1,a1])  # dummy, 유지용

        # 탭 2 – PDF 번역
        with gr.TabItem("📄 PDF"):
            src2=gr.Dropdown(LANG,value="Korean",label="입력")
            tgt2=gr.Dropdown(LANG,value="English",label="출력")
            pdf=gr.File(file_types=[".pdf"])
            btn2=gr.Button("번역")
            o2=gr.Textbox(label="추출 원문",lines=15)
            t2=gr.Textbox(label="번역 결과",lines=15)
            btn2.click(translate_pdf:=translate_pdf,[pdf,src2,tgt2],[o2,t2])

        # 탭 3 – 실시간 1언어
        with gr.TabItem("⏱️ 실시간 1"):
            src3=gr.Dropdown(LANG,value="Korean",label="입력")
            tgt3=gr.Dropdown(LANG,value="English",label="출력")
            mic3=gr.Audio(sources=["microphone"],streaming=True)
            o3=gr.Textbox(label="원문(실시간)",lines=8)
            t3=gr.Textbox(label="번역(실시간)",lines=8)
            st3=gr.State()
            mic3.stream(realtime_single,inputs=[src3,tgt3,st3],
                        outputs=[o3,t3,st3])

        # 탭 4 – 실시간 4언어
        with gr.TabItem("🌏 실시간 4"):
            src4=gr.Dropdown(LANG,value="Korean",label="입력")
            mic4=gr.Audio(sources=["microphone"],streaming=True)
            o4=gr.Textbox(label="원문",lines=8)
            e4=gr.Textbox(label="English",lines=8)
            c4=gr.Textbox(label="Chinese(简体)",lines=8)
            th4=gr.Textbox(label="Thai",lines=8)
            r4=gr.Textbox(label="Russian",lines=8)
            st4=gr.State()
            mic4.stream(realtime_four,inputs=[src4,st4],
                        outputs=[o4,e4,c4,th4,r4,st4])

demo.launch(server_name="0.0.0.0",server_port=7860,debug=True)