app.py

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

# ===== 공통 초기화 =========================================
load_dotenv()
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
    print("⚠️  OPENAI_API_KEY를 .env 파일에 설정하세요!")
else:
    print(f"✅ API Key 로드됨: {api_key[:10]}...")

try:
    client = openai.OpenAI(api_key=api_key)
except Exception as e:
    print(f"❌ OpenAI 클라이언트 초기화 실패: {e}")
    client = None

# ===== 언어 설정 ===========================================
LANGUAGES = [
    "Korean", "English", "Japanese", "Chinese",  # 기존
    "Thai", "Russian", "Vietnamese",             # 추가
    "Spanish", "French"                          # 선택
]

# Whisper용 ISO-639 코드 매핑
LANG_CODE_MAP = {
    "Korean": "ko", "English": "en", "Japanese": "ja", "Chinese": "zh",
    "Thai": "th", "Russian": "ru", "Vietnamese": "vi",
    "Spanish": "es", "French": "fr"
}

# TTS 음성 매핑(OpenAI tts-1: alloy, nova 두 가지)
VOICE_MAP = {
    "Korean": "nova",
    "English": "alloy",
    "Japanese": "nova",
    "Chinese": "nova",
    "Thai": "alloy",
    "Russian": "alloy",
    "Vietnamese": "alloy",
    "Spanish": "alloy",
    "French": "alloy"
}

# ----------------------------------------------------------
#  (1) 음성(STT) → 번역 → 음성(TTS)
# ----------------------------------------------------------
def translate_audio(audio_file, source_lang, target_lang):
    if not audio_file:
        return "⚠️ 오디오 파일을 업로드하거나 녹음하세요.", "", None
    if not api_key or not client:
        return "❌ API 초기화 오류", "", None
    if source_lang == target_lang:
        return "⚠️ 입력 언어와 출력 언어가 같습니다.", "", None

    try:
        # ---------- Whisper STT ----------
        lang_code = LANG_CODE_MAP.get(source_lang, None)
        with open(audio_file, "rb") as f:
            transcript = client.audio.transcriptions.create(
                model="whisper-1",
                file=f,
                language=lang_code if lang_code else None  # 못 찾으면 자동감지
            )
        original_text = transcript.text.strip()
        if not original_text:
            return "⚠️ 음성이 인식되지 않았습니다.", "", None

        # ---------- GPT 번역 ----------
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system",
                 "content": f"You are a professional translator. Translate the following {source_lang} text to {target_lang}. "
                            f"Only provide the translation without any explanation or additional text."},
                {"role": "user", "content": original_text}
            ],
            temperature=0.3,
            max_tokens=2000
        )
        translated_text = response.choices[0].message.content.strip()

        # ---------- TTS ----------
        tts_response = client.audio.speech.create(
            model="tts-1",
            voice=VOICE_MAP.get(target_lang, "alloy"),
            input=translated_text[:4096]
        )
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
            tmp.write(tts_response.content)
            output_audio = tmp.name

        return original_text, translated_text, output_audio

    except Exception as e:
        return f"❌ 오류: {type(e).__name__}: {str(e)}", "", None

# ----------------------------------------------------------
#  (2) PDF / 이미지 → 번역
# ----------------------------------------------------------
def translate_document(file_obj, source_lang, target_lang):
    if not file_obj:
        return "⚠️ 파일을 업로드하세요.", ""
    if not api_key or not client:
        return "❌ API 초기화 오류", ""
    if source_lang == target_lang:
        return "⚠️ 입력 언어와 출력 언어가 같습니다.", ""

    ext = os.path.splitext(file_obj.name)[1].lower()
    try:
        # --- 텍스트 추출 ---
        if ext == ".pdf":
            import pdfplumber
            text_chunks = []
            with pdfplumber.open(file_obj.name) as pdf:
                for page in pdf.pages[:5]:  # 데모: 앞 5쪽만
                    text_chunks.append(page.extract_text() or "")
            original_text = "\n".join(text_chunks).strip()

        elif ext in [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff"]:
            from PIL import Image
            import pytesseract
            original_text = pytesseract.image_to_string(Image.open(file_obj.name))

        else:
            return "⚠️ 지원하지 않는 형식입니다.", ""

        if not original_text:
            return "⚠️ 텍스트를 추출할 수 없습니다.", ""

        # --- 번역 ---
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system",
                 "content": f"You are a professional translator. Translate the following {source_lang} text to {target_lang}. "
                            f"Only provide the translation without any explanation or additional text."},
                {"role": "user", "content": original_text}
            ],
            temperature=0.3,
            max_tokens=4096
        )
        translated_text = response.choices[0].message.content.strip()
        return original_text, translated_text

    except Exception as e:
        return f"❌ 오류: {type(e).__name__}: {str(e)}", ""

# ==========================================================
# Gradio UI
# ==========================================================
with gr.Blocks(title="SMARTok Demo", theme=gr.themes.Soft()) as app:
    with gr.Tabs():
        # ----- 🎙️ 음성 번역 -----
        with gr.TabItem("🎙️ 음성 번역"):
            gr.Markdown("""
            # 🎙️ AI 음성 번역기  
            마이크로 녹음하거나 오디오 파일을 업로드하면 **실시간 자막 + 번역 + 음성합성**까지 한 번에!
            """)

            with gr.Row():
                src_lang_a = gr.Dropdown(LANGUAGES, value="Korean", label="입력 언어")
                tgt_lang_a = gr.Dropdown(LANGUAGES, value="English", label="출력 언어")

            audio_in = gr.Audio(
                sources=["microphone", "upload"],
                type="filepath",
                label="음성 입력 (녹음 또는 파일 업로드)"
            )
            btn_audio = gr.Button("🔄 번역하기")

            with gr.Row():
                stt_text = gr.Textbox(label="📝 원본 텍스트", lines=5)
                tlt_text = gr.Textbox(label="🌐 번역된 텍스트", lines=5)

            audio_out = gr.Audio(label="🔊 번역된 음성", type="filepath", autoplay=True)

            btn_audio.click(
                translate_audio,
                inputs=[audio_in, src_lang_a, tgt_lang_a],
                outputs=[stt_text, tlt_text, audio_out]
            )

        # ----- 📄 자료 번역 -----
        with gr.TabItem("📄 자료 번역"):
            gr.Markdown("""
            # 📄 PDF / 이미지 번역 데모  
            교육자료·발표자료 등 **PDF 최대 5쪽** 또는 이미지 1장을 업로드하면 텍스트 추출 후 번역해줍니다.
            """)

            with gr.Row():
                src_lang_d = gr.Dropdown(LANGUAGES, value="Korean", label="입력 언어")
                tgt_lang_d = gr.Dropdown(LANGUAGES, value="English", label="출력 언어")

            file_in = gr.File(label="PDF / 이미지 업로드")
            btn_doc = gr.Button("🔄 번역하기")

            original_doc = gr.Textbox(label="📝 추출된 원문", lines=15)
            translated_doc = gr.Textbox(label="🌐 번역 결과", lines=15)

            btn_doc.click(
                translate_document,
                inputs=[file_in, src_lang_d, tgt_lang_d],
                outputs=[original_doc, translated_doc]
            )

# ==========================================================
if __name__ == "__main__":
    print("🚀 서버 시작 중...")
    app.launch(server_name="0.0.0.0", server_port=7860, share=False, debug=True)